"""Create a ChatVectorDBChain for question/answering."""
from langchain.callbacks.manager import AsyncCallbackManager
from langchain.callbacks.tracers import LangChainTracer
from langchain.chains import (
    ConversationalRetrievalChain, RetrievalQA
)
# from langchain.chains.chat_vector_db.prompts import (
# CONDENSE_QUESTION_PROMPT, QA_PROMPT)
from src.generators.prompt_templates import (
    QA_PROMPT, CONDENSE_QUESTION_PROMPT
    )
from langchain.chains.llm import LLMChain
from langchain.chains.question_answering import load_qa_chain
from langchain.chat_models import ChatOpenAI
from langchain.vectorstores.base import VectorStore
from langchain.memory import ConversationBufferWindowMemory


def get_chain(
    vectorstore: VectorStore,
    question_handler,
    stream_handler,
    tracing: bool = False
) -> ConversationalRetrievalChain:
    """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(
        model_name="gpt-3.5-turbo",
        streaming=True,
        callback_manager=question_manager,
        verbose=True,
        max_retries=1
    )
    streaming_llm = ChatOpenAI(
        model_name="gpt-3.5-turbo",
        streaming=True,
        callback_manager=stream_manager,
        verbose=True,
        max_retries=1
    )

    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=QA_PROMPT,
        callback_manager=manager
    )

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


def get_chainM(
    vectorstore: VectorStore,
    question_handler,
    stream_handler
) -> ConversationalRetrievalChain:
    """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])
    memory = ConversationBufferWindowMemory(
        memory_key="chat_history", return_messages=True
        )

    question_gen_llm = ChatOpenAI(
        model_name="gpt-3.5-turbo",
        streaming=True,
        callback_manager=question_manager,
        verbose=False,
        max_retries=2
    )
    streaming_llm = ChatOpenAI(
        model_name="gpt-3.5-turbo",
        streaming=True,
        callback_manager=stream_manager,
        verbose=False,
        max_retries=2,
        temperature=0
    )

    question_generator = LLMChain(
        llm=question_gen_llm,
        prompt=CONDENSE_QUESTION_PROMPT,
        callback_manager=manager,
        verbose=False
    )

    doc_chain = load_qa_chain(
        streaming_llm,
        chain_type="stuff",
        prompt=QA_PROMPT,
        callback_manager=manager,
        verbose=False
    )

    qa = ConversationalRetrievalChain(
        retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
        combine_docs_chain=doc_chain,
        callback_manager=manager,
        question_generator=question_generator,
        memory=memory,
        verbose=False
    )
    return qa


def get_chain_RetrievalQA(
    vectorstore: VectorStore, stream_handler, tracing: bool = False
) -> RetrievalQA:
    """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([])
    stream_manager = AsyncCallbackManager([stream_handler])
    if tracing:
        tracer = LangChainTracer()
        tracer.load_default_session()
        manager.add_handler(tracer)
        stream_manager.add_handler(tracer)

    streaming_llm = ChatOpenAI(
        model_name="gpt-3.5-turbo",
        streaming=True,
        callback_manager=stream_manager,
        verbose=True,
        max_retries=1
    )

    qa = RetrievalQA.from_llm(
        streaming_llm,
        retriever=vectorstore.as_retriever(k=2),
        callback_manager=manager,
        prompt=QA_PROMPT
    )
    return qa

langchain.chains.question_answering.load_qa_chain langchain.callbacks.tracers.LangChainTracer langchain.memory.ConversationBufferWindowMemory langchain.callbacks.manager.AsyncCallbackManager langchain.chains.llm.LLMChain langchain.chat_models.ChatOpenAI
[((1070, 1094), 'langchain.callbacks.manager.AsyncCallbackManager', 'AsyncCallbackManager', (['[]'], {}), '([])\n', (1090, 1094), False, 'from langchain.callbacks.manager import AsyncCallbackManager\n'), ((1118, 1158), 'langchain.callbacks.manager.AsyncCallbackManager', 'AsyncCallbackManager', (['[question_handler]'], {}), '([question_handler])\n', (1138, 1158), False, 'from langchain.callbacks.manager import AsyncCallbackManager\n'), ((1180, 1218), 'langchain.callbacks.manager.AsyncCallbackManager', 'AsyncCallbackManager', (['[stream_handler]'], {}), '([stream_handler])\n', (1200, 1218), False, 'from langchain.callbacks.manager import AsyncCallbackManager\n'), ((1456, 1579), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-3.5-turbo"""', 'streaming': '(True)', 'callback_manager': 'question_manager', 'verbose': '(True)', 'max_retries': '(1)'}), "(model_name='gpt-3.5-turbo', streaming=True, callback_manager=\n    question_manager, verbose=True, max_retries=1)\n", (1466, 1579), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1641, 1762), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-3.5-turbo"""', 'streaming': '(True)', 'callback_manager': 'stream_manager', 'verbose': '(True)', 'max_retries': '(1)'}), "(model_name='gpt-3.5-turbo', streaming=True, callback_manager=\n    stream_manager, verbose=True, max_retries=1)\n", (1651, 1762), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1830, 1923), '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', (1838, 1923), False, 'from langchain.chains.llm import LLMChain\n'), ((1967, 2063), 'langchain.chains.question_answering.load_qa_chain', 'load_qa_chain', (['streaming_llm'], {'chain_type': '"""stuff"""', 'prompt': 'QA_PROMPT', 'callback_manager': 'manager'}), "(streaming_llm, chain_type='stuff', prompt=QA_PROMPT,\n    callback_manager=manager)\n", (1980, 2063), False, 'from langchain.chains.question_answering import load_qa_chain\n'), ((2669, 2693), 'langchain.callbacks.manager.AsyncCallbackManager', 'AsyncCallbackManager', (['[]'], {}), '([])\n', (2689, 2693), False, 'from langchain.callbacks.manager import AsyncCallbackManager\n'), ((2717, 2757), 'langchain.callbacks.manager.AsyncCallbackManager', 'AsyncCallbackManager', (['[question_handler]'], {}), '([question_handler])\n', (2737, 2757), False, 'from langchain.callbacks.manager import AsyncCallbackManager\n'), ((2779, 2817), 'langchain.callbacks.manager.AsyncCallbackManager', 'AsyncCallbackManager', (['[stream_handler]'], {}), '([stream_handler])\n', (2799, 2817), False, 'from langchain.callbacks.manager import AsyncCallbackManager\n'), ((2831, 2910), 'langchain.memory.ConversationBufferWindowMemory', 'ConversationBufferWindowMemory', ([], {'memory_key': '"""chat_history"""', 'return_messages': '(True)'}), "(memory_key='chat_history', return_messages=True)\n", (2861, 2910), False, 'from langchain.memory import ConversationBufferWindowMemory\n'), ((2953, 3077), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-3.5-turbo"""', 'streaming': '(True)', 'callback_manager': 'question_manager', 'verbose': '(False)', 'max_retries': '(2)'}), "(model_name='gpt-3.5-turbo', streaming=True, callback_manager=\n    question_manager, verbose=False, max_retries=2)\n", (2963, 3077), False, 'from langchain.chat_models import ChatOpenAI\n'), ((3139, 3276), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-3.5-turbo"""', 'streaming': '(True)', 'callback_manager': 'stream_manager', 'verbose': '(False)', 'max_retries': '(2)', 'temperature': '(0)'}), "(model_name='gpt-3.5-turbo', streaming=True, callback_manager=\n    stream_manager, verbose=False, max_retries=2, temperature=0)\n", (3149, 3276), False, 'from langchain.chat_models import ChatOpenAI\n'), ((3352, 3460), 'langchain.chains.llm.LLMChain', 'LLMChain', ([], {'llm': 'question_gen_llm', 'prompt': 'CONDENSE_QUESTION_PROMPT', 'callback_manager': 'manager', 'verbose': '(False)'}), '(llm=question_gen_llm, prompt=CONDENSE_QUESTION_PROMPT,\n    callback_manager=manager, verbose=False)\n', (3360, 3460), False, 'from langchain.chains.llm import LLMChain\n'), ((3512, 3623), 'langchain.chains.question_answering.load_qa_chain', 'load_qa_chain', (['streaming_llm'], {'chain_type': '"""stuff"""', 'prompt': 'QA_PROMPT', 'callback_manager': 'manager', 'verbose': '(False)'}), "(streaming_llm, chain_type='stuff', prompt=QA_PROMPT,\n    callback_manager=manager, verbose=False)\n", (3525, 3623), False, 'from langchain.chains.question_answering import load_qa_chain\n'), ((4286, 4310), 'langchain.callbacks.manager.AsyncCallbackManager', 'AsyncCallbackManager', (['[]'], {}), '([])\n', (4306, 4310), False, 'from langchain.callbacks.manager import AsyncCallbackManager\n'), ((4332, 4370), 'langchain.callbacks.manager.AsyncCallbackManager', 'AsyncCallbackManager', (['[stream_handler]'], {}), '([stream_handler])\n', (4352, 4370), False, 'from langchain.callbacks.manager import AsyncCallbackManager\n'), ((4560, 4681), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-3.5-turbo"""', 'streaming': '(True)', 'callback_manager': 'stream_manager', 'verbose': '(True)', 'max_retries': '(1)'}), "(model_name='gpt-3.5-turbo', streaming=True, callback_manager=\n    stream_manager, verbose=True, max_retries=1)\n", (4570, 4681), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1252, 1269), 'langchain.callbacks.tracers.LangChainTracer', 'LangChainTracer', ([], {}), '()\n', (1267, 1269), False, 'from langchain.callbacks.tracers import LangChainTracer\n'), ((4404, 4421), 'langchain.callbacks.tracers.LangChainTracer', 'LangChainTracer', ([], {}), '()\n', (4419, 4421), False, 'from langchain.callbacks.tracers import LangChainTracer\n')]