Update rag.py

rag.py

@@ -6,16 +6,17 @@ import zipfile
 import gdown
 import time
 import asyncio
+
 from typing import Dict, List
 
 from google import genai
 
-from langchain_community.vectorstores import Chroma
-from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain.memory import ConversationBufferMemory
-from langchain.retrievers import BM25Retriever, EnsembleRetriever
 from langchain.schema import Document
 
+from langchain_community.vectorstores import Chroma
+from langchain_community.embeddings import HuggingFaceEmbeddings
+
 
 # =====================================================
 # Logging
@@ -41,22 +42,34 @@ MODEL_NAME = "models/gemini-2.5-flash-lite"
 # Download Vector DB
 # =====================================================
 if not os.path.exists("chroma_db"):
+
     print("Downloading vector DB...")
 
-    url = "https://drive.google.com/uc?id=1yelCLmiRD2Qds-_8EsnKzUNY7uPT5xMg"
+    url = (
+        "https://drive.google.com/uc?id="
+        "1yelCLmiRD2Qds-_8EsnKzUNY7uPT5xMg"
+    )
 
-    gdown.download(url, "chroma_db.zip", quiet=False)
+    gdown.download(
+        url,
+        "chroma_db.zip",
+        quiet=False
+    )
 
     print("Extracting DB...")
 
-    with zipfile.ZipFile("chroma_db.zip", "r") as zip_ref:
+    with zipfile.ZipFile(
+        "chroma_db.zip",
+        "r"
+    ) as zip_ref:
+
         zip_ref.extractall(".")
 
     print("Vector DB Ready ✅")
 
 
 # =====================================================
-# Embedding Model
+# Load Embedding Model
 # =====================================================
 embedding = HuggingFaceEmbeddings(
     model_name="sentence-transformers/all-MiniLM-L6-v2"
@@ -75,61 +88,58 @@ print("Vector DB Loaded ✅")
 
 
 # =====================================================
-# Build BM25 Retriever (Hybrid Search)
-# =====================================================
-print("Building BM25 Retriever...")
-
-all_docs = vector_db.similarity_search(
-    "nutrition",
-    k=200
-)
-
-bm25_retriever = BM25Retriever.from_documents(all_docs)
-bm25_retriever.k = 6
-
-vector_retriever = vector_db.as_retriever(
-    search_kwargs={"k": 6}
-)
-
-# Hybrid Search
-hybrid_retriever = EnsembleRetriever(
-    retrievers=[bm25_retriever, vector_retriever],
-    weights=[0.4, 0.6]
-)
-
-print("Hybrid Retriever Ready ✅")
-
-
-# =====================================================
-# Session-based Memory
+# Session Memory
 # =====================================================
 class SessionConversationMemory:
 
     def __init__(self):
-        self.user_memories: Dict[str, ConversationBufferMemory] = {}
 
-    def get_memory(self, session_id: str):
+        self.user_memories: Dict[
+            str,
+            ConversationBufferMemory
+        ] = {}
+
+    # ===============================================
+    # Get User Memory
+    # ===============================================
+    def get_memory(
+        self,
+        session_id: str
+    ):
 
         if session_id not in self.user_memories:
 
-            self.user_memories[session_id] = ConversationBufferMemory(
+            self.user_memories[
+                session_id
+            ] = ConversationBufferMemory(
                 return_messages=True,
                 memory_key="chat_history"
             )
 
-        return self.user_memories[session_id]
+        return self.user_memories[
+            session_id
+        ]
 
-    def clear_history(self, session_id: str):
+    # ===============================================
+    # Clear User History
+    # ===============================================
+    def clear_history(
+        self,
+        session_id: str
+    ):
 
         if session_id in self.user_memories:
-            self.user_memories[session_id].clear()
+
+            self.user_memories[
+                session_id
+            ].clear()
 
 
 conversation = SessionConversationMemory()
 
 
 # =====================================================
-# Gemini Async Call with Retry
+# Async Gemini Call with Retry
 # =====================================================
 async def call_gemini_async(
     prompt: str,
@@ -141,13 +151,16 @@ async def call_gemini_async(
         try:
 
             response = await asyncio.to_thread(
+
                 client.models.generate_content,
+
                 model=MODEL_NAME,
                 contents=prompt
+
             )
 
-            # Streaming handling
             if response and response.text:
+
                 return response.text
 
             return "Empty response from model."
@@ -156,23 +169,36 @@ async def call_gemini_async(
 
             msg = str(e)
 
-            if "429" in msg or "RESOURCE_EXHAUSTED" in msg:
+            # =========================================
+            # Handle Quota Errors
+            # =========================================
+            if (
+                "429" in msg
+                or
+                "RESOURCE_EXHAUSTED" in msg
+            ):
 
                 wait_time = 5 * (attempt + 1)
 
                 logger.warning(
-                    f"Quota hit. Retrying in {wait_time}s..."
+                    f"Quota hit. "
+                    f"Retrying in {wait_time}s..."
                 )
 
                 await asyncio.sleep(wait_time)
 
                 continue
 
-            logger.error(f"Gemini Error: {e}")
+            logger.error(
+                f"Gemini Error: {e}"
+            )
 
             break
 
-    return "Model is temporarily unavailable. Please try again later."
+    return (
+        "Model is temporarily unavailable. "
+        "Please try again later."
+    )
 
 
 # =====================================================
@@ -185,18 +211,29 @@ def rerank_documents(
 
     scored_docs = []
 
-    query_words = set(query.lower().split())
+    query_words = set(
+        query.lower().split()
+    )
 
     for doc in docs:
 
-        content = doc.page_content.lower()
+        content = (
+            doc.page_content.lower()
+        )
 
         score = sum(
-            1 for word in query_words
+
+            1
+
+            for word in query_words
+
             if word in content
+
         )
 
-        scored_docs.append((score, doc))
+        scored_docs.append(
+            (score, doc)
+        )
 
     scored_docs.sort(
         key=lambda x: x[0],
@@ -204,23 +241,40 @@ def rerank_documents(
     )
 
     reranked_docs = [
-        doc for score, doc in scored_docs[:4]
+
+        doc
+
+        for score, doc in scored_docs[:4]
+
     ]
 
     return reranked_docs
 
 
 # =====================================================
-# Get Relevant Context
+# Retrieve Context
 # =====================================================
-def get_relevant_context_from_db(query):
+def get_relevant_context_from_db(
+    query: str
+):
 
     try:
 
-        docs = hybrid_retriever.get_relevant_documents(query)
+        # =============================================
+        # Vector Search
+        # =============================================
+        docs = vector_db.similarity_search(
+            query,
+            k=8
+        )
 
+        # =============================================
         # Reranking
-        docs = rerank_documents(query, docs)
+        # =============================================
+        docs = rerank_documents(
+            query,
+            docs
+        )
 
         context_parts = []
 
@@ -228,7 +282,9 @@ def get_relevant_context_from_db(query):
 
         for doc in docs:
 
-            content = doc.page_content[:1200]
+            content = (
+                doc.page_content[:1200]
+            )
 
             source = doc.metadata.get(
                 "source",
@@ -236,18 +292,25 @@ def get_relevant_context_from_db(query):
             )
 
             context_parts.append(
-                f"[SOURCE: {source}]\n{content}"
+
+                f"[SOURCE: {source}]\n"
+                f"{content}"
+
             )
 
             sources.add(source)
 
-        context = "\n\n".join(context_parts)
+        context = "\n\n".join(
+            context_parts
+        )
 
         return context, list(sources)
 
     except Exception as e:
 
-        logger.error(f"Retriever Error: {e}")
+        logger.error(
+            f"Retriever Error: {e}"
+        )
 
         return "", []
 
@@ -263,7 +326,10 @@ def build_prompt(
 
     history = memory.load_memory_variables(
         {}
-    ).get("chat_history", [])
+    ).get(
+        "chat_history",
+        []
+    )
 
     history_text = "\n".join([
 
@@ -316,31 +382,41 @@ async def generate_chat_answer(
     try:
 
         logger.info(
-            f"Session: {session_id} | Query: {query}"
+            f"Session: {session_id} "
+            f"| Query: {query}"
         )
 
-        memory = conversation.get_memory(session_id)
+        # =============================================
+        # Get User Memory
+        # =============================================
+        memory = conversation.get_memory(
+            session_id
+        )
 
-        # =========================================
+        # =============================================
         # Retrieve Context
-        # =========================================
-        context, sources = get_relevant_context_from_db(
-            query
+        # =============================================
+        context, sources = (
+            get_relevant_context_from_db(
+                query
+            )
         )
 
-        # =========================================
+        # =============================================
         # Build Prompt
-        # =========================================
+        # =============================================
         prompt = build_prompt(
             query=query,
             context=context,
             memory=memory
         )
 
-        # =========================================
-        # Call Gemini
-        # =========================================
-        answer = await call_gemini_async(prompt)
+        # =============================================
+        # Generate Answer
+        # =============================================
+        answer = await call_gemini_async(
+            prompt
+        )
 
         answer = answer.strip()
 
@@ -351,28 +427,39 @@ async def generate_chat_answer(
         )
 
         if not answer:
-            answer = "Sorry, no response generated."
 
-        # =========================================
-        # Save Memory
-        # =========================================
+            answer = (
+                "Sorry, no response generated."
+            )
+
+        # =============================================
+        # Save Conversation Memory
+        # =============================================
         memory.save_context(
+
             {"input": query},
+
             {"output": answer}
+
         )
 
-        # =========================================
+        # =============================================
         # Add Sources
-        # =========================================
+        # =============================================
         if sources:
 
-            unique_sources = list(set(sources))
+            unique_sources = list(
+                set(sources)
+            )
 
             answer += "\n\nSources:\n"
 
             answer += "\n".join([
+
                 f"- {source}"
+
                 for source in unique_sources
+
             ])
 
         return answer
@@ -388,11 +475,15 @@ async def generate_chat_answer(
 
 
 # =====================================================
-# Clear Session Chat
+# Clear Chat History
 # =====================================================
-def clear_chat_history(session_id: str):
+def clear_chat_history(
+    session_id: str
+):
 
-    conversation.clear_history(session_id)
+    conversation.clear_history(
+        session_id
+    )
 
 
 # =====================================================
@@ -405,4 +496,7 @@ def signal_handler(sig, frame):
     sys.exit(0)
 
 
-signal.signal(signal.SIGINT, signal_handler)
+signal.signal(
+    signal.SIGINT,
+    signal_handler
+)