app_modules/llm_qa_chain.py

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: int = 2048):
        super().__init__(llm_loader)
        self.vectorstore = vectorstore

    def create_chain(self) -> Chain:
        qa = ConversationalRetrievalChain.from_llm(
            self.llm_loader.llm,
            self.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