"""Create a ChatVectorDBChain for question/answering."""
from langchain.callbacks.manager import AsyncCallbackManager
from langchain.callbacks.tracers import LangChainTracer
from langchain.chains import ChatVectorDBChain
from langchain.chains.chat_vector_db.prompts import (CONDENSE_QUESTION_PROMPT,
                                                     QA_PROMPT)
from langchain.chains.llm import LLMChain
from langchain.chains.question_answering import load_qa_chain
from langchain.chat_models import ChatOpenAI
from langchain.llms import OpenAI
from langchain.vectorstores import Pinecone
from langchain.chains import ConversationalRetrievalChain

from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

# system_template = """Use the following pieces of context to answer the users question.
# If you don't know the answer, just say that you don't know, don't try to make up an answer.
# ----------------
# {context}"""

template = """You are a helpful AI assistant that answers questions about
an e-commerce company called "Sindabad.com" in a friendly and polite
manner. You will be given a context that will represent Sindabad.com's
product inventory. Users might ask about products, they might want to
know your suggestions as well. Most importantly, they might ask about
specific product and its associated product link. If they want to know
about product links, you will provide it accordingly with the help of the
given "Context". Answer the question in your own words as truthfully as
possible from the context given to you. If you do not know the answer to
the question, simply respond with "I don't know. Could you please rephrase
the question?". If questions are asked where there is no relevant information
available in the context, answer the question with your existing knowledge on
that question and "ignore" the "Context" given to you.

----------------
context: {context}"""


messages = [
    SystemMessagePromptTemplate.from_template(template),
    HumanMessagePromptTemplate.from_template("{question}")
]
prompt = ChatPromptTemplate.from_messages(messages)


def get_chain(
    vectorstore: Pinecone,
    question_handler,
    stream_handler,
    tracing: bool = False
) -> ChatVectorDBChain:
    """Create a ChatVectorDBChain for question/answering."""
    # Construct a ChatVectorDBChain with a streaming llm for combine docs
    # and a separate, non-streaming llm for question generation
    manager = AsyncCallbackManager([])
    question_manager = AsyncCallbackManager([question_handler])
    stream_manager = AsyncCallbackManager([stream_handler])
    if tracing:
        tracer = LangChainTracer()
        tracer.load_default_session()
        manager.add_handler(tracer)
        question_manager.add_handler(tracer)
        stream_manager.add_handler(tracer)

    question_gen_llm = ChatOpenAI(
        temperature=0,
        verbose=True,
        callback_manager=question_manager,
    )
    streaming_llm = ChatOpenAI(
        streaming=True,
        callback_manager=stream_manager,
        verbose=True,
        temperature=0,
    )

    question_generator = LLMChain(
        llm=question_gen_llm,
        prompt=CONDENSE_QUESTION_PROMPT,
        callback_manager=manager
    )
    doc_chain = load_qa_chain(
        streaming_llm,
        chain_type="stuff",
        prompt=prompt,
        callback_manager=manager
    )

    # qa = ChatVectorDBChain(
    #     vectorstore=vectorstore,
    #     combine_docs_chain=doc_chain,
    #     question_generator=question_generator,
    #     callback_manager=manager,
    # )
    qa = ConversationalRetrievalChain(
        retriever=vectorstore.as_retriever(),
        combine_docs_chain=doc_chain,
        question_generator=question_generator,
        callback_manager=manager
    )
    return qa

langchain.chains.question_answering.load_qa_chain langchain.callbacks.tracers.LangChainTracer langchain.callbacks.manager.AsyncCallbackManager langchain.prompts.chat.SystemMessagePromptTemplate.from_template langchain.chat_models.ChatOpenAI langchain.chains.llm.LLMChain langchain.prompts.chat.HumanMessagePromptTemplate.from_template langchain.prompts.chat.ChatPromptTemplate.from_messages
[((2109, 2151), 'langchain.prompts.chat.ChatPromptTemplate.from_messages', 'ChatPromptTemplate.from_messages', (['messages'], {}), '(messages)\n', (2141, 2151), False, 'from langchain.prompts.chat import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate\n'), ((1986, 2037), 'langchain.prompts.chat.SystemMessagePromptTemplate.from_template', 'SystemMessagePromptTemplate.from_template', (['template'], {}), '(template)\n', (2027, 2037), False, 'from langchain.prompts.chat import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate\n'), ((2043, 2097), 'langchain.prompts.chat.HumanMessagePromptTemplate.from_template', 'HumanMessagePromptTemplate.from_template', (['"""{question}"""'], {}), "('{question}')\n", (2083, 2097), False, 'from langchain.prompts.chat import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate\n'), ((2501, 2525), 'langchain.callbacks.manager.AsyncCallbackManager', 'AsyncCallbackManager', (['[]'], {}), '([])\n', (2521, 2525), False, 'from langchain.callbacks.manager import AsyncCallbackManager\n'), ((2549, 2589), 'langchain.callbacks.manager.AsyncCallbackManager', 'AsyncCallbackManager', (['[question_handler]'], {}), '([question_handler])\n', (2569, 2589), False, 'from langchain.callbacks.manager import AsyncCallbackManager\n'), ((2611, 2649), 'langchain.callbacks.manager.AsyncCallbackManager', 'AsyncCallbackManager', (['[stream_handler]'], {}), '([stream_handler])\n', (2631, 2649), False, 'from langchain.callbacks.manager import AsyncCallbackManager\n'), ((2887, 2961), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': '(0)', 'verbose': '(True)', 'callback_manager': 'question_manager'}), '(temperature=0, verbose=True, callback_manager=question_manager)\n', (2897, 2961), False, 'from langchain.chat_models import ChatOpenAI\n'), ((3013, 3105), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'streaming': '(True)', 'callback_manager': 'stream_manager', 'verbose': '(True)', 'temperature': '(0)'}), '(streaming=True, callback_manager=stream_manager, verbose=True,\n    temperature=0)\n', (3023, 3105), False, 'from langchain.chat_models import ChatOpenAI\n'), ((3167, 3260), 'langchain.chains.llm.LLMChain', 'LLMChain', ([], {'llm': 'question_gen_llm', 'prompt': 'CONDENSE_QUESTION_PROMPT', 'callback_manager': 'manager'}), '(llm=question_gen_llm, prompt=CONDENSE_QUESTION_PROMPT,\n    callback_manager=manager)\n', (3175, 3260), False, 'from langchain.chains.llm import LLMChain\n'), ((3303, 3396), 'langchain.chains.question_answering.load_qa_chain', 'load_qa_chain', (['streaming_llm'], {'chain_type': '"""stuff"""', 'prompt': 'prompt', 'callback_manager': 'manager'}), "(streaming_llm, chain_type='stuff', prompt=prompt,\n    callback_manager=manager)\n", (3316, 3396), False, 'from langchain.chains.question_answering import load_qa_chain\n'), ((2683, 2700), 'langchain.callbacks.tracers.LangChainTracer', 'LangChainTracer', ([], {}), '()\n', (2698, 2700), False, 'from langchain.callbacks.tracers import LangChainTracer\n')]