From bf493c60b594a5cf60e8d6a50d3a7b4b25bb6d44 Mon Sep 17 00:00:00 2001 From: ghostiee-11 <168410465+ghostiee-11@users.noreply.github.com> Date: Wed, 18 Mar 2026 12:49:14 +0530 Subject: [PATCH] Fix pie/donut chart label overlap with post-processing Add _fix_arc_labels post-processing step that repositions text labels outside the pie using theta encoding with a radius offset. For charts with >8 categories, labels are dropped in favor of legend + tooltip. --- lumen/ai/agents/vega_lite.py | 97 ++++++++++++++++ lumen/ai/prompts/VegaLiteAgent/main.jinja2 | 16 +++ lumen/tests/ai/test_agents.py | 127 +++++++++++++++++++++ 3 files changed, 240 insertions(+) diff --git a/lumen/ai/agents/vega_lite.py b/lumen/ai/agents/vega_lite.py index d9ac23169..eb9b3222b 100644 --- a/lumen/ai/agents/vega_lite.py +++ b/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) diff --git a/lumen/ai/prompts/VegaLiteAgent/main.jinja2 b/lumen/ai/prompts/VegaLiteAgent/main.jinja2 index b0e368f30..319459a50 100644 --- a/lumen/ai/prompts/VegaLiteAgent/main.jinja2 +++ b/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: `, `from: {data: {...}, key: , 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: +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: diff --git a/lumen/tests/ai/test_agents.py b/lumen/tests/ai/test_agents.py index 928ed7696..5ca5658bd 100644 --- a/lumen/tests/ai/test_agents.py +++ b/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