Update app.py

app.py

@@ -1,187 +1,222 @@
 import os
+import json
+import redis
+import openai
+import numpy as np
+import gradio as gr
 from dotenv import load_dotenv
-from langchain_community.document_loaders import PyPDFLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.schema import HumanMessage
-from langchain_openai import OpenAIEmbeddings, ChatOpenAI
+from pinecone import Pinecone, ServerlessSpec
+from langchain_openai import OpenAIEmbeddings
 from langchain_voyageai import VoyageAIEmbeddings
 from langchain_pinecone import PineconeVectorStore
-from langchain.prompts import PromptTemplate
-from pinecone import Pinecone
-from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.metrics.pairwise import cosine_similarity
-import openai
-import gradio as gr
+from langchain_openai import ChatOpenAI
+from langchain_core.documents import Document
+from langchain_core.output_parsers import StrOutputParser
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.schema import HumanMessage
+from sentence_transformers import CrossEncoder
 
-# Load API keys
+# Load environment variables
 load_dotenv()
-openai.api_key = os.environ.get("OPENAI_API_KEY")
-pinecone_api_key = os.environ.get("PINECONE_API_KEY")
-voyage_api_key = os.environ.get("VOYAGE_API_KEY")
+openai.api_key = os.getenv("OPENAI_API_KEY")
+pinecone_api_key = os.getenv("PINECONE_API_KEY")
+pinecone_environment = os.getenv("PINECONE_ENV")
+voyage_api_key = os.getenv("VOYAGE_API_KEY")
 
 # Initialize Pinecone
 pc = Pinecone(api_key=pinecone_api_key)
+
+# Redis caching for reranking
+redis_client = redis.Redis(host='localhost', port=6379, db=0)
+
+# Initialize embeddings
 embeddings = VoyageAIEmbeddings(voyage_api_key=voyage_api_key, model="voyage-law-2")
 
-# 🔹 Query Expansion using GPT-4
+# Load Cross-Encoder model for reranking
+reranker_model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2")
+
+
+# **1. Optimized Pinecone Index Initialization**
+def initialize_pinecone_index(index_name):
+    """
+    Ensures the Pinecone index is optimized for fast ANN-based search.
+    """
+    if index_name not in pc.list_indexes():
+        pc.create_index(
+            name=index_name,
+            dimension=1024,
+            metric="cosine",
+            spec=ServerlessSpec(cloud="aws", region="us-west-2"),
+            hnsw_config={"ef_construction": 200, "M": 16}  # Fast ANN search
+        )
+    
+    return PineconeVectorStore(index_name=index_name, embedding=embeddings)
+
+
+# **2. Query Expansion**
+QUERY_EXPANSIONS = {
+    "docs": "Find all legal documents related to case law.",
+    "contract": "Find contracts and legal agreements relevant to the query.",
+    "policy": "Retrieve company policies and regulatory guidelines."
+}
+
 def expand_query(query):
-    llm = ChatOpenAI(model="gpt-4", openai_api_key=openai.api_key, temperature=0.3)
-    prompt = f"Rewrite this vague query into a more specific one:\nQuery: {query}\nSpecific Query:"
-    refined_query = llm([HumanMessage(content=prompt)]).content.strip()
-    return refined_query if refined_query else query
-
-# 🔹 Hybrid Search (TF-IDF + Semantic Retrieval)
-def hybrid_search(query, user_groups, index_name="briefmeta", min_score=0, fetch_k=50):
-    vector_store = PineconeVectorStore(index_name=index_name, embedding=embeddings)
-    semantic_results = vector_store.max_marginal_relevance_search(query, k=10, fetch_k=fetch_k)
-
-    all_texts = [doc.page_content for doc in semantic_results]
-    vectorizer = TfidfVectorizer(stop_words="english")
-    tfidf_matrix = vectorizer.fit_transform(all_texts)
-    query_tfidf = vectorizer.transform([query])
-    keyword_scores = cosine_similarity(query_tfidf, tfidf_matrix).flatten()
-
-    combined_results, seen_ids = [], set()
-    for i, doc in enumerate(semantic_results):
-        doc_id, doc_groups = doc.metadata.get("id"), doc.metadata.get("groups", [])
-        semantic_score = float(doc.metadata.get("score", 0))
-        keyword_score = float(keyword_scores[i])
-        final_score = 0.7 * semantic_score + 0.3 * keyword_score  # Hybrid score
-
-        if doc_id not in seen_ids and any(group in user_groups for group in doc_groups) and final_score > min_score:
-            seen_ids.add(doc_id)
-            doc.metadata["final_score"] = final_score
-            combined_results.append(doc)
-
-    combined_results.sort(key=lambda x: x.metadata["final_score"], reverse=True)
-    return [
-        {
-            "doc_id": doc.metadata.get("doc_id", "N/A"),
-            "chunk_id": doc.metadata.get("id", "N/A"),
-            "title": doc.metadata.get("source", "N/A"),
-            "text": doc.page_content,
-            "page_number": str(doc.metadata.get("page_number", "N/A")),
-            "score": str(doc.metadata.get("final_score", "N/A")),
-        }
-        for doc in combined_results
-    ]
-
-# 🔹 Metadata-Weighted Reranking
-def rerank(query, context):
-    reranker = pc.inference.rerank(
-        model="bge-reranker-v2-m3", query=query, documents=context, top_n=10, return_documents=True
-    )
-
-    final_reranked = []
-    for entry in reranker.data:
-        doc, score = entry["document"], float(entry["score"])
-        citation_boost = 1.2 if "high_citations" in doc.get("tags", []) else 1.0
-        recency_boost = 1.1 if "recent_upload" in doc.get("tags", []) else 1.0
-        final_score = score * citation_boost * recency_boost
-        doc["final_score"] = final_score
-        final_reranked.append(doc)
-
-    final_reranked.sort(key=lambda x: x["final_score"], reverse=True)
-    return final_reranked
-
-# 🔹 Intelligent Search Summary Generator
+    """
+    Expands the query efficiently using predefined mappings and LLM if needed.
+    """
+    query = query.strip().lower()
+    
+    if query in QUERY_EXPANSIONS:
+        return QUERY_EXPANSIONS[query]
+    
+    if len(query.split()) < 3:
+        llm = ChatOpenAI(model="gpt-4", openai_api_key=openai.api_key, temperature=0.3)
+        prompt = f"Rewrite the following vague search query into a more specific one:\nQuery: {query}.\nSpecific Query:"
+        refined_query = llm([HumanMessage(content=prompt)]).content.strip()
+        return refined_query if refined_query else query
+    
+    return query
+
+
+# **3. Hybrid Search (Dense + Sparse Fusion)**
+def search_documents(query, user_groups, index_name="briefmeta"):
+    """
+    Hybrid search combining semantic and sparse (keyword) retrieval.
+    """
+    try:
+        vector_store = PineconeVectorStore(index_name=index_name, embedding=embeddings)
+
+        # Dense search (Semantic embeddings)
+        dense_results = vector_store.similarity_search_with_relevance_scores(
+            query=query, k=10, sparse_weight=0.3,  
+            filter={"groups": {"$in": user_groups}}
+        )
+
+        # Sparse search (BM25-style keyword search)
+        sparse_results = vector_store.sparse_search(query=query, k=10)
+
+        # Fusion of results
+        hybrid_results = {}
+        for doc, score in dense_results:
+            hybrid_results[doc.metadata["id"]] = {"doc": doc, "score": score * 0.7}
+
+        for doc, score in sparse_results:
+            if doc.metadata["id"] in hybrid_results:
+                hybrid_results[doc.metadata["id"]]["score"] += score * 0.3
+            else:
+                hybrid_results[doc.metadata["id"]] = {"doc": doc, "score": score * 0.3}
+
+        # Sort by final score
+        final_results = sorted(hybrid_results.values(), key=lambda x: x["score"], reverse=True)
+
+        # Format output
+        search_output = [
+            {
+                "doc_id": item["doc"].metadata.get("doc_id", "N/A"),
+                "title": item["doc"].metadata.get("source", "N/A"),
+                "text": item["doc"].page_content,
+                "score": round(item["score"], 3)
+            }
+            for item in final_results
+        ]
+
+        return search_output
+
+    except Exception as e:
+        return [], f"Error in hybrid search: {str(e)}"
+
+
+# **4. Reranking with Cross-Encoder (Cached)**
+def rerank_results(query, search_results):
+    """
+    Uses a Cross-Encoder for reranking search results.
+    """
+    if not search_results:
+        return search_results  
+
+    cache_key = f"rerank:{query}"
+    cached_result = redis_client.get(cache_key)
+    if cached_result:
+        return json.loads(cached_result)
+
+    # Prepare input pairs for reranking
+    pairs = [(query, doc["text"]) for doc in search_results]
+    scores = reranker_model.predict(pairs)
+
+    # Attach scores and sort
+    for i, score in enumerate(scores):
+        search_results[i]["rerank_score"] = round(float(score), 3)
+
+    sorted_results = sorted(search_results, key=lambda x: x["rerank_score"], reverse=True)
+
+    redis_client.setex(cache_key, 600, json.dumps(sorted_results))  # Cache for 10 min
+    return sorted_results
+
+
+# **5. Intelligent Search Summary**
 def generate_search_summary(search_results, query):
+    """
+    Generates an intelligent search summary.
+    """
     if not search_results:
-        return "No relevant documents found. Try refining your query."
+        return "No relevant documents were found for your search."
 
-    num_results = len(search_results)
-    doc_titles = [doc.get("title", "Unknown Document") for doc in search_results]
-    doc_pages = [doc.get("page_number", "N/A") for doc in search_results]
-    relevance_scores = [float(doc.get("score", 0)) for doc in search_results]
+    top_docs = search_results[:3]
+    doc_titles = [doc["title"] for doc in top_docs]
 
     summary_prompt = f"""
-    Generate a concise 1-3 sentence summary:
+    Generate a **concise** 2-3 sentence summary of the search results.
     - User Query: "{query}"
-    - Matching Documents: {num_results} found
-    - Titles: {", ".join(set(doc_titles))}
-    - Pages Referenced: {", ".join(set(doc_pages))}
-    - Relevance Scores (0-1): {relevance_scores}
-    Provide a clear, user-friendly summary with an action suggestion.
+    - Matching Documents: {len(search_results)} found
+    - Titles: {", ".join(doc_titles)}
+
+    **Summarize in user-friendly language.**
     """
 
-    llm = ChatOpenAI(model="gpt-4", openai_api_key=openai.api_key, temperature=0.5)
+    llm = ChatOpenAI(model="gpt-3.5-turbo", openai_api_key=openai.api_key, temperature=0.5)
     summary = llm([HumanMessage(content=summary_prompt)]).content.strip()
+
     return summary if summary else "No intelligent summary available."
 
-# 🔹 LLM-based Answer Generation
-def generate_output(context, query):
-    if not context.strip():
-        return "No relevant information found. Try refining your query."
-
-    llm = ChatOpenAI(model="gpt-4", openai_api_key=openai.api_key, temperature=0.5)
-    prompt_template = PromptTemplate(
-        template="Use the following context to answer the question:\nContext: {context}\nQuestion: {question}\nAnswer:",
-        input_variables=["context", "question"],
-    )
-    prompt = prompt_template.format(context=context, question=query)
-    response = llm([HumanMessage(content=prompt)]).content.strip()
-    return response if response else "No relevant answer found."
-
-# 🔹 Full Workflow
+
+# **6. Full RAG Workflow**
 def complete_workflow(query, user_groups, index_name="briefmeta"):
+    """
+    Full RAG workflow: Hybrid Search -> Reranking -> Intelligent Summary
+    """
     try:
-        refined_query = expand_query(query)
-        context_data = hybrid_search(refined_query, user_groups)
-        reranked_results = rerank(refined_query, context_data)
+        query = expand_query(query)
+        raw_results = search_documents(query, user_groups, index_name)
+        reranked_results = rerank_results(query, raw_results)
 
-        context_data = [
-            {
-                'chunk_id': doc["chunk_id"],
-                'doc_id': doc["doc_id"],
-                'title': doc["title"],
-                'text': doc["text"],
-                'page_number': str(doc["page_number"]),
-                'score': str(doc["final_score"])
-            }
-            for doc in reranked_results
-        ]
-
-        document_titles = list({os.path.basename(doc["title"]) for doc in context_data})
+        document_titles = list({doc["title"] for doc in reranked_results})
         formatted_titles = " " + "\n".join(document_titles)
-        intelligent_search_summary = generate_search_summary(context_data, refined_query)
+
+        intelligent_search_summary = generate_search_summary(reranked_results, query)
 
         results = {
-            "results": [
-                {
-                    "natural_language_output": generate_output(doc["text"], refined_query),
-                    "chunk_id": doc["chunk_id"],
-                    "document_id": doc["doc_id"],
-                    "title": doc["title"],
-                    "text": doc["text"],
-                    "page_number": doc["page_number"],
-                    "score": doc["score"],
-                }
-                for doc in context_data
-            ],
-            "total_results": len(context_data),
-            "intelligent_search_summary": intelligent_search_summary
+            "results": reranked_results[:5],
+            "total_results": len(reranked_results)
         }
 
         return results, formatted_titles, intelligent_search_summary
-
     except Exception as e:
-        return {"results": [], "total_results": 0, "intelligent_search_summary": "Error generating summary."}, f"Error in workflow: {str(e)}"
+        return {"results": [], "total_results": 0}, f"Error in workflow: {str(e)}"
+
 
-# 🔹 Gradio UI
+# **7. Gradio UI**
 def gradio_app():
     with gr.Blocks() as app:
-        gr.Markdown("### 📄 Intelligent Document Search Prototype-v0.2")
-        user_query = gr.Textbox(label="🔍 Enter Search Query")
-        user_groups = gr.Textbox(label="👥 User Groups", placeholder="e.g., ['KarthikPersonal']")
-        index_name = gr.Textbox(label="📂 Index Name", placeholder="Default: briefmeta")
-        search_btn = gr.Button("🔎 Search")
-        search_summary = gr.Textbox(label="📜 Intelligent Search Summary", interactive=False)
-        result_output = gr.JSON(label="📊 Search Results")
-        titles_output = gr.Textbox(label="📂 Retrieved Document Titles", interactive=False)
+        gr.Markdown("## 🔍 AI-Powered Document Search")
+
+        user_query = gr.Textbox(label="Enter Your Search Query")
+        user_groups = gr.Textbox(label="Enter User Groups", interactive=True)
+        search_btn = gr.Button("Search")
+        results_output = gr.JSON(label="Search Results")
+        search_summary = gr.Textbox(label="Intelligent Search Summary")
 
-        search_btn.click(complete_workflow, inputs=[user_query, user_groups, index_name], outputs=[result_output, titles_output, search_summary])
+        search_btn.click(complete_workflow, inputs=[user_query, user_groups], outputs=[results_output, search_summary])
 
     return app
 
-# Launch the App
 gradio_app().launch()