Update app.py

app.py

@@ -1,6 +1,4 @@
-# ==============================================================================
-# SECTION 1: IMPORTS
-# ==============================================================================
+
 import os
 import json
 import re
@@ -8,238 +6,164 @@ 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
 
-# --- LangChain and related imports ---
+# Make sure you have these files in a 'utils' folder
+from utils.DocsLoader import load_and_chunk
+from utils.Schemas import RunRequest, RunResponse
+
 from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_huggingface import HuggingFaceEmbeddings # Correct new import
 from langchain_chroma import Chroma
 from langchain_community.retrievers import BM25Retriever
 from langchain.retrievers import EnsembleRetriever, ContextualCompressionRetriever
 from langchain.retrievers.document_compressors import CrossEncoderReranker
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
 from langchain.prompts import PromptTemplate
 
-# --- Local utility imports ---
-from utils.DocsLoader import load_and_chunk
-from utils.Schemas import RunRequest, RunResponse
 
 # Load environment variables
 load_dotenv()
 
-# ==============================================================================
-# SECTION 2: LIFESPAN MANAGER (MODEL LOADING BYPASSED FOR DEBUGGING)
-# ==============================================================================
+# --- 1. Lifespan Event Handler (The New, Correct Way) ---
+# This dictionary will hold our loaded models
 ml_models = {}
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """
-    Debug version: This function will bypass model loading to test server startup.
-    """
-    print("🚀 Lifespan startup: Bypassing model loading for this test.")
-
-    # --- ALL MODEL LOADING IS TEMPORARILY COMMENTED OUT ---
-    # GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY") 
-    # 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")
-    # 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-flash", api_key=GOOGLE_API_KEY)
-    # ml_models["prompt_template"] = PromptTemplate.from_template(
-    #     """Your full prompt here..."""
-    # )
-
-    print("✅ Bypassed model loading. The server should now start.")
-
-    yield
-
-    print("Shutting down.")
-
-# ==============================================================================
-# SECTION 3: FASTAPI APP INITIALIZATION
-# ==============================================================================
-app = FastAPI(title="Debug RAG Server", lifespan=lifespan)
-
-
-# ==============================================================================
-# SECTION 4 & 5: SECURITY AND API ENDPOINT
-# ==============================================================================
-# Note: The API endpoint will fail if you call it, because the models are not loaded.
-# We are only testing if the server can START.
-
-@app.get("/")
-def root():
-    return {"message": "Debug server is running! Model loading was bypassed."}
-
-
-# import os
-# import json
-# import re
-# 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
-
-# # Make sure you have these files in a 'utils' folder
-# from utils.DocsLoader import load_and_chunk
-# from utils.Schemas import RunRequest, RunResponse
-
-# from langchain_google_genai import ChatGoogleGenerativeAI
-# from langchain_huggingface import HuggingFaceEmbeddings # Correct new import
-# from langchain_chroma import Chroma
-# from langchain_community.retrievers import BM25Retriever
-# from langchain.retrievers import EnsembleRetriever, 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
-# ml_models = {}
-
-# @asynccontextmanager
-# async def lifespan(app: FastAPI):
-#     # This code runs ONCE when the application starts up
-#     print("🚀 Initializing models and prompt template...")
+    # This code runs ONCE when the application starts up
+    print("🚀 Initializing models and prompt template...")
     
-#     # Use a consistent key name that you set in Hugging Face Secrets
-#     GOOGLE_API_KEY = os.getenv("gemini_api_key") 
-#     if not GOOGLE_API_KEY:
-#         raise RuntimeError("CRITICAL: Missing GOOGLE_API_KEY in environment secrets!")
+    # Use a consistent key name that you set in Hugging Face Secrets
+    GOOGLE_API_KEY = os.getenv("gemini_api_key") 
+    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")
-#     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-flash", api_key=GOOGLE_API_KEY)
-#     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:
-# {context}
-
-# ### 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 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 a perfect and concise.
-# }}
-# Else:
-# {{
-#   "response": "Concise natural language answer"
-#   Make it a perfect and concise.
-# }}
-
-# ##NOTE : Do  not mention document id or its page number just mention clauses if applicable .
-#         """
-#     )
-#     print("✅ Models and prompt loaded successfully!")
+    # Load models into the shared dictionary
+    ml_models["embedder"] = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en-v1.5")
+    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-flash", api_key=GOOGLE_API_KEY)
+    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:
+{context}
+
+### 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 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 a perfect and concise.
+}}
+Else:
+{{
+  "response": "Concise natural language answer"
+  Make it a perfect and concise.
+}}
+
+##NOTE : Do  not mention document id or its page number just mention clauses if applicable .
+        """
+    )
+    print("✅ Models and prompt loaded successfully!")
     
-#     yield
+    yield
     
-#     # Code below yield runs on shutdown (optional)
-#     print("Shutting down and 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]
-#     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 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 f"Unstructured response:\n{content.strip()}"
-
-#     except Exception as e:
-#         return f"An error occurred while processing the response: {str(e)}"
-
-
-# # --- 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)
-
-#     # Create retrievers using the pre-loaded models from our ml_models dictionary
-#     keyword_retriever = BM25Retriever.from_documents(chunks)
-#     dense_retriever = Chroma.from_documents(documents=chunks, embedding=ml_models["embedder"]).as_retriever()
-#     ensemble_retriever = EnsembleRetriever(retrievers=[keyword_retriever, dense_retriever], weights=[0.3, 0.7])
+    # Code below yield runs on shutdown (optional)
+    print("Shutting down and 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]
+    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 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 f"Unstructured response:\n{content.strip()}"
+
+    except Exception as e:
+        return f"An error occurred while processing the response: {str(e)}"
+
+
+# --- 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)
+
+    # Create retrievers using the pre-loaded models from our ml_models dictionary
+    keyword_retriever = BM25Retriever.from_documents(chunks)
+    dense_retriever = Chroma.from_documents(documents=chunks, embedding=ml_models["embedder"]).as_retriever()
+    ensemble_retriever = EnsembleRetriever(retrievers=[keyword_retriever, dense_retriever], weights=[0.3, 0.7])
     
-#     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"]
-#     )
+    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"]
+    )
     
-#     answers = []
-#     for q in req.questions:
-#         try:
-#             result = await hybrid_rag_chain.ainvoke({"full_query": q})
-#             parsed = parse_llm_response(result.content)
-#             answers.append(parsed)
-#         except Exception as e:
-#             return JSONResponse({"error": str(e)}, status_code=500)
-
-#     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."}
+    answers = []
+    for q in req.questions:
+        try:
+            result = await hybrid_rag_chain.ainvoke({"full_query": q})
+            parsed = parse_llm_response(result.content)
+            answers.append(parsed)
+        except Exception as e:
+            return JSONResponse({"error": str(e)}, status_code=500)
+
+    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."}