diff --git a/pkg-py/src/commons/_context_layer.py b/pkg-py/src/commons/_context_layer.py index 9cc829d3..847fb062 100644 --- a/pkg-py/src/commons/_context_layer.py +++ b/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"" + # 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]] = (), diff --git a/pkg-py/tests/test_context_layer.py b/pkg-py/tests/test_context_layer.py index 294e9597..92f94769 100644 --- a/pkg-py/tests/test_context_layer.py +++ b/pkg-py/tests/test_context_layer.py @@ -57,3 +57,116 @@ def test_context_layer_repr_counts_documents(tmp_path): assert repr(context_layer()) == "" assert repr(context_layer(files=[path])) == "" + + +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