from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.agents import AgentExecutor, Tool
from langchain.chains import RetrievalQAWithSourcesChain
from langchain.chat_models import ChatOpenAI
from langchain.memory import ConversationBufferMemory, ConversationSummaryBufferMemory
from langchain.utilities import WikipediaAPIWrapper
from langchain.agents import initialize_agent, AgentType

import gradio as gr

# Add presets for Gradio theme
from app_modules.presets import * 

import os
os.environ["OPENAI_API_KEY"] = os.environ["OPENAI_TOKEN"]

# Define the LLM chat model
#model = 'gpt-3.5-turbo'
model = 'gpt-3.5-turbo-16k'
#model = 'gpt-4'
token_limit = 4000 if model == 'gpt-3.5-turbo' else 16000
temperature = 0
llm = ChatOpenAI(temperature=temperature, model=model)

# Load existing vectorstore
persist_dir = "./chroma"
embeddings = OpenAIEmbeddings()
vectorstore = Chroma(persist_directory=persist_dir, embedding_function=embeddings)
vectorstore.persist()

# Create Retrieval Chain with sources
## It returns a dictionary with at least the 'answer' and the 'sources'
qa = RetrievalQAWithSourcesChain.from_chain_type(
            llm=llm,
            chain_type="stuff",
            retriever=vectorstore.as_retriever(search_type="mmr"),
            #retriever=vectorstore.as_retriever(),
            return_source_documents=True, 
            max_tokens_limit=token_limit
        )

# Define tools
wikipedia = WikipediaAPIWrapper()

tools = [
    Tool(
        name="GBV Q&A Bot System",
        #func=qa,
        func=lambda question: qa({"question": question}, return_only_outputs=True),
        description="Useful for when you need to answer questions about the aspects asked. Input may be a partial or fully formed question.",
        #return_direct=True, # use the agent as a router and directly return the result
    ),
    Tool(
        name='Wikipedia',
        func=wikipedia.run,
        description='You must only use this tool if you cannot find answers with the other tools. Useful for when you need to look for answers in the Wikipedia.'
    )
]

# Create Conversational Buffer Memory
#memory = ConversationBufferMemory(memory_key="chat_history", input_key='input', output_key="output", return_messages=True)
# Create Conversational Summary Buffer Memory
memory = ConversationSummaryBufferMemory(llm=llm, memory_key="chat_history", input_key='input', output_key="output", return_messages=True)

# Initialize Re-Act agent and create Agent Executor Chain
react = initialize_agent(tools, llm, agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION, verbose=True, max_iterations=3, early_stopping_method='generate',  memory=memory)

agent_chain = AgentExecutor.from_agent_and_tools(
                agent=react.agent, tools=tools, verbose=True, memory=memory, return_intermediate_steps=True, return_source_documents=True, handle_parsing_errors=True
            )

# Add custom CSS
with open("assets/custom.css", "r", encoding="utf-8") as f:
    customCSS = f.read()

with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme, title="GBV Q&A Bot") as demo:
    
    gr.Markdown(
        """
        # 🦜🔗 Ask the GBV in Emergencies Q&A Bot!
        This generative model has been trained on various sources covering themes on Gender-Based Violence response in Humanitarian Settings. This AI agent might complement the replies with additional information retrieved from Wikipedia sources. You can maintain a natural language conversation with it in order to retrieve information on this area of knowledge. If you want to obtain the source information, type 'Please provide the source documents' right after your query (see example)

        Example questions:
        - What is GBV?
        - Which agency leads the GBV Response in Emergencies? Please provide the source documents
        - Please outline a strategy to minimize GBV risks in a temporary settlement
        """
    )
    
    # Start chatbot with welcome from bot
    chatbot = gr.Chatbot([(None,'How can I help you?')]).style(height=550)
    msg = gr.Textbox()
    clear = gr.ClearButton([msg, chatbot])

    def user(user_message, history):
        return gr.update(value="", interactive=False), history + [[user_message, None]]

    def bot(history):
        user_message = history[-1][0] # get if from most recent history element
        #user_message = user_message + " Please provide the source documents" # to alter the prompt and provide sources
        response = agent_chain(user_message)
        bot_message = response['output']
        history[-1][1] = ""
        for character in bot_message:
            history[-1][1] += character
            yield history

    response = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    response.then(lambda: gr.update(interactive=True), None, [msg], queue=False)

demo.queue()
demo.launch()