diff --git a/.gitignore b/.gitignore index 63b31ce..9e0a7f9 100644 --- a/.gitignore +++ b/.gitignore @@ -26,3 +26,6 @@ Thumbs.db # Application data *.log + +# Ralph loop +progress.md diff --git a/app.py b/app.py index 686cd85..1ef509d 100644 --- a/app.py +++ b/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="") diff --git a/src/retrieval/bm25_index.py b/src/retrieval/bm25_index.py new file mode 100644 index 0000000..7458b0e --- /dev/null +++ b/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 + ] diff --git a/src/retrieval/fusion.py b/src/retrieval/fusion.py new file mode 100644 index 0000000..7b915c8 --- /dev/null +++ b/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 + ] diff --git a/src/retrieval/hybrid.py b/src/retrieval/hybrid.py new file mode 100644 index 0000000..74e14b8 --- /dev/null +++ b/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, + ) diff --git a/src/retrieval/vector_store.py b/src/retrieval/vector_store.py index fb792e0..7e9124e 100644 --- a/src/retrieval/vector_store.py +++ b/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.""" diff --git a/tests/test_bm25_index.py b/tests/test_bm25_index.py new file mode 100644 index 0000000..e9a234f --- /dev/null +++ b/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 diff --git a/tests/test_fusion.py b/tests/test_fusion.py new file mode 100644 index 0000000..29da76e --- /dev/null +++ b/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 == [] diff --git a/tests/test_hybrid.py b/tests/test_hybrid.py new file mode 100644 index 0000000..06a2620 --- /dev/null +++ b/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