Update app.py

app.py

@@ -1,90 +1,103 @@
+
 import os
 import json
 import re
 import asyncio
-import time
-import numpy as np
-from cachetools import TTLCache
 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 langchain_core.documents import Document
+
 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_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
+
+### to make it faster we are now using our built reranker thats why commenting the imports below
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import CrossEncoderReranker 
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 from langchain.prompts import PromptTemplate
-from fastapi.middleware.cors import CORSMiddleware  # Optional, for external frontend
+
 
 # Load environment variables
 load_dotenv()
 
-# Cache setup
-document_cache = TTLCache(maxsize=5, ttl=300)
-cache_lock = asyncio.Lock()
-
-# --- 1. Lifespan Event Handler ---
+# --- 1. Lifespan Event Handler (The New, Correct Way) ---
+# This dictionary will hold our loaded models
 ml_models = {}
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
+    # This code runs ONCE when the application starts up
     print("🚀 Initializing models and prompt template...")
 
     try:
         GOOGLE_API_KEY = os.getenv("gemini_api_key")
+        print("🔑 gemini_api_key:", "FOUND" if GOOGLE_API_KEY else "NOT FOUND")
+
         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,
-                'show_progress_bar': False
-            }
-        )
-        
+        # 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
+                                                      }
+                                                    show_progress_bar= False
+    )
+        ### to make it faster we are now using our built reranker thats why commenting the code below
+        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"] = PromptTemplate.from_template("""
-**Role**: Insurance Policy Expert
+**Role**: You are an expert decision maker Assistant in the domain such as insurance, legal compliance, human resources, and contract management.
+
 **Context**: 
 {context}
 
 **Query**: {full_query}
 
 **Instructions**:
-1. If query contains age/gender/procedure/location/duration:
+1. If query contains age ,gender,procedure,duration , location ,and query like can i get a knee surgery ,if i am male and duration is 3months or any query is similar to the previus example :
    - Output ONLY JSON: {{"decision":"approved/rejected","amount":"₹X","justification":"Clause reference"}}
 2. Else: Provide concise answer
 3. NEVER mention document sources
 4. If unsure, respond: "Insufficient information"
 
 **Response**:
-""")
+"""
+    )
         print("✅ Models and prompt loaded successfully!")
     except Exception as e:
-        print(f"❌ Lifespan error: {str(e)}")
+        print("❌ Lifespan error:", str(e))
         raise e
 
     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")
 
@@ -92,95 +105,178 @@ def verify_api_key(authorization: str = Header(...)):
     if not authorization.startswith("Bearer "):
         raise HTTPException(status_code=401, detail="Invalid Authorization header format")
     token = authorization.split("Bearer ")[1]
+    # print(token)
+    # print(TEAM_API_KEY)
     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)
-        
-        if "decision" in data:
-            return (
-                f"Decision: {data.get('decision', 'N/A').upper()}\n"
-                f"Amount: {data.get('amount', 'Not specified')}\n"
-                f"Justification: {data.get('justification', 'No justification provided')}"
-            )
-        elif "response" in data:
-            return data["response"]
-        return "Response format error"
+
+        if isinstance(data, dict):
+            if "decision" in data:
+                decision = data.get("decision", "N/A").upper()
+                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 content.strip()
+        return f"Unstructured response:\n{content.strip()}"
+
     except Exception as e:
-        return f"Response processing error: {str(e)}"
-
-# --- 5. Retrieval Optimization ---
-async def get_relevant_docs(question: str, vectorstore: FAISS, keyword_retriever: BM25Retriever):
-    dense_docs, sparse_docs = await asyncio.gather(
-        vectorstore.asimilarity_search(question, k=6),
-        asyncio.get_event_loop().run_in_executor(
-            None, 
-            keyword_retriever.get_relevant_documents, 
-            question
-        )
-    )
-    
-    all_docs = dense_docs + sparse_docs
-    unique_docs = {doc.page_content: doc for doc in all_docs}.values()
-    
-    query_embedding = ml_models["embedder"].embed_query(question)
-    doc_texts = [doc.page_content for doc in unique_docs]
-    doc_embeddings = ml_models["embedder"].embed_documents(doc_texts)
-    
-    similarities = cosine_similarity([query_embedding], doc_embeddings)[0]
-    sorted_indices = np.argsort(similarities)[::-1][:5]  # Top 5
-    
-    return [list(unique_docs)[i] for i in sorted_indices]
+        return f"An error occurred while processing the response: {str(e)}"
+
 
-# --- 6. Main API Endpoint ---
+# --- 5. Main API Endpoint ---
 @app.post("/api/v1/hackrx/run", response_model=RunResponse, dependencies=[Depends(verify_api_key)])
 async def run_hackrx(req: RunRequest):
-    start_time = time.time()
-    
-    async with cache_lock:
-        doc_key = hash(str(req.documents))  # Hash key for safety
-        if doc_key in document_cache:
-            print("♻️ Using cached document")
-            chunks = document_cache[doc_key]
-        else:
-            raw_chunks = load_and_chunk(str(req.documents))
-            # 🔧 FIXED: Convert string chunks to Document objects
-            chunks = [Document(page_content=chunk) for chunk in raw_chunks]
-            document_cache[doc_key] = chunks
-    
+    chunks = load_and_chunk(str(req.documents))
     if not chunks:
-        return JSONResponse({"error": "No documents processed"}, status_code=400)
+        return JSONResponse({"error": "No documents could be processed."}, status_code=400)
+
+        
+    ####code for parallel####################################################################################
+
+    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.6})
+    
 
-    # Create vector store and retrievers
-    vectorstore = await FAISS.afrom_documents(chunks, ml_models["embedder"])
+    # Create retrievers using the pre-loaded models from our ml_models dictionary
     keyword_retriever = BM25Retriever.from_documents(chunks)
-    keyword_retriever.k = 6
+    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])
+    ### 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"]
+    )
 
-    # Process questions in parallel
-    async def process_question(q: str):
-        relevant_docs = await get_relevant_docs(q, vectorstore, keyword_retriever)
-        context = "\n".join([d.page_content for d in relevant_docs])
-        
-        prompt = ml_models["prompt_template"].format_prompt(
-            full_query=q, 
-            context=context
-        )
-        result = await ml_models["llm"].ainvoke(prompt)
-        return parse_llm_response(result.content)
     
-    answers = await asyncio.gather(*(process_question(q) for q in req.questions))
+    # 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"]
+    )
+
+    ######## OUR SELF RERANKER ######################################################################
+    #Embed all questions at once
+    # question_embeddings = ml_models["embedder"].embed_documents(req.questions)
+
+    # For each question, retrieve and rerank with cosine
+    # retrieved_chunks_all = []
+    # for i, question in enumerate(req.questions):
+    #     docs = ensemble_retriever.get_relevant_documents(question)
+    #     doc_texts = [doc.page_content for doc in docs]
+
+    #     doc_embeddings = ml_models["embedder"].embed_documents(doc_texts)
+    #     sims = cosine_similarity([question_embeddings[i]], doc_embeddings)[0]
+
+    #     top_k = 5
+    #     top_indices = np.argsort(sims)[-top_k:][::-1]
+    #     top_chunks = [doc_texts[j] for j in top_indices]
+
+    #     # Join for context
+    #     joined_context = "\n\n".join(top_chunks)
+    #     retrieved_chunks_all.append(joined_context)
+    # def mmr_select(query_embedding, doc_embeddings, k=6, lambda_mult=0.6):
+    #     selected = []
+    #     candidates = list(range(len(doc_embeddings)))
+    #     doc_embeddings = np.array(doc_embeddings)
     
-    proc_time = time.time() - start_time
-    print(f"⏱️ Processed {len(req.questions)} questions in {proc_time:.2f}s")
+    #     # Convert query_embedding to 2D
+    #     query_embedding = np.array(query_embedding).reshape(1, -1)
     
+    #     # Compute similarity between query and all documents
+    #     query_doc_sims = cosine_similarity(query_embedding, doc_embeddings)[0]
+    
+    #     for _ in range(k):
+    #         mmr_score = []
+    #         for idx in candidates:
+    #             if not selected:
+    #                 diversity = 0
+    #             else:
+    #                 selected_embeddings = doc_embeddings[selected]
+    #                 diversity = max(cosine_similarity(
+    #                     doc_embeddings[idx].reshape(1, -1),
+    #                     selected_embeddings
+    #                 )[0])
+    #             score = lambda_mult * query_doc_sims[idx] - (1 - lambda_mult) * diversity
+    #             mmr_score.append(score)
+    #         selected_idx = candidates[np.argmax(mmr_score)]
+    #         selected.append(selected_idx)
+    #         candidates.remove(selected_idx)
+    
+    #     return selected
+    # async def async_retrieve_and_rerank(question: str, q_idx: int):
+    #     docs = await ensemble_retriever.ainvoke(question)
+    #     doc_texts = [doc.page_content for doc in docs]
+    #     doc_embeddings = ml_models["embedder"].embed_documents(doc_texts)
+    #     # sims = cosine_similarity([question_embeddings[q_idx]], doc_embeddings)[0]
+    #     query_embedding = question_embeddings[q_idx]
+    #     selected_indices = mmr_select(
+    #     query_embedding=query_embedding,
+    #     doc_embeddings=doc_embeddings,
+    #     k=6,
+    #     lambda_mult=0.6,
+    # )
+    #     # top_indices = np.argsort(sims)[-top_k:][::-1]
+    #     top_chunks = [doc_texts[j] for j in selected_indices]
+    #     return "\n\n".join(top_chunks)
+    # # Retrieve and rerank all in parallel
+    # retrieved_chunks_all = await asyncio.gather(
+    #     *[async_retrieve_and_rerank(q, i) for i, q in enumerate(req.questions)]
+    # ) 
+
+    ####################################################################################################################
+
+    # tasks = []
+    # for i in range(len(req.questions)):
+    #     prompt_input = {
+    #         "full_query": req.questions[i],
+    #         "context": retrieved_chunks_all[i]
+    #     }
+        # tasks.append(ml_models["llm"].ainvoke(ml_models["prompt_template"].format_prompt(**prompt_input)))
+    # # tasks = [hybrid_rag_chain.ainvoke({"full_query": q}) for q in req.questions]
+    # results = await asyncio.gather(*tasks)
+    # answers = []
+
+    # for msg in results:
+    #     # Safely access the content field
+    #     if hasattr(msg, "content"):
+    #         answers.append(msg.content.strip())
+    #     # Extract the content from each result and parse it
+    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 = []
+
+    for msg in results:
+        # Safely access the content field
+        if hasattr(msg, "content"):
+            response.append(msg.content.strip())    
+
     return JSONResponse({"answers": answers}, status_code=200)
 
 @app.get("/", include_in_schema=False)
 def root():
-    return {"message": "HackRX API operational. Use /api/v1/hackrx/run"}
+    return {"message": "API is running. Go to /docs for documentation."}
+
+
+# def dummy_gradio(): return "✅ API running!"
+# gr.Interface(fn=dummy_gradio, inputs=[], outputs="text").launch(server_name="0.0.0.0", server_port=7860)