refactor: make the gridded subset explicit about the spec

subset_gridded_to_2d and _spec_field_references now take a
kind: Literal['vega-lite', 'deckgl'] instead of sniffing both grammars at
once. The Vega-Lite path reads encoding field values and the deck.gl path
reads @@= accessor expressions, and each agent passes its own kind.

Author: ghostiee-11

lumen/ai/agents/deck_gl.py

@@ -300,7 +300,7 @@ async def respond(
         # DeckGL renders a static layer and cannot page a dimension, so once
         # the spec has picked its axes, collapse every gridded dim it does not
         # reference to a single slice (else a gridded source blows the row cap).
-        pipeline = subset_gridded_to_2d(pipeline, full_dict.get("spec", {}))
+        pipeline = subset_gridded_to_2d(pipeline, full_dict.get("spec", {}), "deckgl")
 
         # Create view and output
         view = self.view_type(pipeline=pipeline, **full_dict)

lumen/ai/agents/vega_lite.py

@@ -630,7 +630,7 @@ async def respond(
         # Vega-Lite cannot page a dimension, so once the spec has picked its
         # axes, collapse every gridded dim it does not reference to a single
         # slice (keeps e.g. a lon/time Hovmoller, drops an unused time on a map).
-        pipeline = subset_gridded_to_2d(pipeline, full_dict.get("spec", {}))
+        pipeline = subset_gridded_to_2d(pipeline, full_dict.get("spec", {}), "vega-lite")
 
         # Step 2: Show complete plot immediately
         view = self.view_type(pipeline=pipeline, **full_dict)

lumen/ai/utils.py

@@ -1532,20 +1532,23 @@ def _is_regular(values) -> bool:
     }
 
 
-def _spec_field_references(spec: Any) -> set[str]:
-    """Column names a generated view spec references: Vega-Lite encoding
-    ``field`` values and deck.gl ``@@=`` accessor identifiers. Lets the
-    gridded subset keep the dims a spec actually plots (an axis, color, ...)
-    and collapse the rest, without hard-coding which dims are spatial.
+def _spec_field_references(
+    spec: Any, kind: Literal['vega-lite', 'deckgl']
+) -> set[str]:
+    """Column names a generated view ``spec`` references, so the gridded subset
+    can keep the dims a spec actually plots (an axis, color, ...) and collapse
+    the rest. The two grammars reference columns differently, so the caller
+    states which it is: a Vega-Lite spec names columns in encoding ``field``
+    values, a deck.gl spec in ``@@=`` accessor expressions.
     """
     refs: set[str] = set()
 
     def walk(node):
         if isinstance(node, dict):
             for key, val in node.items():
-                if key == 'field' and isinstance(val, str):
+                if kind == 'vega-lite' and key == 'field' and isinstance(val, str):
                     refs.add(val)
-                elif isinstance(val, str) and val.startswith('@@='):
+                elif kind == 'deckgl' and isinstance(val, str) and val.startswith('@@='):
                     refs.update(re.findall(r'[A-Za-z_]\w*', val))
                 else:
                     walk(val)
@@ -1557,10 +1560,13 @@ def walk(node):
     return refs
 
 
-def subset_gridded_to_2d(pipeline: Pipeline, spec: Any) -> Pipeline:
+def subset_gridded_to_2d(
+    pipeline: Pipeline, spec: Any, kind: Literal['vega-lite', 'deckgl']
+) -> Pipeline:
     """Pin every gridded dimension the view ``spec`` does not reference to its
     first value, so a view that cannot page a dimension (VegaLite, DeckGL)
-    renders a single 2D slice instead of the full grid.
+    renders a single 2D slice instead of the full grid. ``kind`` selects the
+    spec grammar (see :func:`_spec_field_references`).
 
     Runs after the spec is generated, so the dims the model chose for x/y/color
     are kept -- a lon/time Hovmoller works -- and only the leftover dims (e.g.
@@ -1571,7 +1577,7 @@ def subset_gridded_to_2d(pipeline: Pipeline, spec: Any) -> Pipeline:
     if not md:
         return pipeline
     ds = pipeline.source.dataset
-    referenced = _spec_field_references(spec)
+    referenced = _spec_field_references(spec, kind)
     collapse = [dim for dim in ds.dims if dim not in referenced]
     if not collapse:
         return pipeline

lumen/tests/ai/test_gridded.py

@@ -240,41 +240,41 @@ def test_gridded_metadata_stays_agnostic_about_spatial_dims(simple_dataset_3d_pi
 
 def test_spec_field_references_vegalite_and_deckgl():
     """Field references are pulled from Vega-Lite encoding fields and deck.gl
-    ``@@=`` accessors alike."""
+    ``@@=`` accessors, each under its own explicit ``kind``."""
     vega = {"encoding": {"x": {"field": "lon"}, "y": {"field": "time"},
                          "color": {"field": "air"}}}
-    assert _spec_field_references(vega) == {"lon", "time", "air"}
+    assert _spec_field_references(vega, "vega-lite") == {"lon", "time", "air"}
     deck = {"layers": [{"getPosition": "@@=[lon, lat]", "getElevation": "@@=air"}]}
-    assert {"lon", "lat", "air"} <= _spec_field_references(deck)
+    assert {"lon", "lat", "air"} <= _spec_field_references(deck, "deckgl")
 
 
 def test_subset_collapses_dims_absent_from_spec(simple_dataset_3d_pipeline):
     """Dims the spec does not reference (here time on a lon/lat map) collapse to
     a single slice; the spec drives the subset, not a name guess."""
     spec = {"encoding": {"x": {"field": "lon"}, "y": {"field": "lat"}}}
     n_full = len(simple_dataset_3d_pipeline.data)
-    sub = subset_gridded_to_2d(simple_dataset_3d_pipeline, spec)
+    sub = subset_gridded_to_2d(simple_dataset_3d_pipeline, spec, "vega-lite")
     assert len(sub.data) < n_full
     assert sub.data["time"].nunique() == 1
 
 
 def test_subset_keeps_dims_the_spec_uses(simple_dataset_3d_pipeline):
     """A lon/time spec (Hovmoller) keeps time as an axis; lat collapses instead."""
     spec = {"encoding": {"x": {"field": "lon"}, "y": {"field": "time"}}}
-    sub = subset_gridded_to_2d(simple_dataset_3d_pipeline, spec)
+    sub = subset_gridded_to_2d(simple_dataset_3d_pipeline, spec, "vega-lite")
     assert sub.data["time"].nunique() > 1
     assert sub.data["lat"].nunique() == 1
 
 
 def test_subset_noop_when_spec_uses_all_dims(xarray_pipeline):
     """When the spec references every dim, the pipeline is returned unchanged."""
     spec = {"encoding": {"x": {"field": "lon"}, "y": {"field": "lat"}}}
-    assert subset_gridded_to_2d(xarray_pipeline, spec) is xarray_pipeline
+    assert subset_gridded_to_2d(xarray_pipeline, spec, "vega-lite") is xarray_pipeline
 
 
 def test_subset_noop_for_tabular(tabular_pipeline):
     """A non-xarray pipeline is returned unchanged."""
-    assert subset_gridded_to_2d(tabular_pipeline, {}) is tabular_pipeline
+    assert subset_gridded_to_2d(tabular_pipeline, {}, "deckgl") is tabular_pipeline
 
 
 def test_hvplot_prompt_no_gridded_rules_for_tabular():