Added hybrid search: BM25 + semantic + RRF fusion

.gitignore

@@ -26,3 +26,6 @@ Thumbs.db
 
 # Application data
 *.log
+
+# Ralph loop
+progress.md

app.py

@@ -1,3 +1,6 @@
+import logging
+import time
+
 import chainlit as cl
 import chromadb
 
@@ -6,9 +9,13 @@
 from src.health_check import check_models, check_ollama
 from src.ingestion.chunker import chunk_documents
 from src.ingestion.loader import load_folder
+from src.retrieval.bm25_index import BM25Index
 from src.retrieval.embeddings import make_ollama_embed_fn
+from src.retrieval.hybrid import HybridRetriever
 from src.retrieval.vector_store import VectorStore
 
+logger = logging.getLogger(__name__)
+
 
 @cl.on_chat_start
 async def on_chat_start():
@@ -58,10 +65,23 @@ async def on_chat_start():
             )
             store.add_chunks(chunks)
             doc_count = store.count()
-            await cl.Message(
-                content=f"Ingested {len(docs)} pages into {doc_count} chunks. Ask me anything!"
-            ).send()
-            return
+
+    # Build BM25 index from stored chunks
+    bm25_index = BM25Index()
+    if doc_count > 0:
+        start = time.time()
+        texts, metadatas = store.get_all_texts_and_metadatas()
+        bm25_index.build(texts, metadatas)
+        elapsed = time.time() - start
+        logger.info("BM25 index built from %d chunks in %.2fs", doc_count, elapsed)
+
+    # Create hybrid retriever
+    retriever = HybridRetriever(
+        vector_store=store,
+        bm25_index=bm25_index,
+        config=config.retrieval,
+    )
+    cl.user_session.set("retriever", retriever)
 
     if doc_count > 0:
         await cl.Message(
@@ -76,22 +96,20 @@ async def on_chat_start():
 @cl.on_message
 async def on_message(message: cl.Message):
     config = cl.user_session.get("config")
-    store = cl.user_session.get("store")
+    retriever = cl.user_session.get("retriever")
 
-    if store is None:
-        await cl.Message(content="No document store available. Please restart the app.").send()
+    if retriever is None:
+        await cl.Message(content="No retriever available. Please restart the app.").send()
         return
 
-    if store.count() == 0:
+    store = cl.user_session.get("store")
+    if store is None or store.count() == 0:
         await cl.Message(
             content="No documents indexed yet. Configure your documents folder in `config.yaml` first."
         ).send()
         return
 
-    results = store.search(
-        message.content,
-        k=config.retrieval.semantic_top_k,
-    )
+    results = retriever.retrieve(message.content)
 
     # Stream the answer token by token
     msg = cl.Message(content="")

src/retrieval/bm25_index.py

@@ -0,0 +1,42 @@
+from rank_bm25 import BM25Okapi
+
+from src.models import SearchResult
+
+# --- Not via LangChain: BM25 keyword search uses rank_bm25 directly. ---
+
+
+class BM25Index:
+    """BM25 keyword search index built from chunk texts."""
+
+    def __init__(self) -> None:
+        self._index: BM25Okapi | None = None
+        self._texts: list[str] = []
+        self._metadatas: list[dict] = []
+
+    def build(self, texts: list[str], metadatas: list[dict]) -> None:
+        """Build the BM25 index from chunk texts and their metadata."""
+        self._texts = texts
+        self._metadatas = metadatas
+        tokenised = [text.lower().split() for text in texts]
+        self._index = BM25Okapi(tokenised)
+
+    def search(self, query: str, k: int = 20) -> list[SearchResult]:
+        """Return top-k results ranked by BM25 score."""
+        if self._index is None or not self._texts:
+            return []
+
+        tokenised_query = query.lower().split()
+        scores = self._index.get_scores(tokenised_query)
+
+        # Get top-k indices sorted by score descending
+        ranked = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:k]
+
+        return [
+            SearchResult(
+                text=self._texts[i],
+                metadata=self._metadatas[i],
+                distance=float(scores[i]),
+            )
+            for i in ranked
+            if scores[i] > 0
+        ]

src/retrieval/fusion.py

@@ -0,0 +1,32 @@
+from src.models import SearchResult
+
+
+def reciprocal_rank_fusion(
+    *ranked_lists: list[SearchResult],
+    k: int = 60,
+    top_n: int = 30,
+) -> list[SearchResult]:
+    """Merge ranked lists using Reciprocal Rank Fusion.
+
+    Score for each result = sum(1 / (k + rank)) across all lists it appears in.
+    Deduplicates by chunk text. Returns top_n results sorted by RRF score.
+    """
+    scores: dict[str, float] = {}
+    results_by_text: dict[str, SearchResult] = {}
+
+    for ranked_list in ranked_lists:
+        for rank, result in enumerate(ranked_list):
+            scores[result.text] = scores.get(result.text, 0.0) + 1.0 / (k + rank + 1)
+            if result.text not in results_by_text:
+                results_by_text[result.text] = result
+
+    sorted_texts = sorted(scores, key=lambda t: scores[t], reverse=True)[:top_n]
+
+    return [
+        SearchResult(
+            text=results_by_text[text].text,
+            metadata=results_by_text[text].metadata,
+            distance=scores[text],
+        )
+        for text in sorted_texts
+    ]

src/retrieval/hybrid.py

@@ -0,0 +1,31 @@
+from src.config import RetrievalConfig
+from src.models import SearchResult
+from src.retrieval.bm25_index import BM25Index
+from src.retrieval.fusion import reciprocal_rank_fusion
+from src.retrieval.vector_store import VectorStore
+
+
+class HybridRetriever:
+    """Orchestrates BM25 + semantic search with RRF fusion."""
+
+    def __init__(
+        self,
+        vector_store: VectorStore,
+        bm25_index: BM25Index,
+        config: RetrievalConfig,
+    ) -> None:
+        self._vector_store = vector_store
+        self._bm25_index = bm25_index
+        self._config = config
+
+    def retrieve(self, query: str) -> list[SearchResult]:
+        """Run hybrid search: BM25 + semantic → RRF fusion."""
+        bm25_results = self._bm25_index.search(query, k=self._config.bm25_top_k)
+        semantic_results = self._vector_store.search(query, k=self._config.semantic_top_k)
+
+        return reciprocal_rank_fusion(
+            bm25_results,
+            semantic_results,
+            k=self._config.rrf_k,
+            top_n=30,
+        )

src/retrieval/vector_store.py

@@ -71,10 +71,15 @@ def count(self) -> int:
         """Return the number of chunks in the collection."""
         return self._collection.count()
 
+    def get_all_texts_and_metadatas(self) -> tuple[list[str], list[dict]]:
+        """Return all stored chunk texts and their metadata."""
+        result = self._collection.get()
+        return result["documents"], result["metadatas"]
+
     def get_all_texts(self) -> list[str]:
         """Return all stored chunk texts."""
-        result = self._collection.get()
-        return result["documents"]
+        texts, _ = self.get_all_texts_and_metadatas()
+        return texts
 
     def has_document(self, doc_hash: str) -> bool:
         """Check if any chunk with the given doc_hash exists."""

tests/test_bm25_index.py

@@ -0,0 +1,55 @@
+from src.retrieval.bm25_index import BM25Index
+
+
+def _build_index():
+    """Build a BM25 index from a known corpus."""
+    texts = [
+        "The capital of France is Paris",
+        "Berlin is the capital of Germany",
+        "Python is a programming language",
+        "Machine learning uses neural networks",
+        "Paris has the Eiffel Tower",
+    ]
+    metadatas = [
+        {"filename": "geo.pdf", "doc_type": "pdf", "page_number": 1, "chunk_index": i, "doc_hash": "h1"}
+        for i in range(len(texts))
+    ]
+    index = BM25Index()
+    index.build(texts, metadatas)
+    return index
+
+
+def test_build_and_search_returns_results():
+    """Searching a built index returns non-empty results."""
+    index = _build_index()
+    results = index.search("capital of France")
+    assert len(results) > 0
+
+
+def test_search_returns_k_results():
+    """search() respects the k parameter."""
+    index = _build_index()
+    results = index.search("capital", k=2)
+    assert len(results) == 2
+
+
+def test_search_relevance():
+    """Query about France returns the France chunk first."""
+    index = _build_index()
+    results = index.search("capital of France", k=3)
+    assert "France" in results[0].text or "Paris" in results[0].text
+
+
+def test_search_empty_index():
+    """Searching an unbuilt index returns empty list."""
+    index = BM25Index()
+    results = index.search("anything")
+    assert results == []
+
+
+def test_search_results_have_metadata():
+    """Search results include correct metadata from the corpus."""
+    index = _build_index()
+    results = index.search("Python programming", k=1)
+    assert results[0].metadata["filename"] == "geo.pdf"
+    assert "chunk_index" in results[0].metadata

tests/test_fusion.py

@@ -0,0 +1,54 @@
+from src.models import SearchResult
+from src.retrieval.fusion import reciprocal_rank_fusion
+
+
+def _make_result(text, distance=0.0):
+    return SearchResult(
+        text=text,
+        metadata={"filename": "test.pdf", "doc_type": "pdf", "page_number": 1, "chunk_index": 0, "doc_hash": "h1"},
+        distance=distance,
+    )
+
+
+def test_rrf_merges_two_lists():
+    """RRF merges two ranked lists into a single list."""
+    list_a = [_make_result("A"), _make_result("B")]
+    list_b = [_make_result("C"), _make_result("D")]
+    merged = reciprocal_rank_fusion(list_a, list_b, k=60)
+    texts = [r.text for r in merged]
+    assert "A" in texts
+    assert "C" in texts
+    assert len(merged) == 4
+
+
+def test_rrf_deduplicates():
+    """Same chunk in both lists appears once with combined score."""
+    list_a = [_make_result("A"), _make_result("B")]
+    list_b = [_make_result("B"), _make_result("C")]
+    merged = reciprocal_rank_fusion(list_a, list_b, k=60)
+    texts = [r.text for r in merged]
+    assert texts.count("B") == 1
+    # B should rank highest since it appears in both lists
+    assert merged[0].text == "B"
+
+
+def test_rrf_respects_top_n():
+    """RRF returns at most top_n results."""
+    list_a = [_make_result(f"A{i}") for i in range(10)]
+    list_b = [_make_result(f"B{i}") for i in range(10)]
+    merged = reciprocal_rank_fusion(list_a, list_b, k=60, top_n=5)
+    assert len(merged) == 5
+
+
+def test_rrf_single_list():
+    """RRF works with a single list (passthrough)."""
+    results = [_make_result("A"), _make_result("B")]
+    merged = reciprocal_rank_fusion(results, k=60)
+    assert len(merged) == 2
+    assert merged[0].text == "A"
+
+
+def test_rrf_empty_lists():
+    """RRF with empty lists returns empty."""
+    merged = reciprocal_rank_fusion([], [], k=60)
+    assert merged == []

tests/test_hybrid.py

@@ -0,0 +1,51 @@
+from unittest.mock import MagicMock
+
+from src.config import RetrievalConfig
+from src.models import SearchResult
+from src.retrieval.hybrid import HybridRetriever
+
+
+def _make_result(text, distance=0.0):
+    return SearchResult(
+        text=text,
+        metadata={"filename": "test.pdf", "doc_type": "pdf", "page_number": 1, "chunk_index": 0, "doc_hash": "h1"},
+        distance=distance,
+    )
+
+
+def test_retrieve_combines_bm25_and_semantic():
+    """HybridRetriever calls both BM25 and semantic search, fuses results."""
+    bm25 = MagicMock()
+    bm25.search.return_value = [_make_result("BM25 only"), _make_result("Shared")]
+
+    store = MagicMock()
+    store.search.return_value = [_make_result("Semantic only"), _make_result("Shared")]
+
+    config = RetrievalConfig(bm25_top_k=20, semantic_top_k=20, rrf_k=60)
+    retriever = HybridRetriever(vector_store=store, bm25_index=bm25, config=config)
+
+    results = retriever.retrieve("test query")
+
+    bm25.search.assert_called_once_with("test query", k=20)
+    store.search.assert_called_once_with("test query", k=20)
+    texts = [r.text for r in results]
+    assert "Shared" in texts
+    assert "BM25 only" in texts
+    assert "Semantic only" in texts
+    # Shared should rank highest (appears in both lists)
+    assert results[0].text == "Shared"
+
+
+def test_retrieve_returns_limited_results():
+    """HybridRetriever returns at most the configured number of results."""
+    bm25 = MagicMock()
+    bm25.search.return_value = [_make_result(f"BM25_{i}") for i in range(20)]
+
+    store = MagicMock()
+    store.search.return_value = [_make_result(f"Sem_{i}") for i in range(20)]
+
+    config = RetrievalConfig(bm25_top_k=20, semantic_top_k=20, rrf_k=60)
+    retriever = HybridRetriever(vector_store=store, bm25_index=bm25, config=config)
+
+    results = retriever.retrieve("query")
+    assert len(results) <= 30