Update app.py

app.py

@@ -1,113 +1,88 @@
-
 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 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
 
-
 # Load environment variables
 load_dotenv()
 
-# --- 1. Lifespan Event Handler (The New, Correct Way) ---
-# This dictionary will hold our loaded models
+# Cache setup
+document_cache = TTLCache(maxsize=5, ttl=300)
+cache_lock = asyncio.Lock()
+
+# --- 1. Lifespan Event Handler ---
 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!")
     
-        # Load models into the shared dictionary
-        ml_models["embedder"] = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en-v1.5")
-        ### 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-2.0-flash", api_key=GOOGLE_API_KEY)
+        # Optimized embedding model with batching
+        ml_models["embedder"] = HuggingFaceEmbeddings(
+            model_name="BAAI/bge-base-en-v1.5",
+            encode_kwargs={
+                'batch_size': 64,
+                'show_progress_bar': False
+            }
+        )
+        
+        # Faster LLM with constrained output
+        ml_models["llm"] = ChatGoogleGenerativeAI(
+            model="gemini-1.0-pro",
+            api_key=GOOGLE_API_KEY,
+            temperature=0.1,
+            max_output_tokens=300
+        )
+        
+        # Improved prompt template
         ml_models["prompt_template"] = PromptTemplate.from_template("""
-You are an expert decision maker Assistant in the domain such as insurance, legal compliance, human resources, and contract management.
-
-The customer has submitted a query. If the query has details about age , gender ,procedure ,location , policy duration, then parse it and understand the query properly
-If not, the raw question is provided instead:
-- Query: {full_query}
-
-We retrieved the following policy clauses and rules relevant to this case:
+**Role**: Insurance Policy Expert
+**Context**: 
 {context}
 
-Remember the below three points while answering or responsing the query :
-1:Do  not mention document id or its page number just ,mention clauses if applicable .
-2:Do not make such statement in the response that you need more document or information to answer the query, just answer from retrieved data.
-3:UNDERSTAND THE QUERY ASKED SMARTLY AND ANSWER IT
-4.Do not add any irrelevent information.
-### Task:
-If the details about age , gender , procedure , location , policy duration are available, do all of the following:
-1. Decide whether the procedure is covered.
-2. Estimate the claimable amount.
-3. Justify with the relevant clause.
-and answer the query  precisely as insurance agent.
-
-Otherwise, answer the question concisely , relevently and clearly using the retrieved context.
+**Query**: {full_query}
 
-### Output format:
-If query involes age , gender , procedure , location , policy duration answer like below:
-{{
-  "decision": "approved / rejected",
-  "amount": "INR amount or null",
-  "justification": "Refer to specific clause"
-  Make it  concise ,perfect, relevent and clear .
-}}
-Else:
-{{
-  "response": "Concise natural language answer"
-   Make it  concise ,perfect, relevent and clear .
-}}
+**Instructions**:
+1. If query contains age/gender/procedure/location/duration:
+   - 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("❌ Lifespan error:", str(e))
+        print(f"❌ 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 ---
@@ -117,171 +92,99 @@ 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
+        # Clean JSON response
         content_cleaned = re.sub(r"^```json|```$", "", content.strip(), flags=re.IGNORECASE).strip()
         data = json.loads(content_cleaned)
-
-        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."
-
+        
+        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"
     except json.JSONDecodeError:
-        return f"Unstructured response:\n{content.strip()}"
-
+        return content.strip()
     except Exception as e:
-        return f"An error occurred while processing the response: {str(e)}"
-
+        return f"Response processing error: {str(e)}"
+
+# --- 5. Retrieval Optimization ---
+async def get_relevant_docs(question: str, vectorstore: FAISS, keyword_retriever: BM25Retriever):
+    # Parallel retrieval
+    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
+        )
+    )
+    
+    # Combine and deduplicate
+    all_docs = dense_docs + sparse_docs
+    unique_docs = {doc.page_content: doc for doc in all_docs}.values()
+    
+    # Fast reranking
+    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]
 
-# --- 5. Main API Endpoint ---
+# --- 6. 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)
-
-        
-    ####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})
+    start_time = time.time()
     
-
-    # 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.4, 0.6])
-    ### 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"]
-    )
-
+    # Cache document processing
+    async with cache_lock:
+        if req.documents in document_cache:
+            print("♻️ Using cached document")
+            chunks = document_cache[req.documents]
+        else:
+            chunks = load_and_chunk(str(req.documents))
+            document_cache[req.documents] = chunks
     
-    # 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]
+    if not chunks:
+        return JSONResponse({"error": "No documents processed"}, status_code=400)
 
-    #     top_k = 5
-    #     top_indices = np.argsort(sims)[-top_k:][::-1]
-    #     top_chunks = [doc_texts[j] for j in top_indices]
+    # Create vector store and retrievers
+    vectorstore = await FAISS.afrom_documents(chunks, ml_models["embedder"])
+    keyword_retriever = BM25Retriever.from_documents(chunks)
+    keyword_retriever.k = 6
 
-    #     # 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)
-    
-    #     # Convert query_embedding to 2D
-    #     query_embedding = np.array(query_embedding).reshape(1, -1)
+    # 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])
+        
+        # Generate response
+        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)
     
-    #     # Compute similarity between query and all documents
-    #     query_doc_sims = cosine_similarity(query_embedding, doc_embeddings)[0]
+    answers = await asyncio.gather(*(process_question(q) for q in req.questions))
     
-    #     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)
+    # Performance logging
+    proc_time = time.time() - start_time
+    print(f"⏱️ Processed {len(req.questions)} questions in {proc_time:.2f}s")
     
-    #     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]
-
     return JSONResponse({"answers": answers}, status_code=200)
 
 @app.get("/", include_in_schema=False)
 def root():
-    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)
+    return {"message": "HackRX API operational. Use /api/v1/hackrx/run"}