Update app.py

app.py

@@ -13,11 +13,11 @@ from typing import List, Tuple
 from langchain.schema import BaseRetriever
 from langchain_core.documents import Document
 from langchain_core.runnables import chain
-import gradio as gr
 from pinecone import Pinecone, ServerlessSpec 
 import openai
-from langchain.retrievers import BM25Retriever
 import numpy as np
+from pinecone.grpc import PineconeGRPC as Pinecone
+
     
 load_dotenv()
 
@@ -26,7 +26,7 @@ openai.api_key = os.environ.get("OPENAI_API_KEY")
 pinecone_api_key = os.environ.get("PINECONE_API_KEY")
 pinecone_environment = os.environ.get("PINECONE_ENV")
 voyage_api_key = os.environ.get("VOYAGE_API_KEY")
-pinecone_index_name = "rag-proto012"
+pinecone_index_name = "rag-proto011"
 
 # Initialize Pinecone
 pc = Pinecone(api_key=pinecone_api_key)
@@ -54,7 +54,7 @@ def search_documents(query):
     try:
         vector_store = PineconeVectorStore(index_name=pinecone_index_name, embedding=embeddings)
         
-        results = vector_store.max_marginal_relevance_search(query, k=7, fetch_k=20)  # Adjust fetch_k for more diverse results
+        results = vector_store.max_marginal_relevance_search(query, k=10, fetch_k=30)  # Adjust fetch_k for more diverse results
         
         # Filter results to ensure uniqueness based on metadata.id
         seen_ids = set()
@@ -72,17 +72,30 @@ def search_documents(query):
                 "doc_id": result.metadata.get("doc_id", "N/A"),
                 "chunk_id": result.metadata.get("id", "N/A"),
                 "title": result.metadata.get("source", "N/A"),
-                "relevant_text": result.page_content,
+                "text": result.page_content,
                 "page_number": result.metadata.get("page", "N/A"),
                 "score": result.metadata.get("score", 0.0),  # Score might not be available in all libraries
             })
         
         # Combine the relevant text for additional processing
-        combined_context = "\n\n".join([res["relevant_text"] for res in context])
+        combined_context = "\n\n".join([res["text"] for res in context])
         return context, combined_context
     except Exception as e:
         return [], f"Error searching documents: {str(e)}"
 
+# Reranker
+def rerank(query, context):        
+    result = pc.inference.rerank(
+        model="bge-reranker-v2-m3",
+        query=query,
+        documents=context,
+        top_n=5,
+        return_documents=True,
+        parameters={
+            "truncate": "END"
+        }
+    )
+    return result
 
 def generate_output(context, query):
     try:
@@ -104,6 +117,31 @@ def generate_output(context, query):
 def complete_workflow(query):
     try:
         context_data, combined_context = search_documents(query)
+
+        # print("Context Data")
+        
+        # [print(doc) for doc in context_data]
+
+        reranked = rerank(query, context_data)
+
+        context_data= []
+        
+        # print("\n\n reranked data")
+        # print(reranked.data)
+
+        for i, entry in enumerate(reranked.data):  # Access the 'data' attribute
+            context_data.append({
+                'chunk_id': entry['document']['chunk_id'],
+                'doc_id': entry['document']['doc_id'],
+                'title': entry['document']['title'],
+                'text': entry['document']['text'],
+                'page_number': entry['document']['page_number'],
+                'score': entry['score']
+            })
+
+        # print("\n\n New Context Data")
+        # [print(doc) for doc in context_data]
+
         
         document_titles = list({os.path.basename(doc["title"]) for doc in context_data})  # Get only file names
         formatted_titles = " " + "\n".join(document_titles)
@@ -113,11 +151,11 @@ def complete_workflow(query):
         results = {
             "results": [
                 {
-                    "natural_language_output": generate_output(doc["relevant_text"], query),
+                    "natural_language_output": generate_output(doc["text"], query),
                     "chunk_id": doc["chunk_id"],
                     "document_id": doc["doc_id"],  # Assuming doc_id is the UUID
                     "title": doc["title"],
-                    "relevant_text": doc["relevant_text"],
+                    "text": doc["text"],
                     "page_number": doc["page_number"],
                     "score": doc["score"],
                 }
@@ -130,6 +168,7 @@ def complete_workflow(query):
     except Exception as e:
         return {"results": [], "total_results": 0}, f"Error in workflow: {str(e)}"
 
+
 def gradio_app():
     with gr.Blocks(css=".result-output {width: 150%; font-size: 16px; padding: 10px;}") as app:
         gr.Markdown("### Intelligent Document Search Prototype-v0.1.2 ")