import os
import gradio as gr
import pinecone
from llama_index import GPTIndexMemory, GPTPineconeIndex, QuestionAnswerPrompt
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="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"
os.environ['OPENAI_API_KEY'] = "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"
PINECONE_API_KEY="3fa304c8-f592-4ecb-b960-21a5966e80ab"

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

pindex=pinecone.Index("sejarah")

#define custom QuestionAnswerPrompt- adding clear and concise instruction to the model
query_str = "Bincangkan kesan sekiranya semangat assabiyah diamalkan dalam masyarakat hari ini."


QA_PROMPT_TMPL = (
    "Answer each question truthfully using the Malaysia's Form 1 History textbook data provided."
    "Your answers should be concise and straight to the point.For questions that are open-ended, which require subjective judgment or opinion, you may not find a definitive answer in the textbook."
    "However, you should still address the question's directive based on the data's context. Ideally, your answer should provide 3 points that support your response. You are encouraged to better provide positive suggestions for concepts that are less ethical.Please keep in mind that the scope of the data provided is limited to the content covered in the Malaysia's Form 1 History textbook."
    "If you do not know the answer please say 'I don't know, Please rephrase the question'\n"
    "---------------------\n Context:\n"
    "{context_str}"
    "\n---------------------\n"
    "Given this information, please answer the question: {query_str}\n"
)


QA_PROMPT = QuestionAnswerPrompt(QA_PROMPT_TMPL)


indexed_pinecone=GPTPineconeIndex([], pinecone_index=pindex)

tools = [
    Tool(
        name = "GPT Index",
        func=lambda q: str(indexed_pinecone.query(q,text_qa_template=QA_PROMPT)),
        description="Asking questions to Pandai AI",
        return_direct=True
    )
]

memory = GPTIndexMemory(index=indexed_pinecone, memory_key="chat_history", query_kwargs={"response_mode": "compact"})
llm=OpenAI(temperature=0, model_name="text-davinci-003",max_tokens=1000)
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():
        txt = gr.Textbox(show_label=False, placeholder="Enter text and press enter").style(container=False)

    txt.submit(predict, [txt, state], [chatbot, state])
    # txt.submit(agent_executor.run, [txt, state], [chatbot, state])

demo.launch()