from langchain.memory import ConversationBufferMemory
from langchain.chains.conversational_retrieval.base import ConversationalRetrievalChain


class Chatbot:
    def __init__(self) -> None:
        pass

    def create_conversational_chain(self, llm, db_retriever):
        memory = ConversationBufferMemory(
            memory_key="chat_history", return_messages=True, output_key="answer"
        )

        return ConversationalRetrievalChain.from_llm(
            llm=llm, retriever=db_retriever, memory=memory, return_source_documents=True
        )