"""Create a LangChain chain for question/answering."""
from langchain.callbacks.manager import AsyncCallbackManager
from langchain.callbacks.tracers import LangChainTracer
from langchain.chains import ConversationalRetrievalChain, RetrievalQAWithSourcesChain
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.llms.huggingface_endpoint import HuggingFaceEndpoint
from langchain.vectorstores.base import VectorStore
from langchain.vectorstores import Vectara
from langchain.memory import ConversationBufferMemory
from dotenv import load_dotenv
from langchain.chat_models import ChatOpenAI
from langchain.embeddings import HuggingFaceInferenceAPIEmbeddings
from langchain.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableParallel
from langchain.schema import StrOutputParser
from operator import itemgetter
import os


#
def format_docs(docs) -> str:
    return "\n\n".join(doc.page_content for doc in docs)


def get_chain(
    vectorstore: VectorStore, question_handler, stream_handler
) -> RunnableParallel:
    """Create a chain for question/answering."""

    load_dotenv()
    manager = AsyncCallbackManager([])
    question_manager = AsyncCallbackManager([question_handler])
    stream_manager = AsyncCallbackManager([stream_handler])

    hf_llm = HuggingFaceEndpoint(
        endpoint_url="https://euo6lqs9bqkddhci.us-east-1.aws.endpoints.huggingface.cloud",
        huggingfacehub_api_token=os.environ["HF_TOKEN"],
        task="text-generation",
        model_kwargs={
            "temperature": 0.1,
            "max_new_tokens": 488,
        },
    )

    retriever = vectorstore.as_retriever(search_kwargs={"k": 2})

    prompt_template = """\
    Use the provided context to answer the user's question. If you don't know the answer, say you don't know.

    Context:
    {context}

    Question:
    {question}
    
    Answer in french and do not start with 'Réponse:'
    """

    rag_prompt = ChatPromptTemplate.from_template(prompt_template)

    entry_point_chain = {
        "context": lambda input: format_docs(input["documents"]),
        "question": itemgetter("question"),
    }
    rag_chain = entry_point_chain | rag_prompt | hf_llm | StrOutputParser()

    rag_chain_with_sources = RunnableParallel(
        {"documents": retriever, "question": RunnablePassthrough()}
    ) | {
        "documents": lambda input: [doc.metadata for doc in input["documents"]],
        "answer": rag_chain,
    }
    return rag_chain_with_sources

langchain.llms.huggingface_endpoint.HuggingFaceEndpoint langchain.schema.StrOutputParser langchain.callbacks.manager.AsyncCallbackManager langchain_core.runnables.RunnablePassthrough langchain.prompts.ChatPromptTemplate.from_template
[((1270, 1283), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (1281, 1283), False, 'from dotenv import load_dotenv\n'), ((1298, 1322), 'langchain.callbacks.manager.AsyncCallbackManager', 'AsyncCallbackManager', (['[]'], {}), '([])\n', (1318, 1322), False, 'from langchain.callbacks.manager import AsyncCallbackManager\n'), ((1346, 1386), 'langchain.callbacks.manager.AsyncCallbackManager', 'AsyncCallbackManager', (['[question_handler]'], {}), '([question_handler])\n', (1366, 1386), False, 'from langchain.callbacks.manager import AsyncCallbackManager\n'), ((1408, 1446), 'langchain.callbacks.manager.AsyncCallbackManager', 'AsyncCallbackManager', (['[stream_handler]'], {}), '([stream_handler])\n', (1428, 1446), False, 'from langchain.callbacks.manager import AsyncCallbackManager\n'), ((1461, 1707), 'langchain.llms.huggingface_endpoint.HuggingFaceEndpoint', 'HuggingFaceEndpoint', ([], {'endpoint_url': '"""https://euo6lqs9bqkddhci.us-east-1.aws.endpoints.huggingface.cloud"""', 'huggingfacehub_api_token': "os.environ['HF_TOKEN']", 'task': '"""text-generation"""', 'model_kwargs': "{'temperature': 0.1, 'max_new_tokens': 488}"}), "(endpoint_url=\n    'https://euo6lqs9bqkddhci.us-east-1.aws.endpoints.huggingface.cloud',\n    huggingfacehub_api_token=os.environ['HF_TOKEN'], task='text-generation',\n    model_kwargs={'temperature': 0.1, 'max_new_tokens': 488})\n", (1480, 1707), False, 'from langchain.llms.huggingface_endpoint import HuggingFaceEndpoint\n'), ((2116, 2165), 'langchain.prompts.ChatPromptTemplate.from_template', 'ChatPromptTemplate.from_template', (['prompt_template'], {}), '(prompt_template)\n', (2148, 2165), False, 'from langchain.prompts import ChatPromptTemplate\n'), ((2279, 2301), 'operator.itemgetter', 'itemgetter', (['"""question"""'], {}), "('question')\n", (2289, 2301), False, 'from operator import itemgetter\n'), ((2367, 2384), 'langchain.schema.StrOutputParser', 'StrOutputParser', ([], {}), '()\n', (2382, 2384), False, 'from langchain.schema import StrOutputParser\n'), ((2478, 2499), 'langchain_core.runnables.RunnablePassthrough', 'RunnablePassthrough', ([], {}), '()\n', (2497, 2499), False, 'from langchain_core.runnables import RunnablePassthrough, RunnableParallel\n')]