Add gridded heatmap guidance to VegaLiteAgent prompt

Author: ghostiee-11

lumen/ai/prompts/VegaLiteAgent/main.jinja2

@@ -58,6 +58,14 @@ Legends appear when you map data to `color`, `size`, or `shape` in encoding.
 **Geographic maps**: Use `longitude`/`latitude` encodings (NOT `x`/`y`) with `projection: {type: mercator}` at top level. Base map added automatically.
 
 **Choropleth maps**: Join data to map boundaries - `lookup: <map_field>`, `from: {data: {...}, key: <your_field>, fields: [...]}` (key/fields inside from!)
+{%- if gridded is defined and gridded %}
+
+**Gridded xarray data**: The data is a long-form grid from an xarray dataset.
+- dimensions: {{ gridded.dims | join(', ') }}
+- data variables: {{ gridded.data_vars | join(', ') }}
+Use `mark: rect` with x/y mapped to coordinate dimensions and color mapped to the value column.
+Do not use `longitude`/`latitude` encodings or geographic projections for this data.
+{%- endif %}
 {% endblock %}
 
 {% block examples %}
@@ -203,6 +211,26 @@ layer:
       color: {field: category, type: nominal, legend: {title: "Category", orient: "right"}}
 ```
 
+{%- if gridded is defined and gridded %}
+Gridded heatmap (xarray-derived long-form data):
+```yaml
+data:
+  name: <TABLE_NAME>
+layer:
+  - mark: rect
+    encoding:
+      x: {field: {{ gridded.dims[1] if gridded.dims | length > 1 else gridded.dims[0] }}, type: ordinal}
+      y: {field: {{ gridded.dims[0] }}, type: ordinal}
+      color:
+        field: {{ gridded.data_vars[0] }}
+        type: quantitative
+        scale: {scheme: viridis}
+        legend: {title: "{{ gridded.data_vars[0] }}"}
+config:
+  view: {stroke: null}
+```
+{%- endif %}
+
 Choropleth map:
 ```yaml
 data:

lumen/tests/ai/test_vega_lite_gridded.py

@@ -0,0 +1,55 @@
+import datetime
+
+from lumen.ai.config import PROMPTS_DIR
+from lumen.ai.utils import render_template
+
+
+def _base_context(gridded=None):
+    class _FakeSource:
+        dialect = None
+
+    class _FakePipeline:
+        table = "air"
+        source = _FakeSource()
+
+    class _FakeMemory(dict):
+        pass
+
+    memory = _FakeMemory({
+        "pipeline": _FakePipeline(),
+        "data": "lat,lon,air\n0,10,1.0",
+        "table": "air",
+    })
+    return dict(
+        memory=memory,
+        doc="VegaLiteView description",
+        gridded=gridded,
+        current_datetime=datetime.datetime(2026, 1, 1),
+        actor_name="VegaLiteAgent",
+    )
+
+
+def test_vegalite_prompt_includes_gridded_block():
+    gridded = {
+        "source_type": "xarray",
+        "dims": ["lat", "lon"],
+        "coords": {"lat": [3], "lon": [4]},
+        "data_vars": ["air"],
+    }
+    rendered = render_template(
+        PROMPTS_DIR / "VegaLiteAgent" / "main.jinja2",
+        **_base_context(gridded=gridded),
+    )
+    assert "mark: rect" in rendered
+    assert "viridis" in rendered
+    assert "lat" in rendered
+    assert "air" in rendered
+
+
+def test_vegalite_prompt_no_gridded_block_for_tabular():
+    rendered = render_template(
+        PROMPTS_DIR / "VegaLiteAgent" / "main.jinja2",
+        **_base_context(gridded=None),
+    )
+    assert "Gridded xarray data" not in rendered
+    assert "Gridded heatmap" not in rendered