Fix pie/donut chart label overlap with post-processing

lumen/ai/agents/vega_lite.py

@@ -242,6 +242,102 @@ def _deep_merge_dicts(self, base_dict: dict[str, Any], update_dict: dict[str, An
 
         return result
 
+    _MAX_ARC_LABELS = 8
+
+    def _fix_arc_labels(self, spec: dict[str, Any]) -> dict[str, Any]:
+        """Fix pie/donut chart labels by ensuring proper outside placement.
+
+        When the LLM generates a pie chart with text labels, the labels
+        often overlap on small slices. This method replaces any existing
+        text layer with one that uses theta encoding (same polar coordinate
+        system as the arc) with a large radius to push labels outside.
+
+        For charts with more than ``_MAX_ARC_LABELS`` categories, text
+        labels are removed entirely and the chart relies on the color
+        legend and tooltips instead.
+        """
+        layers = spec.get("layer", [])
+        if not layers:
+            return spec
+
+        arc_layer = None
+        for layer in layers:
+            if self._get_layer_mark_type(layer) == "arc":
+                arc_layer = layer
+                break
+        if arc_layer is None:
+            return spec
+
+        # Find theta and color encodings from arc layer or top-level
+        theta_enc = (
+            arc_layer.get("encoding", {}).get("theta")
+            or spec.get("encoding", {}).get("theta")
+        )
+        color_enc = (
+            arc_layer.get("encoding", {}).get("color")
+            or spec.get("encoding", {}).get("color")
+        )
+        if not theta_enc or not color_enc:
+            return spec
+
+        theta_field = theta_enc.get("field", "")
+        color_field = color_enc.get("field", "")
+        if not theta_field or not color_field:
+            return spec
+
+        # Remove existing text layers
+        spec["layer"] = [l for l in layers if self._get_layer_mark_type(l) != "text"]
+
+        # Normalize arc mark and set outerRadius
+        arc_mark = arc_layer.get("mark", {})
+        if isinstance(arc_mark, str):
+            arc_layer["mark"] = {"type": "arc", "outerRadius": 80}
+        elif isinstance(arc_mark, dict):
+            arc_mark.setdefault("outerRadius", 80)
+
+        outer_r = arc_layer["mark"].get("outerRadius", 80) if isinstance(arc_layer["mark"], dict) else 80
+
+        # Count categories from inline data to decide whether labels fit
+        n_categories = self._count_arc_categories(spec, color_field)
+        if n_categories <= self._MAX_ARC_LABELS:
+            text_layer = {
+                "mark": {
+                    "type": "text",
+                    "radius": outer_r + 30,
+                    "fontSize": 11,
+                },
+                "encoding": {
+                    "theta": dict(theta_enc, stack=True),
+                    "text": {"field": color_field, "type": "nominal"},
+                    "color": {"value": "black"},
+                },
+            }
+            spec["layer"].append(text_layer)
+
+        # Ensure legend is enabled on color encoding
+        for layer in spec["layer"]:
+            enc = layer.get("encoding", {})
+            color = enc.get("color", {})
+            if isinstance(color, dict) and color.get("field"):
+                color.pop("legend", None)
+        top_color = spec.get("encoding", {}).get("color", {})
+        if isinstance(top_color, dict) and top_color.get("field"):
+            top_color.pop("legend", None)
+
+        return spec
+
+    def _count_arc_categories(self, spec: dict, color_field: str) -> int:
+        """Count unique category values from inline data.
+
+        Returns 0 when data is referenced by name (not inline) so that
+        the caller falls back to adding labels by default.
+        """
+        data = spec.get("data", {})
+        values = data.get("values", [])
+        if not values:
+            return 0
+        return len({row.get(color_field) for row in values if color_field in row})
+
     def _get_layer_mark_type(self, layer: dict) -> str | None:
         """Extract mark type from a layer, handling both dict and string formats."""
         if "mark" in layer:
@@ -565,6 +661,7 @@ async def _extract_spec(self, context: TContext, spec: dict[str, Any]):
                 vega_spec["height"] = "container"
 
         self._editor_type.validate_spec(vega_spec)
+        vega_spec = self._fix_arc_labels(vega_spec)
 
         # using string comparison because these keys could be in different nested levels
         vega_spec_str = dump_yaml(vega_spec)

lumen/ai/prompts/VegaLiteAgent/main.jinja2

@@ -55,6 +55,8 @@ Legends appear when you map data to `color`, `size`, or `shape` in encoding.
 
 **Selective emphasis**: Reduce opacity of secondary data (`opacity: 0.3`), keep primary prominent (`opacity: 1, strokeWidth: 2-3`)
 
+**Pie/donut charts**: Use `arc` mark with `theta` encoding and a color legend. Tooltips provide detail on hover. For donuts, add `innerRadius` to the arc mark.
+
 **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!)
@@ -203,6 +205,20 @@ layer:
       color: {field: category, type: nominal, legend: {title: "Category", orient: "right"}}
 ```
 
+Pie/donut chart (ALWAYS use this pattern for pie charts):
+```yaml
+data:
+  name: <TABLE_NAME>
+layer:
+  - mark: {type: arc, outerRadius: 80}
+    encoding:
+      theta: {field: amount, type: quantitative, stack: true}
+      color: {field: category, type: nominal, legend: {title: "Category"}}
+      tooltip:
+        - {field: category, type: nominal, title: "Category"}
+        - {field: amount, type: quantitative, format: ",", title: "Amount"}
+```
+
 Choropleth map:
 ```yaml
 data:

lumen/tests/ai/test_agents.py

@@ -274,3 +274,130 @@ class ExtendedChatAgent(ChatAgent):
         messages = [{"role": "user", "content": "test"}]
         prompt = await agent._render_prompt("main", messages, {})
         assert "Footer appended." in prompt
+
+
+class TestFixArcLabels:
+    """Tests for VegaLiteAgent._fix_arc_labels post-processing."""
+
+    @pytest.fixture
+    def agent(self, llm):
+        return VegaLiteAgent(llm=llm, code_execution="disabled")
+
+    def _make_pie_spec(self, categories=None, outer_radius=None, with_text_layer=False):
+        """Helper to build a pie chart spec for testing."""
+        if categories is None:
+            categories = [
+                {"cat": "A", "val": 10},
+                {"cat": "B", "val": 20},
+                {"cat": "C", "val": 30},
+            ]
+        arc_mark = {"type": "arc"}
+        if outer_radius is not None:
+            arc_mark["outerRadius"] = outer_radius
+        layers = [
+            {
+                "mark": arc_mark,
+                "encoding": {
+                    "theta": {"field": "val", "type": "quantitative", "stack": True},
+                    "color": {"field": "cat", "type": "nominal"},
+                },
+            }
+        ]
+        if with_text_layer:
+            layers.append({
+                "mark": {"type": "text", "radiusOffset": 10},
+                "encoding": {
+                    "theta": {"field": "val", "type": "quantitative", "stack": True},
+                    "text": {"field": "cat", "type": "nominal"},
+                },
+            })
+        return {"data": {"values": categories}, "layer": layers}
+
+    def test_adds_text_layer_with_correct_radius(self, agent):
+        spec = self._make_pie_spec(outer_radius=80)
+        result = agent._fix_arc_labels(spec)
+        text_layers = [l for l in result["layer"] if agent._get_layer_mark_type(l) == "text"]
+        assert len(text_layers) == 1
+        assert text_layers[0]["mark"]["radius"] == 110  # 80 + 30
+
+    def test_replaces_existing_text_layer(self, agent):
+        spec = self._make_pie_spec(with_text_layer=True)
+        result = agent._fix_arc_labels(spec)
+        text_layers = [l for l in result["layer"] if agent._get_layer_mark_type(l) == "text"]
+        assert len(text_layers) == 1
+        # Should not have the old radiusOffset
+        assert "radiusOffset" not in text_layers[0]["mark"]
+
+    def test_preserves_non_arc_specs(self, agent):
+        bar_spec = {"layer": [{"mark": "bar", "encoding": {"x": {"field": "a"}}}]}
+        result = agent._fix_arc_labels(bar_spec)
+        assert result == bar_spec
+
+    def test_no_layers_returns_unchanged(self, agent):
+        spec = {"mark": "arc", "encoding": {"theta": {"field": "v"}}}
+        result = agent._fix_arc_labels(spec)
+        assert result is spec
+
+    def test_normalizes_string_arc_mark(self, agent):
+        spec = {
+            "data": {"values": [{"c": "X", "v": 1}]},
+            "layer": [
+                {
+                    "mark": "arc",
+                    "encoding": {
+                        "theta": {"field": "v", "type": "quantitative"},
+                        "color": {"field": "c", "type": "nominal"},
+                    },
+                }
+            ],
+        }
+        result = agent._fix_arc_labels(spec)
+        arc = [l for l in result["layer"] if agent._get_layer_mark_type(l) == "arc"][0]
+        assert isinstance(arc["mark"], dict)
+        assert arc["mark"]["outerRadius"] == 80
+
+    def test_skips_labels_for_many_categories(self, agent):
+        categories = [{"cat": chr(65 + i), "val": 10} for i in range(10)]
+        spec = self._make_pie_spec(categories=categories)
+        result = agent._fix_arc_labels(spec)
+        text_layers = [l for l in result["layer"] if agent._get_layer_mark_type(l) == "text"]
+        assert len(text_layers) == 0
+
+    def test_adds_labels_for_few_categories(self, agent):
+        categories = [{"cat": chr(65 + i), "val": 10} for i in range(5)]
+        spec = self._make_pie_spec(categories=categories)
+        result = agent._fix_arc_labels(spec)
+        text_layers = [l for l in result["layer"] if agent._get_layer_mark_type(l) == "text"]
+        assert len(text_layers) == 1
+
+    def test_adds_labels_when_data_is_named(self, agent):
+        """When data is referenced by name (not inline), labels are added."""
+        spec = {
+            "data": {"name": "my_table"},
+            "layer": [
+                {
+                    "mark": {"type": "arc", "outerRadius": 80},
+                    "encoding": {
+                        "theta": {"field": "val", "type": "quantitative"},
+                        "color": {"field": "cat", "type": "nominal"},
+                    },
+                }
+            ],
+        }
+        result = agent._fix_arc_labels(spec)
+        text_layers = [l for l in result["layer"] if agent._get_layer_mark_type(l) == "text"]
+        assert len(text_layers) == 1
+
+    def test_removes_legend_none(self, agent):
+        spec = self._make_pie_spec()
+        spec["layer"][0]["encoding"]["color"]["legend"] = None
+        result = agent._fix_arc_labels(spec)
+        arc = [l for l in result["layer"] if agent._get_layer_mark_type(l) == "arc"][0]
+        assert "legend" not in arc["encoding"]["color"]
+
+    def test_donut_preserves_inner_radius(self, agent):
+        spec = self._make_pie_spec(outer_radius=80)
+        spec["layer"][0]["mark"]["innerRadius"] = 40
+        result = agent._fix_arc_labels(spec)
+        arc = [l for l in result["layer"] if agent._get_layer_mark_type(l) == "arc"][0]
+        assert arc["mark"]["innerRadius"] == 40