from langchain.chains import RetrievalQA, ConversationalRetrievalChain, ConversationChain
from langchain.prompts.prompt import PromptTemplate
from langchain.vectorstores.base import VectorStoreRetriever
from langchain.chat_models import ChatOpenAI
from langchain.memory import ConversationBufferMemory
import pickle
import os
from env import OPENAI_API_KEY

os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY

_template = """Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question.
You can assume the question about the most recent state of the union address.

Chat History:
{chat_history}
Follow Up Input: {question}
Standalone question:"""
CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(_template)

template = """You are an AI assistant for answering questions about the most recent state of the union address.
You are given the following extracted parts of a long document and a question. Provide a conversational answer.
If you don't know the answer, just say "Hmm, I'm not sure." Don't try to make up an answer.
If the question is not about the most recent state of the union, politely inform them that you are tuned to only answer questions about the most recent state of the union.
Lastly, answer the question as if you were a pirate from the south seas and are just coming back from a pirate expedition where you found a treasure chest full of gold doubloons.
Question: {question}
=========
{context}
=========
Answer in Markdown:"""
QA_PROMPT = PromptTemplate(template=template, input_variables=[
                           "question", "context"])


pyhealth_template = """
1. role clarification, task definition (high-level instruction), and model encouragement
2. task flow/procedures (input/output clarification)
3. information (chat history, retrieval results: doc + codes) -> separate code doc and txt doc
4. notice (law, regulation, policy, etc.) -> !!!
5. now, give the response.
tricks: 
    - important info should be at the beginning or in the end. knowledge is in the middle.
    - use "sep" to block the prompt
"""


def load_retriever():
    with open("vectorstore.pkl", "rb") as f:
        vectorstore = pickle.load(f)
    retriever = VectorStoreRetriever(vectorstore=vectorstore)
    return retriever


def get_basic_qa_chain():
    llm = ChatOpenAI(model_name="gpt-4", temperature=0)
    retriever = load_retriever()
    memory = ConversationBufferMemory(
        memory_key="chat_history", return_messages=True)
    model = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever,
        memory=memory)
    return model


def get_custom_prompt_qa_chain():
    llm = ChatOpenAI(model_name="gpt-4", temperature=0)
    retriever = load_retriever()
    memory = ConversationBufferMemory(
        memory_key="chat_history", return_messages=True)
    # see: https://github.com/langchain-ai/langchain/issues/6635
    # see: https://github.com/langchain-ai/langchain/issues/1497
    model = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever,
        memory=memory,
        combine_docs_chain_kwargs={"prompt": QA_PROMPT})
    return model


def get_condense_prompt_qa_chain():
    llm = ChatOpenAI(model_name="gpt-4", temperature=0)
    retriever = load_retriever()
    memory = ConversationBufferMemory(
        memory_key="chat_history", return_messages=True)
    # see: https://github.com/langchain-ai/langchain/issues/5890
    model = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever,
        memory=memory,
        condense_question_prompt=CONDENSE_QUESTION_PROMPT,
        combine_docs_chain_kwargs={"prompt": QA_PROMPT})
    return model


def get_qa_with_sources_chain():
    llm = ChatOpenAI(model_name="gpt-4", temperature=0)
    retriever = load_retriever()
    history = []
    model = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever,
        return_source_documents=True)

    def model_func(question):
        # bug: this doesn't work with the built-in memory
        # hacking around it for the tutorial
        # see: https://github.com/langchain-ai/langchain/issues/5630
        new_input = {"question": question['question'], "chat_history": history}
        result = model(new_input)
        history.append((question['question'], result['answer']))
        return result

    return model_func


def get_basic_chain():
    llm = ChatOpenAI(model_name="gpt-4", temperature=0)
    memory = ConversationBufferMemory()
    model = ConversationChain(
        llm=llm,
        memory=memory)
    return model


chain_options = {
    "basic": get_basic_qa_chain,
    "with_sources": get_qa_with_sources_chain,
    "custom_prompt": get_custom_prompt_qa_chain,
    "condense_prompt": get_condense_prompt_qa_chain
}

langchain.chains.ConversationChain langchain.prompts.prompt.PromptTemplate langchain.vectorstores.base.VectorStoreRetriever langchain.chains.ConversationalRetrievalChain.from_llm langchain.memory.ConversationBufferMemory langchain.chat_models.ChatOpenAI langchain.prompts.prompt.PromptTemplate.from_template
[((727, 766), 'langchain.prompts.prompt.PromptTemplate.from_template', 'PromptTemplate.from_template', (['_template'], {}), '(_template)\n', (755, 766), False, 'from langchain.prompts.prompt import PromptTemplate\n'), ((1521, 1595), 'langchain.prompts.prompt.PromptTemplate', 'PromptTemplate', ([], {'template': 'template', 'input_variables': "['question', 'context']"}), "(template=template, input_variables=['question', 'context'])\n", (1535, 1595), False, 'from langchain.prompts.prompt import PromptTemplate\n'), ((2225, 2270), 'langchain.vectorstores.base.VectorStoreRetriever', 'VectorStoreRetriever', ([], {'vectorstore': 'vectorstore'}), '(vectorstore=vectorstore)\n', (2245, 2270), False, 'from langchain.vectorstores.base import VectorStoreRetriever\n'), ((2330, 2375), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-4"""', 'temperature': '(0)'}), "(model_name='gpt-4', temperature=0)\n", (2340, 2375), False, 'from langchain.chat_models import ChatOpenAI\n'), ((2422, 2495), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""', 'return_messages': '(True)'}), "(memory_key='chat_history', return_messages=True)\n", (2446, 2495), False, 'from langchain.memory import ConversationBufferMemory\n'), ((2517, 2604), 'langchain.chains.ConversationalRetrievalChain.from_llm', 'ConversationalRetrievalChain.from_llm', ([], {'llm': 'llm', 'retriever': 'retriever', 'memory': 'memory'}), '(llm=llm, retriever=retriever, memory=\n    memory)\n', (2554, 2604), False, 'from langchain.chains import RetrievalQA, ConversationalRetrievalChain, ConversationChain\n'), ((2688, 2733), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-4"""', 'temperature': '(0)'}), "(model_name='gpt-4', temperature=0)\n", (2698, 2733), False, 'from langchain.chat_models import ChatOpenAI\n'), ((2780, 2853), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""', 'return_messages': '(True)'}), "(memory_key='chat_history', return_messages=True)\n", (2804, 2853), False, 'from langchain.memory import ConversationBufferMemory\n'), ((3005, 3141), 'langchain.chains.ConversationalRetrievalChain.from_llm', 'ConversationalRetrievalChain.from_llm', ([], {'llm': 'llm', 'retriever': 'retriever', 'memory': 'memory', 'combine_docs_chain_kwargs': "{'prompt': QA_PROMPT}"}), "(llm=llm, retriever=retriever, memory=\n    memory, combine_docs_chain_kwargs={'prompt': QA_PROMPT})\n", (3042, 3141), False, 'from langchain.chains import RetrievalQA, ConversationalRetrievalChain, ConversationChain\n'), ((3235, 3280), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-4"""', 'temperature': '(0)'}), "(model_name='gpt-4', temperature=0)\n", (3245, 3280), False, 'from langchain.chat_models import ChatOpenAI\n'), ((3327, 3400), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""', 'return_messages': '(True)'}), "(memory_key='chat_history', return_messages=True)\n", (3351, 3400), False, 'from langchain.memory import ConversationBufferMemory\n'), ((3487, 3678), 'langchain.chains.ConversationalRetrievalChain.from_llm', 'ConversationalRetrievalChain.from_llm', ([], {'llm': 'llm', 'retriever': 'retriever', 'memory': 'memory', 'condense_question_prompt': 'CONDENSE_QUESTION_PROMPT', 'combine_docs_chain_kwargs': "{'prompt': QA_PROMPT}"}), "(llm=llm, retriever=retriever, memory=\n    memory, condense_question_prompt=CONDENSE_QUESTION_PROMPT,\n    combine_docs_chain_kwargs={'prompt': QA_PROMPT})\n", (3524, 3678), False, 'from langchain.chains import RetrievalQA, ConversationalRetrievalChain, ConversationChain\n'), ((3773, 3818), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-4"""', 'temperature': '(0)'}), "(model_name='gpt-4', temperature=0)\n", (3783, 3818), False, 'from langchain.chat_models import ChatOpenAI\n'), ((3881, 3982), 'langchain.chains.ConversationalRetrievalChain.from_llm', 'ConversationalRetrievalChain.from_llm', ([], {'llm': 'llm', 'retriever': 'retriever', 'return_source_documents': '(True)'}), '(llm=llm, retriever=retriever,\n    return_source_documents=True)\n', (3918, 3982), False, 'from langchain.chains import RetrievalQA, ConversationalRetrievalChain, ConversationChain\n'), ((4466, 4511), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-4"""', 'temperature': '(0)'}), "(model_name='gpt-4', temperature=0)\n", (4476, 4511), False, 'from langchain.chat_models import ChatOpenAI\n'), ((4525, 4551), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {}), '()\n', (4549, 4551), False, 'from langchain.memory import ConversationBufferMemory\n'), ((4564, 4605), 'langchain.chains.ConversationChain', 'ConversationChain', ([], {'llm': 'llm', 'memory': 'memory'}), '(llm=llm, memory=memory)\n', (4581, 4605), False, 'from langchain.chains import RetrievalQA, ConversationalRetrievalChain, ConversationChain\n'), ((2194, 2208), 'pickle.load', 'pickle.load', (['f'], {}), '(f)\n', (2205, 2208), False, 'import pickle\n')]