Update app.py

app.py

@@ -1,51 +1,62 @@
+
 import os
 import json
 import re
 import asyncio
 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
 
-# Pinecone imports
-from pinecone import Pinecone
 
 # 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):
-    print("🚀 Initializing Pinecone and models...")
+    # This code runs ONCE when the application starts up
+    print("🚀 Initializing models and prompt template...")
 
     try:
         GOOGLE_API_KEY = os.getenv("gemini_api_key")
-        PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
-        
         print("🔑 gemini_api_key:", "FOUND" if GOOGLE_API_KEY else "NOT FOUND")
-        print("🔑 pinecone_api_key:", "FOUND" if PINECONE_API_KEY else "NOT FOUND")
 
-        if not GOOGLE_API_KEY or not PINECONE_API_KEY:
-            raise RuntimeError("CRITICAL: Missing API keys in environment!")
-        
-        # Initialize Pinecone client
-        ml_models["pc"] = Pinecone(api_key=PINECONE_API_KEY)
-        
-        # Initialize LLM
+        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)
-        
-        # Prompt template
         ml_models["prompt_template"] = PromptTemplate.from_template(
 """You are an expert insurance assistant. Your task is to answer the user's question as concisely as possible using ONLY the provided context.
 
@@ -59,8 +70,8 @@ async def lifespan(app: FastAPI):
 
     Concise Answer:
     """
-        )
-        print("✅ Pinecone and models loaded successfully!")
+    )
+        print("✅ Models and prompt loaded successfully!")
     except Exception as e:
         print("❌ Lifespan error:", str(e))
         raise e
@@ -68,196 +79,156 @@ async def lifespan(app: FastAPI):
     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)
 
-# API Key Verification
+# --- 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")
 
-def create_serverless_index(pinecone_client, index_name):
-    """Creates a Pinecone serverless index if it doesn't already exist."""
-    existing_indexes = [index_info.name for index_info in pinecone_client.list_indexes().indexes]
-    if index_name not in existing_indexes:
-        print(f"🪄 Creating serverless index: {index_name}")
-        pinecone_client.create_index(
-            name=index_name,
-            dimension=768,  # or 384 depending on your embedding model
-            metric='cosine',
-            spec={
-                'serverless': {
-                    'cloud': 'aws',
-                    'region': 'us-east-1'   # ✅ Required for Free Tier
-                }
-            }
-        )
-        print("✅ Index created.")
-    else:
-        print(f"⚠️ Index '{index_name}' already exists.")
-    return pinecone_client.Index(index_name)
-
-
-async def embed_with_pinecone_inference(pc, texts, model="multilingual-e5-large"):
-    """Use Pinecone's hosted inference model for embeddings."""
+
+# --- 4. Parsing Helper ---
+def parse_llm_response(content: str) -> str:
     try:
-        inference_client = pc.inference
-        embeddings = inference_client.embed(
-            model=model,
-            inputs=texts,
-            parameters={"input_type": "passage"}
-        )
-        return [item["values"] for item in embeddings["results"]]
+        # 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:
-        print(f"❌ Pinecone Inference Error: {e}")
-        raise
+        return f"An error occurred while processing the response: {str(e)}"
 
-def generate_sparse_vectors(texts):
-    """Generate sparse vectors using simple term frequency"""
-    from collections import Counter
-    import re
-    
-    sparse_vectors = []
-    for text in texts:
-        # Simple tokenization and term frequency
-        tokens = re.findall(r'\b\w+\b', text.lower())
-        token_counts = Counter(tokens)
-        
-        # Create vocabulary mapping (simplified)
-        vocab = {token: i for i, token in enumerate(set(tokens))}
-        
-        indices = []
-        values = []
-        for token, count in token_counts.items():
-            if token in vocab:
-                indices.append(vocab[token])
-                values.append(float(count))
-        
-        sparse_vectors.append({
-            'indices': indices,
-            'values': values
-        })
-    
-    return sparse_vectors
-
-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)
-
-    query_embedding = np.array(query_embedding).reshape(1, -1)
-    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
 
+# --- 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)
 
-    # Initialize Pinecone index
-    index_name = "insurance-hybrid-search"
-    pc = ml_models["pc"]
-    index = create_serverless_index(pc, index_name)
+        
+    ####code for parallel####################################################################################
 
-    # Prepare document texts
-    doc_texts = [doc.page_content for doc in chunks]
+    vectorstore = await FAISS.afrom_documents(
+        documents=chunks, 
+        embedding=ml_models["embedder"]
+    )
+
+    dense_retriever = vectorstore.as_retriever(search_type="mmr",search_kwargs={"k": 8})
     
-    # Generate embeddings using Pinecone's multilingual-e5-large
-    dense_embeddings = await embed_with_pinecone_inference(pc, doc_texts)
+
+    # 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"]
+    # )
+
     
-    # Generate sparse vectors
-    sparse_embeddings = generate_sparse_vectors(doc_texts)
+    # 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)
     
-    # Upsert to Pinecone with hybrid vectors
-    vectors_to_upsert = []
-    for i, (doc, dense_emb, sparse_emb) in enumerate(zip(chunks, dense_embeddings, sparse_embeddings)):
-        vectors_to_upsert.append({
-            'id': f'doc_{i}',
-            'values': dense_emb,
-            'sparse_values': sparse_emb,
-            'metadata': {
-                'text': doc.page_content,
-                'source': getattr(doc, 'metadata', {}).get('source', 'unknown')
-            }
-        })
+        # Convert query_embedding to 2D
+        query_embedding = np.array(query_embedding).reshape(1, -1)
     
-    # Batch upsert
-    batch_size = 100
-    for i in range(0, len(vectors_to_upsert), batch_size):
-        batch = vectors_to_upsert[i:i+batch_size]
-        index.upsert(vectors=batch)
+        # Compute similarity between query and all documents
+        query_doc_sims = cosine_similarity(query_embedding, doc_embeddings)[0]
     
-    print(f"✅ Upserted {len(vectors_to_upsert)} vectors to Pinecone")
-
-    # Generate question embeddings
-    question_embeddings = await embed_with_pinecone_inference(pc, req.questions)
-    question_sparse = generate_sparse_vectors(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)
+    
+        return selected
     async def async_retrieve_and_rerank(question: str, q_idx: int):
-        # Hybrid search with alpha=0.5 (balanced dense/sparse)
-        alpha = 0.5
-        
-        # Scale embeddings for hybrid search
-        dense_query = [v * alpha for v in question_embeddings[q_idx]]
-        sparse_query = {
-            'indices': question_sparse[q_idx]['indices'],
-            'values': [v * (1 - alpha) for v in question_sparse[q_idx]['values']]
-        }
-        
-        # Query Pinecone with hybrid search
-        results = index.query(
-            vector=dense_query,
-            sparse_vector=sparse_query,
-            top_k=8,
-            include_metadata=True
-        )
-        
-        # Extract contexts and embeddings for MMR
-        contexts = [match['metadata']['text'] for match in results['matches']]
-        context_embeddings = [match['values'] for match in results['matches']]
-        
-        # Apply MMR selection
+        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=context_embeddings,
-            k=6,
-            lambda_mult=0.6,
-        )
-        
-        top_chunks = [contexts[j] for j in selected_indices]
+        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 questions in parallel
+    # 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)]
-    )
+    ) 
+
+    ####################################################################################################################
 
-    # Generate answers
     tasks = []
     for i in range(len(req.questions)):
         prompt_input = {
@@ -265,16 +236,23 @@ async def run_hackrx(req: RunRequest):
             "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
+    # 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."}
+    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)