Update rag_utils.py

rag_utils.py

@@ -5,25 +5,60 @@ from groq import AsyncGroq
 import json
 import re
 from rank_bm25 import BM25Okapi
-from sentence_transformers import SentenceTransformer
+from sentence_transformers import SentenceTransformer, CrossEncoder # Added CrossEncoder
 from sklearn.preprocessing import MinMaxScaler
 import numpy as np
 from typing import Any, List, Tuple
 import asyncio
 import torch
 import time
-from flashrank import Ranker, RerankRequest # Import the FlashRank library
 
 # --- Configuration (can be overridden by the calling app) ---
 CHUNK_SIZE = 1000
 CHUNK_OVERLAP = 200
-TOP_K_CHUNKS = 5  # The final number of chunks to send to the LLM
+TOP_K_CHUNKS = 10  # The final number of chunks to send to the LLM
 # A larger number of initial candidates for reranking
-INITIAL_K_CANDIDATES = 20 
+INITIAL_K_CANDIDATES = 20
 GROQ_MODEL_NAME = "llama3-8b-8192"
-EMBEDDING_MODEL_NAME = "all-MiniLM-L6-v2" # A good general-purpose embedding model
+HYDE_MODEL = "meta-llama/llama-4-scout-17b-16e-instruct"
+EMBEDDING_MODEL_NAME = "all-MiniLM-L6-v2"
+
+# --- Hypothetical Document Generation and EmbeddingClient remain unchanged ---
+async def generate_hypothetical_document(query: str, groq_api_key: str) -> str:
+    """
+    Generates a hypothetical document using the Groq API.
+    This prompt is generic and does not require prior knowledge of the document style.
+    """
+    if not groq_api_key:
+        print("Groq API key not set. Skipping hypothetical document generation.")
+        return ""
+
+    print(f"Starting HyDE generation for query: '{query}'...")
+    client = AsyncGroq(api_key=groq_api_key)
+    prompt = (
+        f"You are a document writer. Your task is to write a brief passage as a section of a document "
+        f"that could answer the following question. The passage should use specific terminology and "
+        f"a formal tone, as if it were an excerpt from a larger document. Do not include the question, "
+        f"and do not add any conversational text. The goal is to create a concise, semantically rich text "
+        f"to guide a search engine to find similarly styled and detailed content.\n\n"
+        f"Question: {query}\n\n"
+        f"Hypothetical Section:"
+    )
+
+    try:
+        chat_completion = await client.chat.completions.create(
+            messages=[{"role": "user", "content": prompt}],
+            model=HYDE_MODEL,
+            temperature=0.7,
+            max_tokens=500,
+        )
+        hyde_doc = chat_completion.choices[0].message.content
+        print("Hypothetical document generated.")
+        return hyde_doc
+    except Exception as e:
+        print(f"An error occurred during HyDE generation: {e}")
+        return ""
 
-# --- Class for managing the Sentence Transformer model ---
 class EmbeddingClient:
     """A client for generating text embeddings using a local, open-source model."""
     def __init__(self, model_name: str = EMBEDDING_MODEL_NAME):
@@ -51,8 +86,9 @@ class HybridSearchManager:
         self.embedding_client = EmbeddingClient(model_name=embedding_model_name)
         self.document_chunks = []
         self.document_embeddings = None
-        self.reranker = Ranker() # Initialize the FlashRank reranker
-        print("FlashRank reranker initialized.")
+        # Initialize BGE reranker model
+        self.reranker = CrossEncoder('BAAI/bge-reranker-base', device='cuda' if torch.cuda.is_available() else 'cpu')
+        print("BGE Reranker initialized.")
 
     async def initialize_models(self, documents: list[Document]):
         self.document_chunks = documents
@@ -67,90 +103,85 @@ class HybridSearchManager:
         print(f"Computing and storing document embeddings on {self.embedding_client.device}...")
         self.document_embeddings = self.embedding_client.get_embeddings(corpus)
         print("Document embeddings computed.")
-    
-    async def perform_hybrid_search(self, query: str, top_k: int) -> Tuple[List[dict], float]:
+
+    async def retrieve_candidates(self, query: str, hyde_doc: str) -> List[dict]:
         """
-        Performs a hybrid search, then reranks the results, and returns the top chunks
-        along with the time taken for reranking.
+        Performs a HyDE-enhanced hybrid search to retrieve initial candidates
+        without reranking.
         """
         if self.bm25_model is None or self.document_embeddings is None:
             raise ValueError("Hybrid search models are not initialized. Call initialize_models first.")
-        print(f"Performing hybrid search for query: '{query}' (top_k={top_k})...")
+        print(f"Performing hybrid search for candidate retrieval for query: '{query}'...")
 
-        # Get a larger number of chunks for a better reranking pool
+        hyde_query = f"{query}\n\n{hyde_doc}" if hyde_doc else query
         tokenized_query = query.split(" ")
         bm25_scores = self.bm25_model.get_scores(tokenized_query)
-
-        query_embedding = self.embedding_client.get_embeddings([query])
+        query_embedding = self.embedding_client.get_embeddings([hyde_query])
         from torch.nn.functional import cosine_similarity
         dense_scores = cosine_similarity(query_embedding, self.document_embeddings)
         dense_scores = dense_scores.cpu().numpy()
 
         if len(bm25_scores) == 0 or len(dense_scores) == 0:
-            return [], 0.0
+            return []
 
         scaler = MinMaxScaler()
         normalized_bm25_scores = scaler.fit_transform(bm25_scores.reshape(-1, 1)).flatten()
         normalized_dense_scores = scaler.fit_transform(dense_scores.reshape(-1, 1)).flatten()
         combined_scores = 0.5 * normalized_bm25_scores + 0.5 * normalized_dense_scores
-        
-        # We now get `INITIAL_K_CANDIDATES` documents from the combined search
+
         ranked_indices = np.argsort(combined_scores)[::-1]
         top_initial_indices = ranked_indices[:INITIAL_K_CANDIDATES]
-        
+
         retrieved_results = []
         for idx in top_initial_indices:
             doc = self.document_chunks[idx]
             retrieved_results.append({
                 "content": doc.page_content,
-                "document_metadata": doc.metadata
+                "document_metadata": doc.metadata,
+                "initial_score": combined_scores[idx] # Optionally store the initial score
             })
-        
-        print(f"Retrieved {len(retrieved_results)} initial chunks. Starting reranking...")
-        
-        # --- Reranking Step with Timing ---
-        start_time_rerank = time.perf_counter()
+
+        print(f"Retrieved {len(retrieved_results)} initial candidates for reranking.")
+        return retrieved_results
+
+    async def rerank_results(self, query: str, retrieved_results: List[dict], top_k: int) -> List[dict]:
+        """
+        Performs reranking on a list of retrieved candidate documents.
+        """
         if not retrieved_results:
-            return [], 0.0
+            return []
 
-        # FlashRank expects a list of dictionaries with a "text" key
-        passages = [{"text": chunk["content"]} for chunk in retrieved_results]
-        
-        # The reranker takes a query and a list of passages and returns a reranked list
-        reranked_results = await asyncio.to_thread(
-            self.reranker.rerank, RerankRequest(query=query, passages=passages)
+        print(f"Reranking {len(retrieved_results)} candidates for query: '{query}'...")
+        start_time_rerank = time.perf_counter()
+
+        rerank_input = [[query, chunk["content"]] for chunk in retrieved_results]
+        rerank_scores = await asyncio.to_thread(
+            self.reranker.predict, rerank_input, show_progress_bar=False
         )
         
         end_time_rerank = time.perf_counter()
         rerank_time = end_time_rerank - start_time_rerank
-        
-        # Re-map the reranked results back to our original document format
+
+        scored_results = list(zip(retrieved_results, rerank_scores))
+        scored_results.sort(key=lambda x: x[1], reverse=True)
+
         final_chunks = []
-        for res in reranked_results:
-            # Find the original chunk based on the text
-            original_chunk_data = next(
-                (c for c in retrieved_results if c["content"] == res["text"]),
-                None
-            )
-            if original_chunk_data:
-                final_chunks.append({
-                    "content": original_chunk_data["content"],
-                    "document_metadata": original_chunk_data["document_metadata"],
-                    "rerank_score": res["score"]
-                })
-        
-        # Return the top_k reranked chunks and the timing information
-        print(f"Reranking completed in {rerank_time:.4f} seconds. Retrieved {len(final_chunks[:top_k])} top chunks.")
-        return final_chunks[:top_k], rerank_time
+        for res, score in scored_results[:top_k]:
+            final_chunks.append({
+                "content": res["content"],
+                "document_metadata": res["document_metadata"],
+                "rerank_score": score
+            })
 
+        print(f"Reranking completed in {rerank_time:.4f} seconds. Returning top {len(final_chunks)} chunks.")
+        return final_chunks, rerank_time
 
-# --- Helper Functions (remain unchanged) ---
-def process_markdown_with_manual_sections(
+# --- Other helper functions (process_markdown_with_recursive_chunking, generate_answer_with_groq) remain unchanged ---
+def process_markdown_with_recursive_chunking(
     md_file_path: str,
-    headings_json: dict,
     chunk_size: int,
-    chunk_overlap: int):
-    all_chunks_with_metadata = []
+    chunk_overlap: int) -> List[Document]:
+    all_chunks = []
     full_text = ""
     if not os.path.exists(md_file_path):
         print(f"Error: File not found at '{md_file_path}'")
@@ -169,57 +200,22 @@ def process_markdown_with_manual_sections(
     if not full_text:
         print("Input markdown file is empty.")
         return []
+
     text_splitter = RecursiveCharacterTextSplitter(
         chunk_size=chunk_size,
         chunk_overlap=chunk_overlap,
         length_function=len,
         is_separator_regex=False,
     )
-    heading_texts = headings_json.get("headings", [])
-    print(f"Identified headings for segmentation: {heading_texts}")
-    heading_positions = []
-    for heading in heading_texts:
-        pattern = re.compile(r'\s*'.join(re.escape(word) for word in heading.split()), re.IGNORECASE)
-        match = pattern.search(full_text)
-        if match:
-            heading_positions.append({"heading_text": heading, "start_index": match.start()})
-        else:
-            print(f"Warning: Heading '{heading}' not found in the markdown text using regex. This section might be missed.")
-    heading_positions.sort(key=lambda x: x["start_index"])
-    segments_with_headings = []
-    if heading_positions and heading_positions[0]["start_index"] > 0:
-        preface_text = full_text[:heading_positions[0]["start_index"]].strip()
-        if preface_text:
-            segments_with_headings.append({
-                "section_heading": "Document Start/Preface",
-                "section_text": preface_text
-            })
-    for i, current_heading_info in enumerate(heading_positions):
-        start_index = current_heading_info["start_index"]
-        heading_text = current_heading_info["heading_text"]
-        end_index = len(full_text)
-        if i + 1 < len(heading_positions):
-            end_index = heading_positions[i+1]["start_index"]
-        section_content = full_text[start_index:end_index].strip()
-        if section_content:
-            segments_with_headings.append({
-                "section_heading": heading_text,
-                "section_text": section_content
-            })
-    print(f"Created {len(segments_with_headings)} segments based on provided headings.")
-    for segment in segments_with_headings:
-        section_heading = segment["section_heading"]
-        section_text = segment["section_text"]
-        if section_text:
-            chunks = text_splitter.split_text(section_text)
-            for chunk in chunks:
-                metadata = {
-                    "document_part": "Section",
-                    "section_heading": section_heading,
-                }
-                all_chunks_with_metadata.append(Document(page_content=chunk, metadata=metadata))
-        print(f"Created {len(all_chunks_with_metadata)} chunks with metadata from segmented sections.")
-    return all_chunks_with_metadata
+
+    chunks = text_splitter.split_text(full_text)
+
+    for chunk in chunks:
+        all_chunks.append(Document(page_content=chunk, metadata={"document_part": "Whole Document"}))
+
+    print(f"Created {len(all_chunks)} chunks from the entire document.")
+    return all_chunks
+
 
 async def generate_answer_with_groq(query: str, retrieved_results: list[dict], groq_api_key: str) -> str:
     """
@@ -244,12 +240,13 @@ async def generate_answer_with_groq(query: str, retrieved_results: list[dict], g
         )
     context = "\n\n".join(context_parts)
     prompt = (
-        f"You are a specialized document analyzer assistant. Your task is to answer the user's question "
-        f"solely based on the provided context. Pay close attention to the section heading and document part "
-        f"for each context chunk. Ensure your answer incorporates all relevant details, including any legal nuances "
-        f"and conditions found in the context, and is concise, limited to one or two sentences. "
-        f"Do not explicitly mention the retrieved chunks. If the answer cannot be found in the provided context, "
-        f"clearly state that you do not have enough information.\n\n"
+        f"You are an expert on the provided document. Your task is to answer the user's question "
+        f"based only on the information given. Your answers should be brief, concise, and in a similar style to these examples:\n"
+        f"- Yes, outpatient consultations and diagnostic tests are covered, provided they are medically necessary and fall within the specified sub-limits under the plan.\n"
+        f"- The policy does not cover any expenses incurred during the first 30 days from the inception of the policy, except in the case of accidents.\n"
+        f"- Room rent is covered up to a single private AC room per day unless otherwise specified in the policy schedule.\n"
+        f"- Yes, the policy allows for mid-term inclusion of newly married spouses and newborn children, subject to notification and payment of additional premium within the stipulated time frame.\n"
+        f"Do not mention or refer to the document or the context in your final answer. If the information required to answer the question is not available in the provided context, state that you do not have enough information.\n\n"
         f"Context:\n{context}\n\n"
         f"Question: {query}\n\n"
         f"Answer:"
@@ -271,4 +268,4 @@ async def generate_answer_with_groq(query: str, retrieved_results: list[dict], g
         return answer
     except Exception as e:
         print(f"An error occurred during Groq API call: {e}")
-        return "Could not generate an answer due to an API error."
+        return "Could not generate an answer due to an API error."