Python context layer: lazy BM25 retrieval and prewarm

pkg-py/README.md

@@ -2,6 +2,6 @@
 
 `commons` is a constructor for trustworthy data agents. Once implemented, this package will give an LLM data, semantic, and context layers to work with, tools for querying them, and A/B/C provenance tags so every answer carries a classification as to its trustworthiness.
 
-**Status: pre-alpha.** The package exports the data and semantic layers (`data_source`, `list_tables`, `measure`, `semantic_layer`, and their types) and the context layer (`context_layer()` and the `ContextLayer` it returns) for reading text or Markdown files; the agent constructor that ties them together is not implemented yet. Python 3.11 or later is required.
+**Status: pre-alpha.** The package exports the data and semantic layers (`data_source`, `list_tables`, `measure`, `semantic_layer`, and their types) and the context layer (`context_layer()` and the `ContextLayer` it returns), which reads text or Markdown files and retrieves from them; the agent constructor that ties them together is not implemented yet. Python 3.11 or later is required.
 
 Behavior that both implementations must agree on belongs in [`tests/shared/`](https://github.com/posit-dev/commons/tree/main/tests/shared) at the repository root, which that directory's README defines as the authority. The provenance tag rules and display copy, the citation dialect, and the context layer's frontmatter handling are governed that way; both suites run those cases.

pkg-py/src/commons/_context_layer.py

@@ -10,8 +10,13 @@
 
 import os
 import re
+import threading
 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 +39,8 @@ class ContextLayer:
 
     def __init__(self, docs: Iterable[str] = ()) -> None:
         self._docs = tuple(docs)
+        self._store_cache: DuckDBStore | None = None
+        self._store_lock = threading.Lock()
 
     @property
     def docs(self) -> tuple[str, ...]:
@@ -44,6 +51,65 @@ 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. The lock keeps
+    # concurrent first searches on a shared layer from each building a
+    # store and discarding all but one.
+    def _store(self) -> DuckDBStore:
+        with self._store_lock:
+            if self._store_cache is None:
+                # TODO: emit the commons_context_store_build span the R
+                # store build emits, once pkg-py has its tracing module.
+                store = DuckDBStore.create(location=":memory:", embed=None)
+                chunker = MarkdownChunker()
+                # ingest() requires each document's origin to be distinct
+                # and non-empty, so each gets a distinct synthetic origin.
+                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.
+
+        Optional and idempotent; worth calling when a search is known to be
+        coming, so its cost does not land on the first user turn.
+        """
+        if self._docs:
+            self._store()
+
+    # Public ahead of the R counterpart: context_search() in
+    # pkg-r/R/context-layer.R is internal there and spells the limit `n`.
+    def search(self, query: str, top_k: int = 3) -> list[str]:
+        """Retrieve the chunks most relevant to ``query``.
+
+        Returns chunk texts, best match first, at most ``top_k`` of them.
+        An empty layer, or a query nothing matches, returns an empty list.
+        ``top_k`` must be at least 1.
+        """
+        if top_k < 1:
+            raise ValueError(f"top_k must be at least 1, not {top_k}.")
+        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

@@ -1,4 +1,8 @@
+import time
+from concurrent.futures import ThreadPoolExecutor
+
 import pytest
+from raghilda.store import DuckDBStore
 
 from commons import ContextLayer, context_layer
 from commons._context_layer import strip_frontmatter
@@ -69,3 +73,162 @@ 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
+
+
+def test_search_rejects_a_non_positive_top_k(tmp_path):
+    path = tmp_path / "notes.md"
+    path.write_text("# Revenue\nRevenue excludes tax.")
+    layer = context_layer(files=[path])
+
+    with pytest.raises(ValueError, match="top_k"):
+        layer.search("revenue", top_k=0)
+
+
+def test_search_results_are_the_same_after_prewarm(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."
+    )
+
+    cold = context_layer(files=[path]).search("revenue")
+    warm_layer = context_layer(files=[path])
+    warm_layer.prewarm()
+
+    assert warm_layer.search("revenue") == cold
+
+
+def test_concurrent_first_searches_build_one_store(tmp_path, monkeypatch):
+    path = tmp_path / "notes.md"
+    path.write_text("# Revenue\nRevenue excludes tax.")
+    layer = context_layer(files=[path])
+
+    builds = 0
+    real_create = DuckDBStore.create
+
+    def counting_create(*args, **kwargs):
+        nonlocal builds
+        builds += 1
+        time.sleep(0.05)  # widen the race window
+        return real_create(*args, **kwargs)
+
+    monkeypatch.setattr(DuckDBStore, "create", counting_create)
+
+    with ThreadPoolExecutor(max_workers=8) as pool:
+        results = list(pool.map(lambda _: layer.search("revenue"), range(8)))
+
+    assert builds == 1
+    assert all(result == results[0] for result in results)