Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import os
 import json
 import re
@@ -7,40 +6,32 @@ 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_huggingface import HuggingFaceEmbeddings
 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.retrievers.document_compressors import CrossEncoderReranker
 from langchain_community.cross_encoders import HuggingFaceCrossEncoder
-from langchain.prompts import PromptTemplate
+from langchain.prompts import ChatPromptTemplate
 
-
-# Load environment variables
 load_dotenv()
 
-# --- 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:
@@ -49,63 +40,47 @@ async def lifespan(app: FastAPI):
 
         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",
-                                                      encode_kwargs={
-                                                          "batch_size": 64,
-                                                          # "normalize_embeddings": True
-                                                      }
-                                                    
-    )
-        ### to make it faster we are now using our built reranker thats why commenting the code below
+
+        ml_models["embedder"] = HuggingFaceEmbeddings(
+            model_name="BAAI/bge-base-en-v1.5",
+            encode_kwargs={"batch_size": 64}
+        )
+
         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
-                           )
-        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.
-Do not add any irrelevent information, also do not use bullet points, make it in a perfect sentence answer.
-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:
+            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.
+**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.
-### 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.
-
-### 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 .
-}}
+**Query**: {full_query}
+
+**Response**:
+
+Example how to answer for thr query:
+1. query:What is the grace period for premium payment under the National Parivar Mediclaim Plus Policy?
+   response: A grace period of thirty days is provided for premium payment after the due date to renew or continue the policy without losing continuity benefits.
+
+2. query: What is the waiting period for pre-existing diseases (PED) to be covered?
+   response: There is a waiting period of thirty-six (36) months of continuous coverage from the first policy inception for pre-existing diseases and their direct complications to be covered.
+
+3. query: Are the medical expenses for an organ donor covered under this policy?
+   response : Yes, the policy indemnifies the medical expenses for the organ donor's hospitalization for the purpose of harvesting the organ, provided the organ is for an insured person and the donation complies with the Transplantation of Human Organs Act, 1994.
+
+4. query: Does this policy cover maternity expenses, and what are the conditions?
+   response: Yes, the policy covers maternity expenses, including childbirth and lawful medical termination of pregnancy. To be eligible, the female insured person must have been continuously covered for at least 24 months. The benefit is limited to two deliveries or terminations during the policy period.
 
 """
-    )
+)
+
         print("✅ Models and prompt loaded successfully!")
     except Exception as e:
         print("❌ Lifespan error:", str(e))
@@ -114,27 +89,20 @@ Else:
     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(...)):
     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)
 
@@ -144,159 +112,58 @@ 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)
 
-        
-    ####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.5})
-    
-
-    # Create retrievers using the pre-loaded models from our ml_models dictionary
+    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})
     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])
-    ### 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"]
-    )
-
-    ######## 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)
-    
-    #     # 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"):
-    #         answers.append(msg.content.strip())    
+    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)))
 
+    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)