UnknownPixel/askPESU

1"""Retrieval-Augmented Generation (RAG) pipeline implementation using LangChain and Qdrant."""2 3import os4 5import yaml6from dotenv import load_dotenv7from langchain.chat_models import init_chat_model8from langchain.retrievers.multi_query import MultiQueryRetriever9from langchain_core.documents.base import Document10from langchain_core.language_models import BaseChatModel11from langchain_core.messages import AIMessage, HumanMessage12from langchain_core.output_parsers import StrOutputParser13from langchain_core.prompts import ChatPromptTemplate14from langchain_core.runnables import RunnablePassthrough, RunnableSerializable15from langchain_huggingface.embeddings import HuggingFaceEmbeddings16from langchain_qdrant import QdrantVectorStore17from qdrant_client import QdrantClient18 19load_dotenv()20 21 22class RetrievalAugmentedGenerator:23    """A class that encapsulates the Retrieval-Augmented Generation (RAG) pipeline."""24 25    def __init__(self, config_path: str = "conf/config.yaml") -> None:26        """Initialize the RAG pipeline with configuration from a YAML file.27 28        Args:29            config_path (str): Path to the configuration YAML file.30        """31        # Load configuration from YAML file32        with open(config_path) as file:33            self.config = yaml.safe_load(file)34 35        # Initialize embeddings36        self.embedding = HuggingFaceEmbeddings(model_name=self.config["rag"]["embedding"])37 38        # Initialize Qdrant client and vector store39        self.qdrant_client = QdrantClient(url=os.getenv("QDRANT_URL"), api_key=os.getenv("QDRANT_API_KEY"))40        self.vector_store = QdrantVectorStore(41            collection_name=self.config["rag"]["qdrant_collection"],42            embedding=self.embedding,43            client=self.qdrant_client,44        )45 46        # Initialize LLM47        self.llm_primary = init_chat_model(48            model=self.config["rag"]["llm"]["primary"],49            model_provider="google_genai",50            google_api_key=os.getenv("GEMINI_API_KEY"),51        )52        # Initialize secondary LLM if specified53        self.llm_thinking = None54        if self.config["rag"]["llm"].get("thinking"):55            self.llm_thinking = init_chat_model(56                model=self.config["rag"]["llm"]["thinking"],57                model_provider="google_genai",58                google_api_key=os.getenv("GEMINI_API_KEY"),59            )60 61        # Initialize the prompt template62        self.prompt = ChatPromptTemplate.from_messages(63            [64                ("system", self.config["rag"]["system_prompt"]),65                ("human", "Question: {question}\nContext: {context}\nAnswer:"),66            ]67        )68 69        self.frame_qn_prompt = ChatPromptTemplate.from_messages(70            [71                (72                    "system",73                    "You are a question rewriting assistant. Your job is to rewrite the user's "74                    "question into an independent, self-contained question.\n\n"75                    "Rewrite rules:\n"76                    "1.ONLY use the chat history if the user's question is ambiguous or refers to previous context "77                    "(e.g., pronouns like 'he', 'she', 'it', 'they', 'that').\n"78                    "2.If the question is clear on its own, return it EXACTLY as it is.\n"79                    "3.When resolving a follow-up question, ALWAYS prioritize the most recent topic in the chat history"80                    "Do NOT pull context from older, unrelated parts of the conversation.\n"81                    "4.If the question could refer to multiple topics, choose the MOST RECENT plausible topic.\n"82                    "5.Do NOT invent or assume connections between unrelated topics.\n"83                    "6.Do NOT answer the question — only rewrite it.\n\n"84                    "Chat History:\n{chat_history}",85                ),86                ("human", "{input}"),87            ]88        )89 90        # Build the RAG chains91        self.retriever = self.vector_store.as_retriever(search_kwargs=self.config["rag"]["search_kwargs"])92        self.rag_chain_primary = self._build_chain(self.llm_primary)93        self.rag_chain_thinking = self._build_chain(self.llm_thinking) if self.llm_thinking else None94 95    def _build_chain(self, llm: BaseChatModel) -> RunnableSerializable[str, str]:96        """Build the RAG chain using the specified LLM.97 98        Args:99            llm: The language model to use in the RAG chain.100 101        Returns:102            RunnableSerializable: The constructed RAG chain.103        """104        # Initialize multiquery retriever105        multiquery_retriever = MultiQueryRetriever.from_llm(106            retriever=self.retriever,107            llm=llm,108        )109 110        history_aware_retriever = (111            {"input": RunnablePassthrough(), "chat_history": RunnablePassthrough()}112            | self.frame_qn_prompt113            | llm114            | StrOutputParser()115            | multiquery_retriever116        )117 118        return (119            {120                "context": history_aware_retriever | self.format_docs,121                "question": RunnablePassthrough(),122            }123            | self.prompt124            | llm125            | StrOutputParser()126        )127 128    @staticmethod129    def format_docs(docs: list[Document]) -> str:130        """Format the retrieved documents into a single string."""131        return "\n\n".join(f"{doc.metadata['url']}\n{doc.page_content}" for doc in docs)132 133    async def generate(self, query: str, thinking: bool, history: list) -> str:134        """Generate a response for the given query using the RAG chain.135 136        Args:137            query (str): The input query.138            thinking (bool): Flag to indicate if the model should 'think' before answering.139            history (list): The entire chat history until the current query140 141        Returns:142            str: The generated response.143        """144        chat_history = []145 146        for convo in history:147            if query != convo.query:  # Prevents repeating the same question when using the thinking model.148                chat_history.append(HumanMessage(convo.query))149                chat_history.append(AIMessage(convo.answer))150 151        rag_chain = (152            self.rag_chain_thinking if thinking and self.rag_chain_thinking is not None else self.rag_chain_primary153        )154 155        return await rag_chain.ainvoke({"input": query, "question": query, "chat_history": chat_history})156