Merge pull request #251 from Exodus2004/feat/issue-114-hybrid-search

backend/app/rag/bm25.py

@@ -0,0 +1,146 @@
+"""
+BM25 Keyword Search implementation using rank_bm25.
+Stores a BM25 index per document to allow easy updates and deletions.
+"""
+import os
+import glob
+import pickle
+import logging
+from typing import List, Dict, Any, Optional
+
+from rank_bm25 import BM25Okapi
+from app.config import get_settings
+
+logger = logging.getLogger(__name__)
+settings = get_settings()
+
+def get_bm25_dir(user_id: str) -> str:
+    """Get the directory path for a user's BM25 indexes."""
+    clean_id = user_id.replace("-", "_")
+    path = os.path.join(settings.CHROMA_PERSIST_DIR, "bm25", clean_id)
+    os.makedirs(path, exist_ok=True)
+    return path
+
+def get_bm25_path(user_id: str, document_id: str) -> str:
+    """Get the file path for a specific document's BM25 index."""
+    return os.path.join(get_bm25_dir(user_id), f"{document_id}.pkl")
+
+def tokenize(text: str) -> List[str]:
+    """Simple tokenization for BM25."""
+    # Convert to lowercase and split by whitespace
+    return text.lower().split()
+
+def store_bm25_index(chunks: List[Dict[str, Any]], document_id: str, filename: str, user_id: str):
+    """
+    Build and store a BM25 index for the given document chunks.
+    """
+    if not chunks:
+        return
+
+    texts = [chunk["text"] for chunk in chunks]
+    tokenized_texts = [tokenize(text) for text in texts]
+    bm25 = BM25Okapi(tokenized_texts)
+
+    # Format chunks to match vectorstore output
+    formatted_chunks = []
+    for chunk in chunks:
+        formatted_chunks.append({
+            "text": chunk["text"],
+            "filename": filename,
+            "document_id": document_id,
+            "page": chunk.get("page", 1),
+        })
+
+    data = {
+        "bm25": bm25,
+        "chunks": formatted_chunks
+    }
+
+    path = get_bm25_path(user_id, document_id)
+    try:
+        with open(path, "wb") as f:
+            pickle.dump(data, f)
+        logger.info(f"Stored BM25 index for document {document_id}")
+    except Exception as e:
+        logger.error(f"Failed to store BM25 index for {document_id}: {e}")
+
+def _query_single_index(path: str, tokenized_query: List[str], top_k: int) -> List[Dict[str, Any]]:
+    """Query a single BM25 index file."""
+    if not os.path.exists(path):
+        return []
+
+    try:
+        with open(path, "rb") as f:
+            data = pickle.load(f)
+    except Exception as e:
+        logger.error(f"Failed to load BM25 index from {path}: {e}")
+        return []
+
+    bm25 = data["bm25"]
+    chunks = data["chunks"]
+
+    scores = bm25.get_scores(tokenized_query)
+    
+    # Get top_k indices sorted by score
+    top_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:top_k]
+
+    results = []
+    for i in top_indices:
+        if scores[i] > 0:
+            chunk = chunks[i].copy()
+            # Normalize BM25 score to 0-1 range roughly, or just keep raw.
+            # BM25 scores are usually > 0, often 1-10.
+            # We keep the raw score for now, RRF will handle the ranking.
+            chunk["score"] = float(scores[i])
+            results.append(chunk)
+
+    return results
+
+def query_bm25(
+    query: str,
+    user_id: str,
+    document_id: Optional[str] = None,
+    top_k: int = 10,
+) -> List[Dict[str, Any]]:
+    """
+    Query BM25 index(es) for relevant chunks.
+    """
+    tokenized_query = tokenize(query)
+    
+    if document_id:
+        path = get_bm25_path(user_id, document_id)
+        return _query_single_index(path, tokenized_query, top_k)
+    
+    # If no document_id, query all documents for this user
+    user_dir = get_bm25_dir(user_id)
+    all_results = []
+    
+    for path in glob.glob(os.path.join(user_dir, "*.pkl")):
+        results = _query_single_index(path, tokenized_query, top_k)
+        all_results.extend(results)
+        
+    # Sort all results by score and take top_k
+    all_results.sort(key=lambda x: x["score"], reverse=True)
+    return all_results[:top_k]
+
+def delete_bm25_index(document_id: str, user_id: str):
+    """Delete a specific document's BM25 index."""
+    path = get_bm25_path(user_id, document_id)
+    if os.path.exists(path):
+        try:
+            os.remove(path)
+            logger.info(f"Deleted BM25 index for document {document_id}")
+        except Exception as e:
+            logger.warning(f"Error deleting BM25 index: {e}")
+
+def delete_user_bm25_indexes(user_id: str):
+    """Delete all BM25 indexes for a user."""
+    user_dir = get_bm25_dir(user_id)
+    if os.path.exists(user_dir):
+        try:
+            for path in glob.glob(os.path.join(user_dir, "*.pkl")):
+                os.remove(path)
+            os.rmdir(user_dir)
+            logger.info(f"Deleted BM25 directory for user {user_id}")
+        except Exception as e:
+            logger.warning(f"Error deleting BM25 directory for user {user_id}: {e}")

backend/app/rag/retriever.py

@@ -1,10 +1,18 @@
 """
-Two-stage retrieval: ChromaDB similarity search + cross-encoder reranking.
+Two-stage retrieval: Hybrid Ensemble (ChromaDB + BM25) + cross-encoder reranking.
 """
 import json
 import logging
 import re
 from typing import List, Dict, Any, Optional
+
+# In LangChain 1.3.2+, EnsembleRetriever moved to langchain_classic (imported by langchain_community)
+from langchain_classic.retrievers import EnsembleRetriever
+from langchain_core.retrievers import BaseRetriever
+from langchain_core.documents import Document as LangchainDocument
+from langchain_core.callbacks import CallbackManagerForRetrieverRun
+from pydantic import Field
+
 from app.config import get_settings
 from app.rag.embeddings import embed_query
 from app.rag.tracing import trace_function
@@ -35,6 +43,42 @@ def get_reranker():
     return _reranker if _reranker != "disabled" else None
 
 
+class CustomVectorRetriever(BaseRetriever):
+    user_id: str = Field(description="User ID")
+    document_id: Optional[str] = Field(default=None, description="Document ID")
+    top_k: int = Field(default=10, description="Top K results")
+
+    def _get_relevant_documents(
+        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
+    ) -> List[LangchainDocument]:
+        query_vector = embed_query(query)
+        candidates = query_chunks(
+            query_embedding=query_vector,
+            user_id=self.user_id,
+            document_id=self.document_id,
+            top_k=self.top_k,
+        )
+        return [LangchainDocument(page_content=c["text"], metadata=c) for c in candidates]
+
+
+class CustomBM25Retriever(BaseRetriever):
+    user_id: str = Field(description="User ID")
+    document_id: Optional[str] = Field(default=None, description="Document ID")
+    top_k: int = Field(default=10, description="Top K results")
+
+    def _get_relevant_documents(
+        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
+    ) -> List[LangchainDocument]:
+        from app.rag.bm25 import query_bm25
+        candidates = query_bm25(
+            query=query,
+            user_id=self.user_id,
+            document_id=self.document_id,
+            top_k=self.top_k,
+        )
+        return [LangchainDocument(page_content=c["text"], metadata=c) for c in candidates]
+
+
 def transform_query(query: str) -> List[str]:
     """Rewrite a user question into multiple retrieval-friendly search queries."""
     original_query = query.strip()
@@ -183,28 +227,43 @@ def retrieve(
 ) -> List[Dict[str, Any]]:
     """
     Two-stage retrieval pipeline:
-    1. ChromaDB similarity search (top-K broad)
+    1. Hybrid Search (Vector + BM25 via EnsembleRetriever with RRF) with Query Transformation
     2. Cross-encoder reranking (top-K refined)
 
     Returns chunks with confidence scores.
     """
-    # ── Stage 1: Query transformation + embedding search ─────────────
-    candidates = []
-    for search_query in transform_query(query):
-        query_vector = embed_query(search_query)
-        candidates.extend(
-            query_chunks(
-                query_embedding=query_vector,
-                user_id=user_id,
-                document_id=document_id,
-                top_k=settings.TOP_K_RETRIEVAL,
-            )
-        )
+    # ── Stage 1: Hybrid Search with Query Transformation ─────────────
+    vector_retriever = CustomVectorRetriever(
+        user_id=user_id,
+        document_id=document_id,
+        top_k=settings.TOP_K_RETRIEVAL,
+    )
+
+    bm25_retriever = CustomBM25Retriever(
+        user_id=user_id,
+        document_id=document_id,
+        top_k=settings.TOP_K_RETRIEVAL,
+    )
 
-    if not candidates:
+    ensemble_retriever = EnsembleRetriever(
+        retrievers=[vector_retriever, bm25_retriever],
+        weights=[0.6, 0.4]
+    )
+
+    all_candidates = []
+    for search_query in transform_query(query):
+        docs = ensemble_retriever.invoke(search_query)
+        for i, doc in enumerate(docs):
+            chunk = doc.metadata.copy()
+            # Preserve a mock score based on rank for fallback if reranker fails
+            # We use 1.0/(i+1) as a base RRF-like score
+            chunk["score"] = 1.0 / (i + 1)
+            all_candidates.append(chunk)
+
+    if not all_candidates:
         return []
 
-    candidates = _merge_candidates(candidates)
+    candidates = _merge_candidates(all_candidates)
 
     # ── Stage 2: Cross-encoder reranking ─────────────
     reranker = get_reranker()
@@ -223,8 +282,9 @@ def retrieve(
             candidates.sort(key=lambda x: x.get("rerank_score", 0), reverse=True)
 
         except Exception as e:
-            logger.warning(f"Reranking failed, using embedding scores: {e}")
+            logger.warning(f"Reranking failed, using hybrid scores: {e}")
 
+    # Ensure candidates are sorted by best available score
     candidates.sort(key=lambda x: x.get("rerank_score", x.get("score", 0)), reverse=True)
 
     # ── Take top-K after reranking ─────────���─────────

backend/app/rag/vectorstore.py

@@ -49,12 +49,19 @@ def store_chunks(
     user_id: str,
 ) -> int:
     """
-    Embed and store document chunks in ChromaDB.
+    Embed and store document chunks in ChromaDB, and build a local BM25 index.
     Returns the number of chunks stored.
     """
     if not chunks:
         return 0
 
+    # Build and store BM25 index
+    from app.rag.bm25 import store_bm25_index
+    try:
+        store_bm25_index(chunks, document_id, filename, user_id)
+    except Exception as e:
+        logger.error(f"Could not build BM25 index: {e}")
+
     # Generate captions for any extracted images before embedding
     try:
         from app.rag.vision import generate_captions_for_chunks
@@ -178,6 +185,12 @@ def delete_document_chunks(document_id: str, user_id: str):
     client = get_chroma_client()
     collection_name = get_collection_name(user_id)
 
+    try:
+        from app.rag.bm25 import delete_bm25_index
+        delete_bm25_index(document_id, user_id)
+    except Exception as e:
+        logger.warning(f"Error deleting BM25 index: {e}")
+
     try:
         collection = client.get_collection(name=collection_name)
         # Get all IDs for this document
@@ -197,6 +210,12 @@ def delete_user_collection(user_id: str):
     client = get_chroma_client()
     collection_name = get_collection_name(user_id)
 
+    try:
+        from app.rag.bm25 import delete_user_bm25_indexes
+        delete_user_bm25_indexes(user_id)
+    except Exception as e:
+        logger.warning(f"Error deleting user BM25 indexes: {e}")
+
     try:
         client.delete_collection(name=collection_name)
         logger.info(f"Deleted collection {collection_name}")

backend/requirements.txt

@@ -30,10 +30,12 @@ python-docx
 
 # LangChain & RAG
 langchain
+langchain-classic
 langchain-community
 langchain-huggingface
 langchain-text-splitters
 langsmith
+rank-bm25
 
 # Embeddings & ML
 sentence-transformers

backend/tests/test_documents.py

@@ -95,6 +95,7 @@ def test_ingest_document_builds_and_saves_graph(db_session, monkeypatch, tmp_pat
 
 def test_delete_document_removes_knowledge_graph(client, auth_headers, ready_document, monkeypatch):
     deleted = {}
+    doc_id = ready_document.id
 
     monkeypatch.setattr("app.routes.documents.delete_document_chunks", lambda **kwargs: None)
     monkeypatch.setattr(
@@ -105,9 +106,9 @@ def test_delete_document_removes_knowledge_graph(client, auth_headers, ready_doc
     )
 
     response = client.delete(
-        f"/api/v1/documents/{ready_document.id}",
+        f"/api/v1/documents/{doc_id}",
         headers=auth_headers,
     )
 
     assert response.status_code == 200
-    assert deleted["document_id"] == ready_document.id
+    assert deleted["document_id"] == doc_id

backend/tests/test_graphrag_agent.py

@@ -34,7 +34,7 @@ def test_generate_answer_appends_graph_context_without_changing_sources(monkeypa
         }
     ]
 
-    monkeypatch.setattr(agent, "get_llm_client", lambda: client)
+    monkeypatch.setattr(agent, "get_llm_client", lambda hf_token=None: client)
     monkeypatch.setattr(agent, "retrieve", lambda **kwargs: chunks)
     monkeypatch.setattr(
         agent,
@@ -66,7 +66,7 @@ def test_generate_answer_stream_appends_graph_context(monkeypatch):
             captured["messages"] = messages
             return iter([])
 
-    monkeypatch.setattr(agent, "get_llm_client", lambda: StreamingClient())
+    monkeypatch.setattr(agent, "get_llm_client", lambda hf_token=None: StreamingClient())
     monkeypatch.setattr(
         agent,
         "retrieve",

backend/tests/test_retriever.py

@@ -72,6 +72,6 @@ def test_retrieve_fans_out_transformed_queries_and_merges_duplicates(monkeypatch
     chunks = retriever.retrieve("How do taxes and healthcare work?", user_id="user-1")
 
     assert searched_queries == ["embedding:taxes", "embedding:healthcare"]
-    assert [chunk["id"] for chunk in chunks] == ["shared", "healthcare", "taxes"]
-    assert chunks[0]["score"] == 0.9
+    assert [chunk["id"] for chunk in chunks] == ["shared", "taxes", "healthcare"]
+    assert chunks[0]["score"] == 1.0
     assert chunks[0]["confidence"] == 100.0