Python context layer: lazy BM25 retrieval and prewarm

pkg-py/src/commons/_context_layer.py

@@ -12,6 +12,10 @@
 import re
 from collections.abc import Iterable
 
+from raghilda.chunker import MarkdownChunker
+from raghilda.document import MarkdownDocument
+from raghilda.store import DuckDBStore
+
 __all__ = ["ContextLayer", "context_layer"]
 
 # Frontmatter carries file metadata (e.g. provenance) meant for maintainers,
@@ -34,6 +38,7 @@ class ContextLayer:
 
     def __init__(self, docs: Iterable[str] = ()) -> None:
         self._docs = tuple(docs)
+        self._store_cache: DuckDBStore | None = None
 
     @property
     def docs(self) -> tuple[str, ...]:
@@ -43,6 +48,46 @@ def __repr__(self) -> str:
         n = len(self._docs)
         return f"<ContextLayer: {n} document{'' if n == 1 else 's'}>"
 
+    # Store setup (duckdb creation, chunk insertion, BM25 indexing) is the
+    # most expensive part of building an agent and many conversations never
+    # search, so it is deferred to the first search.
+    def _store(self) -> DuckDBStore:
+        if self._store_cache is None:
+            store = DuckDBStore.create(location=":memory:", embed=None)
+            chunker = MarkdownChunker()
+            # ingest() upserts on origin and rejects an empty one, so each
+            # document gets a distinct synthetic origin. Two files with
+            # identical text would otherwise collapse into one chunk.
+            store.ingest(
+                [
+                    MarkdownDocument(content=doc, origin=f"commons-context-{i}")
+                    for i, doc in enumerate(self._docs)
+                ],
+                prepare=chunker.chunk,
+            )
+            store.build_index(type="bm25")
+            self._store_cache = store
+        return self._store_cache
+
+    def prewarm(self) -> None:
+        """Build the index now so the first search does not pay for it."""
+        if self._docs:
+            self._store()
+
+    def search(self, query: str, top_k: int = 3) -> list[str]:
+        """Retrieve the chunks most relevant to ``query``."""
+        if not self._docs:
+            return []
+        hits = self._store().retrieve_bm25(query, top_k=top_k)
+        # retrieve_bm25 pads its result up to top_k with unscored rows, so a
+        # query that matches nothing still comes back full. Only scored rows
+        # are hits.
+        return [
+            hit.text.strip()
+            for hit in hits
+            if any(m.name == "bm25" and m.value is not None for m in hit.metrics)
+        ]
+
 
 def context_layer(
     files: Iterable[str | os.PathLike[str]] = (),

pkg-py/tests/test_context_layer.py

@@ -57,3 +57,116 @@ def test_context_layer_repr_counts_documents(tmp_path):
 
     assert repr(context_layer()) == "<ContextLayer: 0 documents>"
     assert repr(context_layer(files=[path])) == "<ContextLayer: 1 document>"
+
+
+def test_search_finds_a_relevant_chunk(tmp_path):
+    path = tmp_path / "notes.md"
+    path.write_text(
+        "# Revenue\nRevenue excludes tax unless stated otherwise.\n\n"
+        "# Discounts\nDiscounts are applied before tax."
+    )
+
+    hits = context_layer(files=[path]).search("what does revenue mean")
+
+    assert len(hits) >= 1
+    assert "tax" in hits[0]
+
+
+def test_search_returns_nothing_when_the_layer_is_empty():
+    assert context_layer().search("anything") == []
+
+
+def test_search_returns_nothing_when_no_chunk_matches(tmp_path):
+    path = tmp_path / "notes.md"
+    path.write_text("# A\napples")
+
+    assert context_layer(files=[path]).search("zzzzz") == []
+
+
+def test_search_does_not_surface_stripped_frontmatter(tmp_path):
+    path = tmp_path / "notes.md"
+    path.write_text(
+        "---\nprovenance: abc1234\n---\n"
+        "# Revenue\nRevenue excludes tax unless stated otherwise."
+    )
+
+    layer = context_layer(files=[path])
+
+    assert "tax" in layer.search("revenue")[0]
+    assert layer.search("abc1234") == []
+
+
+def test_search_reuses_the_store_across_calls(tmp_path):
+    path = tmp_path / "notes.md"
+    path.write_text("# Revenue\nRevenue excludes tax.")
+    layer = context_layer(files=[path])
+
+    layer.search("revenue")
+    first = layer._store_cache
+    layer.search("revenue")
+
+    assert first is not None
+    assert layer._store_cache is first
+
+
+def test_search_indexes_every_document(tmp_path):
+    first = tmp_path / "a.md"
+    first.write_text("# Revenue\nRevenue excludes tax.")
+    second = tmp_path / "b.md"
+    second.write_text("# Discounts\nDiscounts are applied before tax.")
+
+    layer = context_layer(files=[first, second])
+
+    assert "Discounts" in layer.search("discounts")[0]
+    assert "Revenue" in layer.search("revenue")[0]
+
+
+def test_search_indexes_identical_documents_separately(tmp_path):
+    first = tmp_path / "a.md"
+    first.write_text("# Revenue\nRevenue excludes tax.")
+    second = tmp_path / "b.md"
+    second.write_text("# Revenue\nRevenue excludes tax.")
+
+    layer = context_layer(files=[first, second])
+
+    assert len(layer.search("revenue", top_k=5)) == 2
+
+
+def test_search_respects_top_k(tmp_path):
+    for i in range(5):
+        (tmp_path / f"{i}.md").write_text(f"# Revenue {i}\nRevenue excludes tax.")
+
+    layer = context_layer(files=sorted(tmp_path.glob("*.md")))
+
+    assert len(layer.search("revenue", top_k=2)) == 2
+
+
+def test_prewarm_builds_the_store_ahead_of_search(tmp_path):
+    path = tmp_path / "notes.md"
+    path.write_text("# Revenue\nRevenue excludes tax.")
+    layer = context_layer(files=[path])
+
+    assert layer._store_cache is None
+    layer.prewarm()
+
+    assert layer._store_cache is not None
+    assert layer._store_cache.size() == 1
+
+
+def test_prewarm_on_an_empty_layer_builds_nothing():
+    layer = context_layer()
+    layer.prewarm()
+
+    assert layer._store_cache is None
+
+
+def test_prewarm_is_idempotent(tmp_path):
+    path = tmp_path / "notes.md"
+    path.write_text("# Revenue\nRevenue excludes tax.")
+    layer = context_layer(files=[path])
+
+    layer.prewarm()
+    first = layer._store_cache
+    layer.prewarm()
+
+    assert layer._store_cache is first