feat(rag): RAGRetriever (load + search → chunks with scores)

src/rag/retrieve.py

@@ -0,0 +1,40 @@
+"""Query → top-k chunks. Encapsulates the embedder + store pair so callers
+don't have to assemble both. Loads from disk lazily.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from src.core.logger import get_logger
+from src.rag.embed import EMBEDDING_DIM, Embedder
+from src.rag.store import FAISSStore
+
+logger = get_logger(__name__)
+
+
+class RAGRetriever:
+    """Bundle (embedder, store). Use `RAGRetriever.load(dir)` to construct."""
+
+    def __init__(self, store: FAISSStore, embedder: Embedder) -> None:
+        self._store = store
+        self._embedder = embedder
+
+    @classmethod
+    def load(cls, index_dir: Path) -> "RAGRetriever":
+        store = FAISSStore.load(Path(index_dir), dim=EMBEDDING_DIM)
+        return cls(store=store, embedder=Embedder())
+
+    def __len__(self) -> int:
+        return len(self._store)
+
+    def search(self, query: str, k: int = 5) -> list[dict]:
+        """Return up to `k` chunks most relevant to `query`, sorted by score desc.
+
+        Each chunk dict carries `text`, `source`, `chunk_index`, `score`.
+        Returns [] for empty query or empty store.
+        """
+        if not query.strip() or len(self._store) == 0:
+            return []
+        vec = self._embedder.encode([query])
+        hits = self._store.search(vec[0], k=k)
+        return [{**chunk, "score": score} for chunk, score in hits]

tests/rag/test_retrieve.py

@@ -0,0 +1,45 @@
+"""Tests for src.rag.retrieve — query → top-k chunks."""
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+from src.rag.ingest import ingest_directory
+from src.rag.retrieve import RAGRetriever
+
+
+_FIXTURE_KB = Path(__file__).parent.parent / "fixtures" / "kb_sample"
+
+
+class TestRAGRetriever:
+    @pytest.fixture(scope="class")
+    def retriever(self, tmp_path_factory: pytest.TempPathFactory) -> RAGRetriever:
+        idx_dir = tmp_path_factory.mktemp("rag_idx")
+        ingest_directory(_FIXTURE_KB, idx_dir)
+        return RAGRetriever.load(idx_dir)
+
+    def test_bbb_query_returns_lipinski_chunk(self, retriever: RAGRetriever) -> None:
+        hits = retriever.search("Why does ethanol cross the blood-brain barrier?", k=3)
+        assert len(hits) == 3
+        sources = [h["source"] for h in hits]
+        assert "lipinski_rule_of_five.md" in sources
+        # top hit should be from lipinski
+        assert hits[0]["source"] == "lipinski_rule_of_five.md"
+
+    def test_combat_query_returns_combat_chunk(self, retriever: RAGRetriever) -> None:
+        hits = retriever.search("How does ComBat remove scanner bias from MRI data?", k=2)
+        assert hits[0]["source"] == "combat_harmonization_primer.md"
+
+    def test_eeg_query_returns_ica_chunk(self, retriever: RAGRetriever) -> None:
+        hits = retriever.search("How do you remove eye blink artifacts from EEG?", k=2)
+        assert hits[0]["source"] == "mne_ica_basics.md"
+
+    def test_search_includes_score_and_text(self, retriever: RAGRetriever) -> None:
+        hits = retriever.search("BBB permeability", k=1)
+        h = hits[0]
+        assert "text" in h
+        assert "source" in h
+        assert "score" in h
+        assert isinstance(h["score"], float)
+        assert 0.0 <= h["score"] <= 1.0