import os
import gradio as gr
import pinecone
from gpt_index import GPTIndexMemory, GPTPineconeIndex
from langchain.agents import Tool
from langchain.chains.conversation.memory import ConversationBufferMemory
from langchain import OpenAI
from langchain.agents import initialize_agent

OPENAI_API_KEY=os.environ["OPENAI_API_KEY"]
PINECONE_API_KEY=os.environ["PINECONE_API_KEY"]

pinecone.init(api_key=PINECONE_API_KEY, environment="us-east1-gcp")

pindex=pinecone.Index("mahathir")
indexed_pinecone=GPTPineconeIndex([], pinecone_index=pindex)

tools = [
    Tool(
        name = "GPT Index",
        func=lambda q: str(indexed_pinecone.query(q)),
        description="useful for when you want to answer questions about the author. The input to this tool should be a complete english sentence.",
        return_direct=True
    ),
]
memory = GPTIndexMemory(index=indexed_pinecone, memory_key="chat_history", query_kwargs={"response_mode": "compact"})
llm=OpenAI(temperature=0)
agent_chain = initialize_agent(tools, llm, agent="conversational-react-description", memory=memory, verbose=True)

def predict(input, history=[]):
    response = agent_chain.run(input)
    history = history + [(input, response)]
    response = history
    # response = [response] 
    # return response, response
    return response, response

with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    state = gr.State([])

    with gr.Row():
        gr.Markdown("<h3><center>Train on three books and one wiki. A Doctor in The House (Memoir), his wikipedia, Malaysian Maverick, and The Malay Dillema </center></h3>")

    with gr.Row():
        txt = gr.Textbox(show_label=False, placeholder="Enter text and press enter").style(container=False)

    gr.Examples(
        examples=[
            "What are your opinion on malaysian resiliency during covid 19 from the perspective of Mahathir Mohamad",
            "how close was anwar ibrahim with mahathir before? any interesting stories?",
            "What is the focus for mahathir to move malaysia towards a more modern economy",
        ],
        inputs=txt,
    )
    
    txt.submit(predict, [txt, state], [chatbot, state])
    # txt.submit(agent_executor.run, [txt, state], [chatbot, state])

demo.launch()