Update app.py

app.py

@@ -16,6 +16,8 @@ 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
     
 load_dotenv()
 
@@ -48,6 +50,79 @@ embeddings = VoyageAIEmbeddings(
     voyage_api_key=voyage_api_key, model="voyage-law-2"
 )
 
+def hybrid_search_documents(query):
+    try:
+        vector_store = PineconeVectorStore(index_name=pinecone_index_name, embedding=embeddings)
+
+        vector_results = vector_store.similarity_search_with_score(query, k=15)  # Fetch top 15 results
+
+        bm25_retriever = BM25Retriever.from_documents(uploaded_docs)
+        
+        keyword_results = bm25_retriever.get_relevant_documents(query)[:10]  # Fetch top 10 keyword-based results
+
+        # Combine results while avoiding duplicates
+        seen_ids = set()
+        hybrid_results = []
+
+        def process_result(result, score, method):
+            unique_id = result.metadata.get("id")
+            if unique_id not in seen_ids:
+                seen_ids.add(unique_id)
+                hybrid_results.append({
+                    "doc_id": result.metadata.get("doc_id", "N/A"),
+                    "chunk_id": unique_id,
+                    "title": result.metadata.get("source", "N/A"),
+                    "relevant_text": result.page_content,
+                    "page_number": result.metadata.get("page", "N/A"),
+                    "score": score,
+                    "method": method  # Vector or BM25
+                })
+
+        # Add dense results
+        for res, score in vector_results:
+            process_result(res, score, "vector")
+
+        # Add BM25 results with an arbitrary score
+        for res in keyword_results:
+            process_result(res, score=0.85, method="bm25")  # BM25 scores aren't normalized, so we use an approximation
+
+        # 🔹 Step 3: Re-Ranking with LLM (GPT-4)
+        llm = ChatOpenAI(model="gpt-4", openai_api_key=openai.api_key, temperature=0.3)
+
+        ranking_prompt = """
+        You are a document retrieval assistant. Given the following query and retrieved documents, 
+        rank them based on their relevance to the query.
+
+        Query: {query}
+
+        Documents:
+        {documents}
+
+        Return a ranked list of document IDs in order of relevance.
+        """
+
+        doc_texts = "\n".join([f"ID: {doc['chunk_id']}, Text: {doc['relevant_text']}" for doc in hybrid_results])
+        prompt = ranking_prompt.format(query=query, documents=doc_texts)
+        response = llm([HumanMessage(content=prompt)]).content.strip()
+
+        # Extract ordered ranking from LLM response
+        ordered_ids = response.split("\n")  # Assuming LLM returns sorted IDs line-by-line
+        hybrid_results = sorted(hybrid_results, key=lambda x: ordered_ids.index(x["chunk_id"]) if x["chunk_id"] in ordered_ids else 999)
+
+        # Normalize Scores for Consistency
+        scores = [doc["score"] for doc in hybrid_results]
+        min_score, max_score = min(scores), max(scores)
+        for doc in hybrid_results:
+            doc["score"] = (doc["score"] - min_score) / (max_score - min_score + 1e-6)  # Normalize scores between 0 and 1
+
+        # Combine context for query generation
+        combined_context = "\n\n".join([res["relevant_text"] for res in hybrid_results])
+
+        return hybrid_results, combined_context
+
+    except Exception as e:
+        return [], f"Error in hybrid search: {str(e)}"
+        
 def search_documents(query):
     try:
         vector_store = PineconeVectorStore(index_name=pinecone_index_name, embedding=embeddings)
@@ -127,6 +202,53 @@ def complete_workflow(query):
         return {"results": []}, f"Error in workflow: {str(e)}"
 
 
+import os
+from langchain_openai import ChatOpenAI
+from langchain.schema import HumanMessage
+
+def complete_workflow(query):
+    try:
+        # 🔹 Step 1: Perform Hybrid Search (Vector + BM25)
+        context_data, combined_context = hybrid_search_documents(query)
+
+        # 🔹 Step 2: Generate LLM-based Natural Language Output
+        llm = ChatOpenAI(model="gpt-4", openai_api_key=openai.api_key, temperature=0.7)
+        prompt_template = """
+        Use the following context to answer the question as accurately as possible:
+        
+        Context: {context}
+        Question: {question}
+        
+        Answer:
+        """
+        prompt = prompt_template.format(context=combined_context, question=query)
+        response = llm([HumanMessage(content=prompt)])
+
+        # 🔹 Step 3: Format Results
+        document_titles = list({os.path.basename(doc["title"]) for doc in context_data})  # Extract unique file names
+        formatted_titles = "\n".join(document_titles)
+
+        results = {
+            "results": [
+                {
+                    "natural_language_output": response.content,
+                    "chunk_id": doc["chunk_id"],
+                    "document_id": doc["doc_id"],
+                    "title": doc["title"],
+                    "relevant_text": doc["relevant_text"],
+                    "page_number": doc["page_number"],
+                    "score": doc["score"],
+                    "method": doc["method"],  # "vector" or "bm25"
+                }
+                for doc in context_data
+            ],
+            "total_results": len(context_data),  # Return total number of retrieved results
+        }
+
+        return results, formatted_titles  # Return both results and formatted document titles
+    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 ")