feat(retrieval): implement vectorless hybrid search with zero-dependency BM25 and weighted RRF

app/core/bm25.py

@@ -0,0 +1,118 @@
+import math
+import re
+import time
+from typing import List, Dict, Tuple, Optional, Any, Set
+from collections import Counter
+from utils.logger import app_logger
+
+
+STOPWORDS: Set[str] = {
+    # English stopwords
+    'the', 'and', 'for', 'that', 'this', 'with', 'from', 'what', 'which', 'where',
+    'when', 'how', 'why', 'are', 'was', 'were', 'been', 'have', 'has', 'had', 'does',
+    'did', 'about', 'some', 'any', 'each', 'into', 'over', 'after', 'before', 'such',
+    'a', 'an', 'in', 'on', 'at', 'to', 'by', 'of', 'is', 'it', 'its', 'as', 'be',
+    # Arabic stopwords
+    'في', 'من', 'على', 'إلى', 'عن', 'مع', 'هذا', 'هذه', 'ذلك', 'تلك', 'التي', 'الذي',
+    'ما', 'ماذا', 'كيف', 'لماذا', 'أين', 'متى', 'كم', 'هل', 'هو', 'هي', 'أن', 'إن', 'كان'
+}
+
+
+class BM25NodeIndex:
+    """
+    Lightweight, in-memory Okapi BM25 index over flattened document nodes.
+    Zero external dependencies. Highly optimized for section-level retrieval.
+    """
+    def __init__(self, nodes: List[Dict[str, Any]], k1: float = 1.5, b: float = 0.75):
+        self.k1 = k1
+        self.b = b
+        self.node_ids: List[str] = []
+        self.tokenized_corpus: List[List[str]] = []
+        self.doc_lens: List[int] = []
+        self.doc_freqs: Counter = Counter()
+        self.N: int = 0
+        self.avgdl: float = 0.0
+
+        self._build_index(nodes)
+
+    @staticmethod
+    def tokenize(text: str) -> List[str]:
+        if not text:
+            return []
+        tokens = re.findall(r'[\w\-]+', text.lower())
+        return [t for t in tokens if len(t) >= 2 and t not in STOPWORDS]
+
+    def _build_index(self, nodes: List[Dict[str, Any]]):
+        self.node_ids = []
+        self.tokenized_corpus = []
+        self.doc_lens = []
+        self.doc_freqs = Counter()
+
+        for node in nodes:
+            nid = node.get("node_id")
+            if not nid:
+                continue
+
+            # Weighted field representation:
+            # Title: 3x boost
+            # Summary: 2x boost
+            # Text: 1x (first 3000 chars)
+            title = node.get("title") or ""
+            summary = node.get("summary") or ""
+            text = (node.get("text") or node.get("content") or "")[:3000]
+
+            title_tokens = self.tokenize(title) * 3
+            summary_tokens = self.tokenize(summary) * 2
+            text_tokens = self.tokenize(text)
+
+            combined_tokens = title_tokens + summary_tokens + text_tokens
+            if not combined_tokens:
+                combined_tokens = self.tokenize(title or nid) or ["document"]
+
+            self.node_ids.append(nid)
+            self.tokenized_corpus.append(combined_tokens)
+            dl = len(combined_tokens)
+            self.doc_lens.append(dl)
+            self.doc_freqs.update(set(combined_tokens))
+
+        self.N = len(self.node_ids)
+        self.avgdl = (sum(self.doc_lens) / self.N) if self.N > 0 else 1.0
+
+    def search(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
+        """
+        Executes Okapi BM25 scoring over indexed nodes.
+        Returns [(node_id, score), ...] sorted descending.
+        """
+        if self.N == 0:
+            return []
+
+        q_tokens = self.tokenize(query)
+        if not q_tokens:
+            # Fallback to raw tokens if all words were stopwords
+            raw_tokens = re.findall(r'[\w\-]+', query.lower())
+            q_tokens = [t for t in raw_tokens if len(t) >= 2]
+            if not q_tokens:
+                return []
+
+        scores: List[Tuple[str, float]] = []
+
+        for idx, doc_tokens in enumerate(self.tokenized_corpus):
+            nid = self.node_ids[idx]
+            dl = self.doc_lens[idx]
+            tf = Counter(doc_tokens)
+            score = 0.0
+
+            for q in q_tokens:
+                if q in tf:
+                    n_q = self.doc_freqs[q]
+                    # Robertson-Spärck Jones IDF
+                    idf = math.log((self.N - n_q + 0.5) / (n_q + 0.5) + 1.0)
+                    f_qd = tf[q]
+                    denom = f_qd + self.k1 * (1.0 - self.b + self.b * (dl / self.avgdl))
+                    score += idf * (f_qd * (self.k1 + 1.0)) / denom
+
+            if score > 0.0:
+                scores.append((nid, score))
+
+        scores.sort(key=lambda x: x[1], reverse=True)
+        return scores[:top_k]

app/core/rag_pipeline.py

@@ -1,13 +1,15 @@
 import asyncio
 import re
 import time
-from typing import Optional, List, Dict, Any
+from typing import Optional, List, Dict, Any, Tuple
 from utils.logger import app_logger
+from core.bm25 import BM25NodeIndex
 
 
 class VectorlessRAG:
     def __init__(self, llm_service):
         self.llm = llm_service
+        self._bm25_cache: Dict[str, BM25NodeIndex] = {}
 
     @staticmethod
     def _parse_paragraphs_with_page_markers(text: str, default_page: Any) -> List[Dict[str, Any]]:
@@ -184,6 +186,74 @@ class VectorlessRAG:
 
         return pruned_nodes
 
+    @staticmethod
+    def _get_doc_nodes(doc: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """Extracts flattened node dictionaries containing node_id, title, summary, text."""
+        node_map = doc.get("node_map") or {}
+        flat_sections = doc.get("flat_sections") or []
+        if flat_sections and node_map:
+            nodes = [node_map[s["node_id"]] for s in flat_sections if s.get("node_id") in node_map]
+            if nodes:
+                return nodes
+        if node_map:
+            return list(node_map.values())
+
+        # Collect from tree recursively
+        tree = doc.get("tree") or []
+        collected = []
+        def collect(items):
+            for it in items:
+                if isinstance(it, dict):
+                    collected.append(it)
+                    if it.get("nodes"):
+                        collect(it["nodes"])
+        collect(tree)
+        return collected
+
+    @staticmethod
+    def reciprocal_rank_fusion(
+        ranked_lists: List[Tuple[List[str], float]],
+        rrf_k: int = 60,
+        top_k: int = 4,
+    ) -> List[Tuple[str, float]]:
+        """
+        Weighted Reciprocal Rank Fusion (RRF):
+        Combines multiple ranked lists into a single consensus ranking.
+        score(d) = sum_L (w_L / (rrf_k + rank_L(d)))
+        """
+        rrf_scores: Dict[str, float] = {}
+        for items, weight in ranked_lists:
+            for rank, item_id in enumerate(items, start=1):
+                if not item_id:
+                    continue
+                rrf_scores[item_id] = rrf_scores.get(item_id, 0.0) + (weight / (rrf_k + rank))
+
+        sorted_items = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
+        return sorted_items[:top_k]
+
+    def bm25_search_nodes(
+        self,
+        query: str,
+        nodes: List[Dict[str, Any]],
+        doc_id: Optional[str] = None,
+        top_k: int = 5,
+    ) -> Tuple[List[Tuple[str, float]], float]:
+        """
+        Executes Okapi BM25 lexical search with in-memory index caching.
+        Returns (matches, latency_ms).
+        """
+        t0 = time.monotonic()
+        cache_key = doc_id or str(id(nodes))
+        index = self._bm25_cache.get(cache_key)
+        if index is None or index.N != len(nodes):
+            index = BM25NodeIndex(nodes)
+            if doc_id:
+                self._bm25_cache[cache_key] = index
+
+        matches = index.search(query, top_k=top_k)
+        elapsed_ms = (time.monotonic() - t0) * 1000
+        return matches, elapsed_ms
+
     async def answer_multi_doc_query_stream(
         self,
         query: str,
@@ -214,12 +284,39 @@ class VectorlessRAG:
                 node_map = doc.get("node_map")
                 flat_sections = doc.get("flat_sections")
                 doc_name = doc.get("filename") or "Document"
+                doc_id = doc.get("doc_id")
+
+                doc_nodes = self._get_doc_nodes(doc)
+
+                # Parallel Retrieval: LLM Hierarchical Tree Search + Lexical BM25
+                async def run_bm25():
+                    return self.bm25_search_nodes(query, doc_nodes, doc_id=doc_id, top_k=5)
 
-                search_res = await self.llm.llm_tree_search(
+                llm_task = self.llm.llm_tree_search(
                     query, tree, chat_history=chat_history, flat_sections=flat_sections
                 )
-                node_ids = search_res.get("node_list", [])
-                nodes = self.llm.find_nodes_by_ids(tree, node_ids, node_map=node_map)
+                bm25_task = run_bm25()
+
+                search_res, (bm25_matches, bm25_latency) = await asyncio.gather(
+                    llm_task, bm25_task, return_exceptions=False
+                )
+
+                llm_node_ids = search_res.get("node_list", [])
+                bm25_node_ids = [nid for nid, sc in bm25_matches if sc > 0]
+
+                # Weighted Reciprocal Rank Fusion: 0.6 LLM Tree Search, 0.4 BM25
+                fused_results = self.reciprocal_rank_fusion(
+                    ranked_lists=[(llm_node_ids, 0.6), (bm25_node_ids, 0.4)],
+                    top_k=4 if deep_analysis else 3,
+                )
+                fused_node_ids = [nid for nid, _ in fused_results]
+
+                app_logger.info(
+                    f"[HYBRID-RETRIEVAL] doc='{doc_name}' llm_nodes={llm_node_ids} "
+                    f"bm25_nodes={bm25_node_ids} fused={fused_node_ids} bm25_latency={bm25_latency:.2f}ms"
+                )
+
+                nodes = self.llm.find_nodes_by_ids(tree, fused_node_ids, node_map=node_map)
 
                 # Fallback to leading sections for overview/comparative queries if no nodes matched
                 if not nodes:
@@ -247,7 +344,7 @@ class VectorlessRAG:
                     n_copy = dict(n)
                     n_copy["doc_name"] = doc_name
                     n_copy["filename"] = doc_name
-                    n_copy["doc_id"] = doc.get("doc_id")
+                    n_copy["doc_id"] = doc_id
                     tagged_nodes.append(n_copy)
                 return tagged_nodes
 

app/core/storage.py

@@ -25,6 +25,7 @@ class MemoryStorage:
         self._documents: Dict[str, Dict[str, Any]] = {}
         self._node_maps: Dict[str, Dict[str, Any]] = {}
         self._flat_sections: Dict[str, List[Dict[str, Any]]] = {}
+        self._bm25_indices: Dict[str, Any] = {}
         # hash → doc_id quick lookup (in-memory only; rebuilt from DB on demand)
         self._hash_to_doc_id: Dict[str, str] = {}
         # active background indexing tasks
@@ -175,6 +176,12 @@ class MemoryStorage:
     def set_flat_sections(self, doc_id: str, sections: List[Dict[str, Any]]) -> None:
         self._flat_sections[doc_id] = sections
 
+    def get_bm25_index(self, doc_id: str) -> Optional[Any]:
+        return self._bm25_indices.get(doc_id)
+
+    def set_bm25_index(self, doc_id: str, index: Any) -> None:
+        self._bm25_indices[doc_id] = index
+
     # ──────────────────────────────────────────────────────────────────────
     # Document lifecycle
     # ──────────────────────────────────────────────────────────────────────

app/routers/documents.py

@@ -59,6 +59,13 @@ async def run_background_indexing(
         # 3. Store in-memory acceleration indexes
         storage.set_node_map(doc_id, node_map)
         storage.set_flat_sections(doc_id, flat_sections)
+        try:
+            from core.bm25 import BM25NodeIndex
+            doc_nodes = [node_map[s["node_id"]] for s in flat_sections if s.get("node_id") in node_map]
+            if doc_nodes:
+                storage.set_bm25_index(doc_id, BM25NodeIndex(doc_nodes))
+        except Exception as bm_ex:
+            app_logger.warning(f"Could not pre-warm BM25 index for {doc_id}: {bm_ex}")
 
         total_duration = round(time.monotonic() - t_start, 2)
 

test_phase2_hybrid_bm25_rrf.py

@@ -0,0 +1,178 @@
+import os
+import sys
+import time
+import asyncio
+
+# Ensure app path is discoverable
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "app"))
+
+from core.bm25 import BM25NodeIndex
+from core.rag_pipeline import VectorlessRAG
+
+
+def test_bm25_node_index_scoring_and_latency():
+    """Verify BM25NodeIndex scores nodes correctly and executes in under 20ms."""
+    nodes = [
+        {
+            "node_id": "sec_01",
+            "title": "Introduction to Vectorless RAG",
+            "summary": "Overview of tree-based hierarchical indexing without vector databases or chunking.",
+            "text": "Traditional RAG depends on embedding models and vector similarity. Vectorless RAG navigates the document tree directly."
+        },
+        {
+            "node_id": "sec_02",
+            "title": "Knowledge Graph Representation Learning",
+            "summary": "Covers TransE, ERNIE, and entity embedding alignments.",
+            "text": "ERNIE fuses entity representations into text tokens by utilizing TransE knowledge-graph embeddings with multi-layer encoders."
+        },
+        {
+            "node_id": "sec_03",
+            "title": "Evaluation Benchmark & Cost",
+            "summary": "Benchmark results on 80 questions with 87.7% correctness.",
+            "text": "Evaluated on an 80-question benchmark with 1.8s avg latency and $0.0009 cost per query."
+        },
+    ]
+
+    t0 = time.monotonic()
+    index = BM25NodeIndex(nodes)
+    build_time_ms = (time.monotonic() - t0) * 1000
+
+    # Search query specifically targeting Section 2
+    query = "How does ERNIE integrate TransE embeddings into language models?"
+    t_search_0 = time.monotonic()
+    results = index.search(query, top_k=3)
+    search_time_ms = (time.monotonic() - t_search_0) * 1000
+
+    print(f"BM25 build time: {build_time_ms:.2f}ms, search time: {search_time_ms:.2f}ms")
+    assert search_time_ms < 20.0, f"Lexical search exceeded 20ms limit: {search_time_ms}ms"
+    assert len(results) > 0, "Expected at least 1 BM25 match"
+    
+    # Top match must be sec_02 (ERNIE / TransE)
+    top_nid, top_score = results[0]
+    assert top_nid == "sec_02", f"Expected sec_02 to rank highest, got {top_nid}"
+    assert top_score > 0.0
+
+    print("✓ Test 1 Passed: BM25NodeIndex accurately scored relevant section in under 20ms!")
+
+
+def test_bm25_arabic_support_and_title_boost():
+    """Verify Arabic tokenization and title field weighting."""
+    nodes = [
+        {
+            "node_id": "ar_01",
+            "title": "بنية الذاكرة المتعددة",
+            "summary": "نظام الذاكرة ثلاثي الطبقات على PostgreSQL.",
+            "text": "عزل محادثات الجلسات وسجلات المستخدمين وتوليد ملخصات تلقائية."
+        },
+        {
+            "node_id": "ar_02",
+            "title": "البحث الحي عبر FastMCP",
+            "summary": "استرجاع نتائج الويب والتحقق من صحة الروابط.",
+            "text": "فحص الروابط الحية والتأكد من توافر المحتوى المطلوب."
+        }
+    ]
+    index = BM25NodeIndex(nodes)
+    results = index.search("ما هو نظام الذاكرة في قاعدة البيانات؟", top_k=2)
+    assert len(results) > 0
+    assert results[0][0] == "ar_01", f"Expected ar_01 to rank highest for memory query, got {results[0][0]}"
+    print("✓ Test 2 Passed: Arabic tokenization and field weighting verified!")
+
+
+def test_weighted_reciprocal_rank_fusion():
+    """Verify Weighted Reciprocal Rank Fusion combines LLM and BM25 results with 0.6 and 0.4 weights."""
+    llm_ranked = ["sec_02", "sec_01"]
+    bm25_ranked = ["sec_02", "sec_03", "sec_01"]
+
+    fused = VectorlessRAG.reciprocal_rank_fusion(
+        ranked_lists=[(llm_ranked, 0.6), (bm25_ranked, 0.4)],
+        rrf_k=60,
+        top_k=3,
+    )
+    
+    fused_nids = [nid for nid, sc in fused]
+    # sec_02 is Rank 1 in both lists -> must be top fused
+    assert fused_nids[0] == "sec_02", f"Expected sec_02 at top of RRF, got {fused_nids[0]}"
+    
+    # Calculate expected mathematical RRF scores:
+    # sec_02: 0.6/(60+1) + 0.4/(60+1) = 1.0 / 61 = ~0.016393
+    # sec_01: 0.6/(60+2) + 0.4/(60+3) = 0.6/62 + 0.4/63 = ~0.009677 + ~0.006349 = ~0.016026
+    # sec_03: 0.0 + 0.4/(60+2) = 0.4/62 = ~0.006452
+    score_sec02 = dict(fused)["sec_02"]
+    score_sec01 = dict(fused)["sec_01"]
+    score_sec03 = dict(fused)["sec_03"]
+    assert score_sec02 > score_sec01 > score_sec03
+    print(f"✓ Test 3 Passed: Weighted RRF correctly prioritized agreed nodes (Scores: sec_02={score_sec02:.6f}, sec_01={score_sec01:.6f}, sec_03={score_sec03:.6f})!")
+
+
+def test_parallel_hybrid_retrieval_multi_doc():
+    """Verify concurrent execution of tree search + BM25 across multiple attached documents."""
+    class MockLLMService:
+        async def llm_tree_search(self, query, tree, chat_history=None, flat_sections=None):
+            # Simulate ~30ms async LLM search
+            await asyncio.sleep(0.03)
+            # LLM picks section 1
+            return {"node_list": ["doc1_sec1"]}
+
+        def find_nodes_by_ids(self, tree, target_ids, node_map=None):
+            if node_map:
+                return [node_map[nid] for nid in target_ids if nid in node_map]
+            return [{"node_id": nid, "title": f"Title {nid}", "text": "Content"} for nid in target_ids]
+
+    mock_llm = MockLLMService()
+    rag = VectorlessRAG(mock_llm)
+
+    doc1 = {
+        "doc_id": "doc_1",
+        "filename": "Doc1.pdf",
+        "tree": [],
+        "node_map": {
+            "doc1_sec1": {"node_id": "doc1_sec1", "title": "Overview", "text": "General overview"},
+            "doc1_sec2": {"node_id": "doc1_sec2", "title": "ERNIE Details", "text": "ERNIE TransE architecture"},
+        },
+        "flat_sections": [
+            {"node_id": "doc1_sec1", "title": "Overview", "summary": "General overview"},
+            {"node_id": "doc1_sec2", "title": "ERNIE Details", "summary": "ERNIE TransE architecture"},
+        ],
+    }
+
+    doc2 = {
+        "doc_id": "doc_2",
+        "filename": "Doc2.pdf",
+        "tree": [],
+        "node_map": {
+            "doc2_sec1": {"node_id": "doc2_sec1", "title": "Financial Results", "text": "Revenue and profit"},
+        },
+        "flat_sections": [
+            {"node_id": "doc2_sec1", "title": "Financial Results", "summary": "Revenue and profit"},
+        ],
+    }
+
+    async def run_test():
+        query = "What is ERNIE TransE architecture?"
+        t0 = time.monotonic()
+        results_generator = rag.answer_multi_doc_query_stream(
+            query=query,
+            docs=[doc1, doc2],
+            chat_history=[],
+        )
+        # Collect first yielded stage event to measure retrieval latency
+        stage_event = await results_generator.__anext__()
+        total_retrieval_ms = (time.monotonic() - t0) * 1000
+        print(f"Parallel hybrid retrieval across 2 docs took {total_retrieval_ms:.2f}ms")
+        assert "STAGE" in stage_event
+        assert total_retrieval_ms < 150.0  # Concurrently waited for LLM ~30ms + BM25 ~1ms
+        print("✓ Test 4 Passed: Multi-document parallel hybrid retrieval verified!")
+
+    asyncio.run(run_test())
+
+
+if __name__ == "__main__":
+    print("\n--- RUNNING PHASE 2 HYBRID RETRIEVAL (BM25 + RRF) TEST SUITE ---")
+    test_bm25_node_index_scoring_and_latency()
+    test_bm25_arabic_support_and_title_boost()
+    test_weighted_reciprocal_rank_fusion()
+    test_parallel_hybrid_retrieval_multi_doc()
+    print("\n=======================================================")
+    print(" ALL 4 PHASE 2 HYBRID RETRIEVAL TESTS PASSED PERFECTLY!")
+    print("=======================================================\n")