Update app.py

app.py

@@ -1,62 +1,51 @@
-
 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):
-    # This code runs ONCE when the application starts up
-    print("🚀 Initializing models and prompt template...")
+    print("🚀 Initializing Pinecone and models...")
 
     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:
-            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)
+        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
         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.
 
@@ -70,8 +59,8 @@ async def lifespan(app: FastAPI):
 
     Concise Answer:
     """
-    )
-        print("✅ Models and prompt loaded successfully!")
+        )
+        print("✅ Pinecone and models loaded successfully!")
     except Exception as e:
         print("❌ Lifespan error:", str(e))
         raise e
@@ -79,156 +68,190 @@ 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)
 
-# --- 3. API Key Verification ---
+# 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(pc, index_name):
+    """Create a serverless index for hybrid search"""
+    existing_indexes = [index.name for index in pc.list_indexes()]
+    
+    if index_name not in existing_indexes:
+        pc.create_index(
+            name=index_name,
+            dimension=1024,  # multilingual-e5-large dimension
+            metric="cosine"
+        )
+        print(f"✅ Created serverless index: {index_name}")
+    else:
+        print(f"📋 Index {index_name} already exists")
+    
+    return pc.Index(index_name)
 
-# --- 4. Parsing Helper ---
-def parse_llm_response(content: str) -> str:
+async def embed_with_pinecone_inference(pc, texts, model="multilingual-e5-large"):
+    """Use Pinecone's hosted multilingual-e5-large model for embeddings"""
     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()}"
-
+        # Use Pinecone Inference for embeddings
+        embeddings = pc.inference.embed(
+            model=model,
+            inputs=texts,
+            parameters={"input_type": "passage"}
+        )
+        return [embedding['values'] for embedding in embeddings]
     except Exception as e:
-        return f"An error occurred while processing the response: {str(e)}"
+        print(f"❌ Embedding error: {e}")
+        raise
 
+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)
 
-        
-    ####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})
-    
-
-    # 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"]
-    # )
+    # Initialize Pinecone index
+    index_name = "insurance-hybrid-search"
+    pc = ml_models["pc"]
+    index = create_serverless_index(pc, index_name)
 
+    # Prepare document texts
+    doc_texts = [doc.page_content for doc in 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]
-
-    #     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)
+    # Generate embeddings using Pinecone's multilingual-e5-large
+    dense_embeddings = await embed_with_pinecone_inference(pc, doc_texts)
     
-        # Convert query_embedding to 2D
-        query_embedding = np.array(query_embedding).reshape(1, -1)
+    # Generate sparse vectors
+    sparse_embeddings = generate_sparse_vectors(doc_texts)
     
-        # Compute similarity between query and all documents
-        query_doc_sims = cosine_similarity(query_embedding, doc_embeddings)[0]
+    # 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')
+            }
+        })
     
-        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)
+    # 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)
     
-        return selected
+    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)
+
     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]
+        # 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
         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]
+            query_embedding=query_embedding,
+            doc_embeddings=context_embeddings,
+            k=6,
+            lambda_mult=0.6,
+        )
+        
+        top_chunks = [contexts[j] for j in selected_indices]
         return "\n\n".join(top_chunks)
-    # Retrieve and rerank all in parallel
+
+    # Retrieve and rerank all questions 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 = {
@@ -236,23 +259,16 @@ 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."}
-
-
-# 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": "API is running. Go to /docs for documentation."}