| from langchain.chains import ConversationalRetrievalChain | |
| from langchain.chains.base import Chain | |
| from langchain.vectorstores.base import VectorStore | |
| from app_modules.llm_inference import LLMInference | |
| class QAChain(LLMInference): | |
| vectorstore: VectorStore | |
| def __init__(self, vectorstore, llm_loader, doc_id_to_vectorstore_mapping=None): | |
| super().__init__(llm_loader) | |
| self.vectorstore = vectorstore | |
| self.doc_id_to_vectorstore_mapping = doc_id_to_vectorstore_mapping | |
| def get_chain(self, inputs) -> Chain: | |
| return self.create_chain(inputs) | |
| def create_chain(self, inputs) -> Chain: | |
| vectorstore = self.vectorstore | |
| if "chat_id" in inputs: | |
| if inputs["chat_id"] in self.doc_id_to_vectorstore_mapping: | |
| vectorstore = self.doc_id_to_vectorstore_mapping[inputs["chat_id"]] | |
| qa = ConversationalRetrievalChain.from_llm( | |
| self.llm_loader.llm, | |
| vectorstore.as_retriever(search_kwargs=self.llm_loader.search_kwargs), | |
| max_tokens_limit=self.llm_loader.max_tokens_limit, | |
| return_source_documents=True, | |
| ) | |
| return qa | |