Python context layer: fold source prose into retrieval

pkg-py/src/commons/_context_layer.py

@@ -20,6 +20,7 @@
 
 if TYPE_CHECKING:
     from ._data_dictionary import DataDictionary
+    from ._data_source import DataSource
 
 __all__ = ["ContextLayer", "context_layer"]
 
@@ -126,6 +127,32 @@ def _dictionary_chunks(dictionary: DataDictionary | None) -> list[str]:
     return dictionary.context_chunks()
 
 
+def augment_context_layer(
+    layer: ContextLayer | None, sources: Iterable[DataSource]
+) -> ContextLayer | None:
+    """Fold each source's prose into a layer the agent can retrieve from.
+
+    Returns a new layer, leaving the caller's untouched: source enrichment
+    belongs to the agent that owns the sources, so mutating the argument
+    would leak one agent's sources into the next agent built from the same
+    layer. With nothing to add, the argument comes back as it went in,
+    ``None`` included, so an agent with neither context nor a dictionary
+    has no layer rather than an empty one.
+    """
+    chunks: list[str] = []
+    for source in sources:
+        chunks.extend(_dictionary_chunks(source.dictionary))
+        # A warehouse's own semantic models contribute their retrieval prose
+        # here too, as they do in pkg-r/R/context-layer.R. This package has
+        # no semantic-model surface yet, so there is nothing to fold in.
+
+    if not chunks:
+        return layer
+
+    existing = layer.docs if layer is not None else ()
+    return ContextLayer([*existing, *chunks])
+
+
 def context_layer(
     files: Iterable[str | os.PathLike[str]] = (),
 ) -> ContextLayer:

pkg-py/tests/test_context_layer.py

@@ -1,11 +1,16 @@
 import time
 from concurrent.futures import ThreadPoolExecutor
 
+import pandas as pd
 import pytest
 from raghilda.store import DuckDBStore
 
-from commons import ContextLayer, context_layer
-from commons._context_layer import _dictionary_chunks, strip_frontmatter
+from commons import ContextLayer, context_layer, data_source
+from commons._context_layer import (
+    _dictionary_chunks,
+    augment_context_layer,
+    strip_frontmatter,
+)
 from commons._data_dictionary import DataDictionary
 
 from ._shared import load_shared_fixture
@@ -248,3 +253,107 @@ def test_dictionary_context_chunks_shared_cases(case):
     dictionary = None if spec is None else DataDictionary.model_validate(spec)
 
     assert _dictionary_chunks(dictionary) == case["expected"]
+
+
+def a_source_with_dictionary(spec):
+    """A real source over an in-memory DuckDB, carrying the given dictionary."""
+    return data_source(
+        notes=pd.DataFrame({"n": [1]}),
+        dictionary=DataDictionary.model_validate(spec),
+    )
+
+
+def test_augment_appends_dictionary_chunks_to_a_new_layer(tmp_path):
+    path = tmp_path / "notes.md"
+    path.write_text("# Revenue\nRevenue excludes tax.")
+    layer = context_layer(files=[path])
+    source = a_source_with_dictionary({"details": "Orders from the retail system."})
+
+    augmented = augment_context_layer(layer, [source])
+
+    assert augmented is not None
+    assert augmented is not layer
+    assert layer.docs == ("# Revenue\nRevenue excludes tax.",)
+    assert augmented.docs == (
+        "# Revenue\nRevenue excludes tax.",
+        "Orders from the retail system.",
+    )
+
+
+def test_augment_creates_a_layer_when_there_was_none():
+    source = a_source_with_dictionary({"details": "Orders from the retail system."})
+
+    augmented = augment_context_layer(None, [source])
+
+    assert augmented is not None
+    assert augmented.docs == ("Orders from the retail system.",)
+
+
+def test_augment_returns_the_argument_unchanged_when_there_is_nothing_to_add(tmp_path):
+    path = tmp_path / "notes.md"
+    path.write_text("# Revenue\nRevenue excludes tax.")
+    layer = context_layer(files=[path])
+    bare = data_source(notes=pd.DataFrame({"n": [1]}))
+
+    assert augment_context_layer(layer, [bare]) is layer
+    assert augment_context_layer(None, [bare]) is None
+    assert augment_context_layer(layer, []) is layer
+    assert augment_context_layer(None, []) is None
+
+
+def test_augment_gives_the_new_layer_a_fresh_store(tmp_path):
+    path = tmp_path / "notes.md"
+    path.write_text("# Revenue\nRevenue excludes tax.")
+    layer = context_layer(files=[path])
+    layer.prewarm()
+    source = a_source_with_dictionary({"details": "Orders from the retail system."})
+
+    augmented = augment_context_layer(layer, [source])
+
+    assert augmented is not None
+    assert augmented._store_cache is None
+    assert "Orders" in augmented.search("orders retail")[0]
+
+
+# An empty case list would make the parametrized test below vacuously pass.
+def test_the_augment_fixture_is_not_empty():
+    assert SHARED["augment_context_layer"]["cases"]
+
+
+@pytest.mark.parametrize(
+    "case", SHARED["augment_context_layer"]["cases"], ids=lambda c: c["name"]
+)
+def test_augment_context_layer_shared_cases(case):
+    layer = None if case["docs"] is None else ContextLayer(case["docs"])
+    sources = [
+        data_source(
+            notes=pd.DataFrame({"n": [1]}),
+            dictionary=None if spec is None else DataDictionary.model_validate(spec),
+        )
+        for spec in case["dictionaries"]
+    ]
+
+    augmented = augment_context_layer(layer, sources)
+
+    if case["expected_docs"] is None:
+        assert augmented is None
+    else:
+        assert augmented is not None
+        assert list(augmented.docs) == case["expected_docs"]
+
+
+def test_a_dictionary_alone_is_searchable_context():
+    source = a_source_with_dictionary(
+        {
+            "details": None,
+            "tables": {},
+            "glossary": {"AOV": "Average order value, revenue divided by order count."},
+        }
+    )
+
+    layer = augment_context_layer(None, [source])
+
+    assert layer is not None
+    assert layer.search("average order value") == [
+        "AOV: Average order value, revenue divided by order count."
+    ]

pkg-r/tests/testthat/fixtures/shared/context_layer.json

@@ -171,5 +171,102 @@
         ]
       }
     ]
+  },
+  "augment_context_layer": {
+    "description": "Folding each source's prose into an agent's context layer. The sources' chunks follow the caller's own documents and then each other in the order the sources were given, so the order a reader would expect is preserved. With no chunks to add, the layer is handed back as it was, absent included: an agent with neither context files nor dictionary prose has no layer rather than an empty one. `docs` and `expected_docs` are null for an absent layer and an empty list for a layer holding no documents; the two are not the same thing. Each entry in `dictionaries` is one source, null for a source carrying no dictionary. Whether the returned layer is a new object, and whether it carries a fresh index, is implementation detail and is not pinned here.",
+    "cases": [
+      {
+        "name": "source chunks follow the layer's own documents",
+        "docs": [
+          "Booked revenue excludes tax."
+        ],
+        "dictionaries": [
+          {
+            "details": "Orders from the retail system.",
+            "glossary": {
+              "AOV": "Average order value."
+            }
+          }
+        ],
+        "expected_docs": [
+          "Booked revenue excludes tax.",
+          "Orders from the retail system.",
+          "AOV: Average order value."
+        ]
+      },
+      {
+        "name": "a dictionary alone becomes the layer",
+        "docs": null,
+        "dictionaries": [
+          {
+            "details": "Orders from the retail system."
+          }
+        ],
+        "expected_docs": [
+          "Orders from the retail system."
+        ]
+      },
+      {
+        "name": "every source contributes, in the order given",
+        "docs": null,
+        "dictionaries": [
+          {
+            "details": "Orders from the retail system."
+          },
+          null,
+          {
+            "glossary": {
+              "AOV": "Average order value."
+            }
+          }
+        ],
+        "expected_docs": [
+          "Orders from the retail system.",
+          "AOV: Average order value."
+        ]
+      },
+      {
+        "name": "a source with no dictionary leaves the layer alone",
+        "docs": [
+          "Booked revenue excludes tax."
+        ],
+        "dictionaries": [
+          null
+        ],
+        "expected_docs": [
+          "Booked revenue excludes tax."
+        ]
+      },
+      {
+        "name": "a dictionary with no prose leaves the layer absent",
+        "docs": null,
+        "dictionaries": [
+          {
+            "details": null,
+            "tables": {
+              "orders": {
+                "description": null
+              }
+            },
+            "glossary": {}
+          }
+        ],
+        "expected_docs": null
+      },
+      {
+        "name": "an empty layer stays an empty layer",
+        "docs": [],
+        "dictionaries": [
+          null
+        ],
+        "expected_docs": []
+      },
+      {
+        "name": "no sources and no layer stays absent",
+        "docs": null,
+        "dictionaries": [],
+        "expected_docs": null
+      }
+    ]
   }
 }

pkg-r/tests/testthat/test-context-layer.R

@@ -27,6 +27,35 @@ test_that("dictionary_context_chunks matches the shared cases", {
   }
 })
 
+test_that("augment_context_layer matches the shared cases", {
+  cases <- shared_fixture("context_layer")$augment_context_layer$cases
+  expect_gt(length(cases), 0)
+
+  for (case in cases) {
+    layer <- if (is.null(case$docs)) {
+      NULL
+    } else {
+      new_context_layer(as.character(unlist(case$docs)))
+    }
+    sources <- lapply(case$dictionaries, function(spec) {
+      dictionary <- if (is.null(spec)) NULL else new_data_dictionary(spec)
+      suppressMessages(data_source(sales = test_sales(), dictionary = dictionary))
+    })
+
+    augmented <- augment_context_layer(layer, sources)
+
+    if (is.null(case$expected_docs)) {
+      expect_null(augmented, info = case$name)
+    } else {
+      expect_identical(
+        context_layer_state(augmented)$docs,
+        as.character(unlist(case$expected_docs)),
+        info = case$name
+      )
+    }
+  }
+})
+
 test_that("context_layer indexes files and finds relevant chunks", {
   path <- withr::local_tempfile(fileext = ".md")
   writeLines(

pkg-r/tests/testthat/test-data-dictionary.R

@@ -473,20 +473,6 @@ test_that("dictionary prose is searchable via the context layer", {
   )
 })
 
-test_that("augmenting keeps existing context docs", {
-  skip_if_not_installed("yaml")
-  path <- withr::local_tempfile(fileext = ".md")
-  writeLines("Booked revenue excludes tax.", path)
-  layer <- context_layer(files = path)
-  augmented <- augment_context_layer(layer, list(local_dict_source()))
-
-  expect_true("Booked revenue excludes tax." %in% context_layer_state(augmented)$docs)
-  expect_gt(length(context_layer_state(augmented)$docs), length(context_layer_state(layer)$docs))
-})
-
-test_that("augmenting without dictionaries is a no-op", {
-  expect_null(augment_context_layer(NULL, list(test_source())))
-})
 
 test_that("agent tools share first-touch state", {
   skip_if_not_installed("yaml")