Update app.py

app.py

@@ -1,3 +1,4 @@
+
 import os
 import json
 import re
@@ -6,32 +7,38 @@ from contextlib import asynccontextmanager
 from dotenv import load_dotenv
 from operator import itemgetter
 
+# import gradio as gr
 from fastapi import FastAPI, Depends, HTTPException, Header
 from fastapi.responses import JSONResponse
 
 from utils.DocsLoader import load_and_chunk
 from utils.Schemas import RunRequest, RunResponse
-
+# from concurrent.futures import ThreadPoolExecutor
 from langchain_community.vectorstores import FAISS
 from langchain.schema import Document
 from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_huggingface import HuggingFaceEmbeddings 
+# from langchain_chroma import Chroma
 from langchain_community.retrievers import BM25Retriever
 from langchain.retrievers import EnsembleRetriever
 from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
 
 from langchain.retrievers import ContextualCompressionRetriever
-from langchain.retrievers.document_compressors import CrossEncoderReranker
+from langchain.retrievers.document_compressors import CrossEncoderReranker 
 from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 from langchain.prompts import ChatPromptTemplate
 
+
+# Load environment variables
 load_dotenv()
 
+
 ml_models = {}
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
+    # This code runs ONCE when the application starts up
     print("🚀 Initializing models and prompt template...")
 
     try:
@@ -40,21 +47,24 @@ async def lifespan(app: FastAPI):
 
         if not GOOGLE_API_KEY:
             raise RuntimeError("CRITICAL: Missing GOOGLE_API_KEY in environment secrets!")
-
-        ml_models["embedder"] = HuggingFaceEmbeddings(
-            model_name="BAAI/bge-base-en-v1.5",
-            encode_kwargs={"batch_size": 64}
-        )
-
+    
+        # Load models into the shared dictionary
+        ml_models["embedder"] = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en-v1.5",
+                                                      encode_kwargs={
+                                                          "batch_size": 64,
+                                                          # "normalize_embeddings": True
+                                                      }
+                                                    
+    )
         cross_encoder_model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base")
+        # cross_encoder_model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-large")
         ml_models["reranker_compressor"] = CrossEncoderReranker(model=cross_encoder_model, top_n=5)
         ml_models["llm"] = ChatGoogleGenerativeAI(
-            model="gemini-1.5-pro",
-            api_key=GOOGLE_API_KEY,
-            temperature=0.1,
-            max_output_tokens=300
-        )
-
+                               model="gemini-1.5-pro",
+                               api_key=GOOGLE_API_KEY,
+                               temperature=0.1,
+                               max_output_tokens=300
+                           )
         ml_models["prompt_template"] = ChatPromptTemplate.from_template("""
 **Role**: You are an expert decision maker Assistant in the domain such as insurance, legal compliance, human resources, and contract management using ONLY the provided context.
 Do not use your own knowledge.
@@ -80,7 +90,6 @@ Example how to answer for thr query:
 
 """
 )
-
         print("✅ Models and prompt loaded successfully!")
     except Exception as e:
         print("❌ Lifespan error:", str(e))
@@ -89,9 +98,11 @@ Example how to answer for thr query:
     yield
     print("🧹 Cleaning up.")
     ml_models.clear()
-
+# --- 2. FastAPI App Instance ---
+# We pass the lifespan function to the FastAPI constructor
 app = FastAPI(title="HackRX RAG Server", lifespan=lifespan)
 
+# --- 3. API Key Verification ---
 TEAM_API_KEY = os.getenv("TEAM_API_KEY")
 
 def verify_api_key(authorization: str = Header(...)):
@@ -101,8 +112,11 @@ def verify_api_key(authorization: str = Header(...)):
     if token != TEAM_API_KEY:
         raise HTTPException(status_code=403, detail="Invalid or missing API key")
 
+
+# --- 4. Parsing Helper ---
 def parse_llm_response(content: str) -> str:
     try:
+        # Remove code fences and clean up
         content_cleaned = re.sub(r"^```json|```$", "", content.strip(), flags=re.IGNORECASE).strip()
         data = json.loads(content_cleaned)
 
@@ -112,54 +126,54 @@ def parse_llm_response(content: str) -> str:
                 amount = data.get("amount", "Not specified")
                 justification = data.get("justification", "No justification provided.")
                 return f"Decision: {decision}\nAmount: {amount}\nJustification: {justification}"
+
             elif "response" in data:
                 return data["response"]
 
         return "The response was parsed but didn't match expected structure."
+
     except json.JSONDecodeError:
         return f"Unstructured response:\n{content.strip()}"
+
     except Exception as e:
         return f"An error occurred while processing the response: {str(e)}"
 
-# --- Post-retrieval Filters ---
-def structure_filter(docs, min_length=30):
-    return [doc for doc in docs if len(doc.page_content.strip()) >= min_length and not doc.page_content.isspace()]
-
-def semantic_filter(query, docs, embedder, threshold=0.45):
-    query_emb = embedder.embed_query(query)
-    doc_texts = [doc.page_content for doc in docs]
-    doc_embs = embedder.embed_documents(doc_texts)
-    sims = cosine_similarity([query_emb], doc_embs)[0]
-    return [doc for doc, sim in zip(docs, sims) if sim >= threshold]
-
-def universal_filter(query, docs, embedder):
-    docs = structure_filter(docs)
-    docs = semantic_filter(query, docs, embedder)
-    return docs if docs else docs  # fallback to raw if filtering removes all
 
+# --- 5. Main API Endpoint ---
 @app.post("/api/v1/hackrx/run", response_model=RunResponse, dependencies=[Depends(verify_api_key)])
 async def run_hackrx(req: RunRequest):
     chunks = load_and_chunk(str(req.documents))
     if not chunks:
         return JSONResponse({"error": "No documents could be processed."}, status_code=400)
 
-    vectorstore = await FAISS.afrom_documents(documents=chunks, embedding=ml_models["embedder"])
-    dense_retriever = vectorstore.as_retriever(search_type="mmr", search_kwargs={"k": 8, "lambda_mult": 0.5})
+    vectorstore = await FAISS.afrom_documents(
+        documents=chunks, 
+        embedding=ml_models["embedder"]
+    )
+
+    # dense_retriever = vectorstore.as_retriever(search_type="mmr",search_kwargs={"k": 8})
+    dense_retriever = vectorstore.as_retriever(search_type="mmr",search_kwargs={"k": 8 ,"lambda_mult": 0.5})
+    
+
+    # Create retrievers using the pre-loaded models from our ml_models dictionary
     keyword_retriever = BM25Retriever.from_documents(chunks)
     keyword_retriever.k = 5
+    # dense_retriever = Chroma.from_documents(documents=chunks, embedding=ml_models["embedder"]).as_retriever()
     ensemble_retriever = EnsembleRetriever(retrievers=[keyword_retriever, dense_retriever], weights=[0.35, 0.65])
-
-    tasks = []
-    for query in req.questions:
-        raw_docs = await ensemble_retriever.ainvoke(query)
-        filtered_docs = universal_filter(query, raw_docs, ml_models["embedder"])
-        if not filtered_docs:
-            filtered_docs = raw_docs  # fallback
-        reranked_docs = ml_models["reranker_compressor"].compress_documents(filtered_docs, query=query)
-        context = "\n\n".join([doc.page_content for doc in reranked_docs])
-        prompt_input = {"full_query": query, "context": context}
-        tasks.append(ml_models["llm"].ainvoke(ml_models["prompt_template"].format_prompt(**prompt_input)))
-
+    ### to make it faster we are now using our built reranker thats why commenting the code below
+    compression_retriever = ContextualCompressionRetriever(
+        base_retriever=ensemble_retriever, base_compressor=ml_models["reranker_compressor"]
+    )
+
+    
+    # Define the RAG chain using pre-loaded components
+    hybrid_rag_chain = (
+        {"context": itemgetter("full_query") | compression_retriever, "full_query": itemgetter("full_query")}
+        | ml_models["prompt_template"]
+        | ml_models["llm"]
+    )
+
+    tasks = [hybrid_rag_chain.ainvoke({"full_query": q}) for q in req.questions]
     results = await asyncio.gather(*tasks)
     # answers = [parse_llm_response(result.content) for result in results]
     answers = []
@@ -167,7 +181,8 @@ async def run_hackrx(req: RunRequest):
     for msg in results:
         # Safely access the content field
         if hasattr(msg, "content"):
-            answers.append(msg.content.strip())
+            answers.append(msg.content.strip())    
+
     return JSONResponse({"answers": answers}, status_code=200)
 
 @app.get("/", include_in_schema=False)