Generate matplotlib charts with AutoGen code agents

AutoGen’s code-executing agents let you offload matplotlib chart generation to an LLM that writes, runs, and debugs its own Python code. This autogen matplotlib code execution tutorial walks through a production-ready setup: a UserProxyAgent with a Docker sandbox, a coder agent that emits only code blocks, and a feedback loop that catches import errors, style issues, and rendering problems before you ever see a broken PNG.

Step 1: Provision a Docker sandbox for code execution

AutoGen’s UserProxyAgent can execute code locally, but a container isolates dependencies, prevents filesystem accidents, and matches your deployment environment. Create a Dockerfile that pins Python and installs the scientific stack:

# Dockerfile
FROM python:3.11-slim

RUN apt-get update && apt-get install -y --no-install-recommends \
    libfreetype6-dev pkg-config && \
    rm -rf /var/lib/apt/lists/*

WORKDIR /workspace
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Non-root user for safety
RUN useradd -m -u 1000 autogen && chown -R autogen:autogen /workspace
USER autogen

# requirements.txt
autogen-agentchat==0.2.3
autogen-ext==0.2.3
matplotlib==3.8.4
numpy==1.26.4
pandas==2.2.2

Build and tag the image:

docker build -t autogen-matplotlib:latest .

Verify the container starts and can import matplotlib headless:

docker run --rm autogen-matplotlib:latest python -c "import matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt; print('OK')"

You should see OK printed. The Agg backend is critical — it renders to memory without a display server.

Step 2: Define the coder agent with a strict system prompt

The assistant agent must output only executable Python inside fenced code blocks. Any explanatory text breaks the executor. Create agents.py:

# agents.py
from autogen_agentchat.agents import AssistantAgent
from autogen_ext.models.openai import OpenAIChatCompletionClient

model_client = OpenAIChatCompletionClient(model="gpt-4o-mini")

coder = AssistantAgent(
    name="chart_coder",
    model_client=model_client,
    system_message=(
        "You are a matplotlib expert. Write only the Python code needed to "
        "generate the requested chart. Output a single fenced code block "
        "with language tag 'python'. Do not include comments outside the code. "
        "Use matplotlib's 'Agg' backend. Save the figure to '/workspace/chart.png' "
        "with dpi=150 and bbox_inches='tight'. Import pandas as pd and numpy as np "
        "if needed. Assume data is provided as a pandas DataFrame named 'df' "
        "already loaded in the workspace."
    ),
)

The system message constrains the model: one code block, Agg backend, fixed output path, known variable name df. This reduces round-trips.

Step 3: Configure the UserProxyAgent with the Docker executor

UserProxyAgent handles code extraction, container execution, and result capture. Wire it to the image you built:

# executor.py
import tempfile
import shutil
from pathlib import Path
from autogen_agentchat.agents import UserProxyAgent
from autogen_agentchat.code_executors import DockerCommandLineCodeExecutor

# Persistent workspace shared with container
workspace = Path(tempfile.mkdtemp(prefix="autogen_ws_"))
print(f"Workspace: {workspace}")

executor = DockerCommandLineCodeExecutor(
    image="autogen-matplotlib:latest",
    work_dir=workspace,
    bind_dir=workspace,
    timeout=60,
    auto_remove=True,
)

user_proxy = UserProxyAgent(
    name="chart_executor",
    code_executor=executor,
    human_input_mode="NEVER",
    max_consecutive_auto_reply=3,
)

bind_dir mounts the host workspace into the container at the same path, so /workspace/chart.png written inside the container appears on the host. max_consecutive_auto_reply=3 lets the coder retry up to three times after seeing execution errors.

Step 4: Load data into the workspace before the first turn

The coder assumes df exists. Write a small loader script and execute it once during initialization:

# data_prep.py
import pandas as pd
import numpy as np
from pathlib import Path

def prepare_sample_data(workspace: Path) -> Path:
    """Create a sample dataset and return the CSV path."""
    np.random.seed(42)
    n = 200
    df = pd.DataFrame({
        "timestamp": pd.date_range("2024-01-01", periods=n, freq="h"),
        "temperature": 20 + 5 * np.sin(np.linspace(0, 4*np.pi, n)) + np.random.normal(0, 0.5, n),
        "humidity": 60 + 15 * np.cos(np.linspace(0, 2*np.pi, n)) + np.random.normal(0, 2, n),
        "location": np.random.choice(["indoor", "outdoor"], n, p=[0.6, 0.4]),
    })
    csv_path = workspace / "sensor_data.csv"
    df.to_csv(csv_path, index=False)
    return csv_path

Then in your main script, run this once and inject a bootstrap execution that loads df:

# main.py
import asyncio
from pathlib import Path
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.conditions import MaxMessageTermination
from agents import coder
from executor import user_proxy, workspace
from data_prep import prepare_sample_data

async def main():
    csv_path = prepare_sample_data(workspace)
    print(f"Data written to {csv_path}")

    # Bootstrap: load df into the executor's namespace
    bootstrap_code = f"""
import pandas as pd
df = pd.read_csv("{csv_path.name}")
df["timestamp"] = pd.to_datetime(df["timestamp"])
print(f"Loaded {{len(df)}} rows")
print(df.dtypes)
"""
    result = await user_proxy.run_code(bootstrap_code)
    print(result.output)

    # Now start the chat
    team = RoundRobinGroupChat(
        participants=[coder, user_proxy],
        termination_condition=MaxMessageTermination(max_messages=6),
    )

    task = (
        "Create a dual-axis line chart: temperature (left axis, red) and humidity "
        "(right axis, blue) over time. Use the 'timestamp' column for x-axis. "
        "Add a legend, grid, and title 'Environmental Sensor Readings'. "
        "Save to '/workspace/chart.png'."
    )

    async for msg in team.run_stream(task=task):
        print(f"[{msg.source}] {msg.content[:200]}...")

    # Verify output
    chart_path = workspace / "chart.png"
    if chart_path.exists():
        print(f"Success: chart saved to {chart_path} ({chart_path.stat().st_size} bytes)")
    else:
        print("Error: chart.png not found")

    # Cleanup
    shutil.rmtree(workspace, ignore_errors=True)

if __name__ == "__main__":
    asyncio.run(main())

Run it:

python main.py

You should see the bootstrap output showing 200 rows loaded, then the coder emitting a code block, the executor running it, and finally Success: chart saved to ....

Step 5: Inspect the generated code and iterate on failures

The first attempt often works, but when it doesn’t, the error message flows back to the coder. Check the console for the actual code executed. A typical failure looks like:

[chart_executor] Traceback (most recent call last):
  File "<string>", line 12, in <module>
    ax2 = ax1.twinx()
AttributeError: 'AxesSubplot' object has no attribute 'twinx'

The coder sees this traceback and corrects itself in the next turn (hence max_consecutive_auto_reply=3). Common fixes it learns:

• plt.subplots() returns (fig, ax), not ax directly
• ax.twinx() exists on the axes, not the figure
• Date formatting requires matplotlib.dates.DateFormatter

To debug, add a callback that logs every code block before execution:

# debug_executor.py
from autogen_agentchat.code_executors import DockerCommandLineCodeExecutor

class LoggingExecutor(DockerCommandLineCodeExecutor):
    async def execute(self, code: str, **kwargs):
        print("=== CODE TO EXECUTE ===")
        print(code)
        print("========================")
        return await super().execute(code, **kwargs)

# Use LoggingExecutor instead of DockerCommandLineCodeExecutor

This lets you copy the failing block into a local REPL for rapid iteration.

Step 6: Parameterize chart requests with a structured schema

Hard-coded tasks don’t scale. Define a Pydantic model for chart specs and have a planner agent emit JSON that the coder consumes:

# schemas.py
from pydantic import BaseModel, Field
from typing import Literal, Optional

class ChartSpec(BaseModel):
    chart_type: Literal["line", "bar", "scatter", "hist", "dual_axis_line"]
    x: str
    y: list[str] = Field(min_length=1)
    y2: Optional[list[str]] = None
    title: str
    xlabel: Optional[str] = None
    ylabel: Optional[str] = None
    y2label: Optional[str] = None
    color_map: Optional[dict[str, str]] = None
    group_by: Optional[str] = None
    agg: Optional[Literal["mean", "sum", "count"]] = None

Update the coder’s system message to read a spec.json file:

# agents.py (updated coder)
coder = AssistantAgent(
    name="chart_coder",
    model_client=model_client,
    system_message=(
        "You are a matplotlib expert. Read '/workspace/spec.json' for the chart specification. "
        "Write only the Python code to generate the chart. Output a single fenced code block "
        "tagged 'python'. Use 'Agg' backend. Save to '/workspace/chart.png' with dpi=150, "
        "bbox_inches='tight'. DataFrame 'df' is pre-loaded. Spec fields: chart_type, x, y, y2, "
        "title, xlabel, ylabel, y2label, color_map, group_by, agg. Handle missing optional fields gracefully."
    ),
)

Add a planner agent that converts natural language to ChartSpec:

# planner.py
from autogen_agentchat.agents import AssistantAgent
from autogen_ext.models.openai import OpenAIChatCompletionClient
from schemas import ChartSpec
import json

model_client = OpenAIChatCompletionClient(model="gpt-4o-mini")

planner = AssistantAgent(
    name="chart_planner",
    model_client=model_client,
    system_message=(
        "Convert the user's request into a ChartSpec JSON object. "
        "Output ONLY the JSON, no markdown, no explanation. "
        f"Schema: {ChartSpec.model_json_schema()}"
    ),
)

Wire them in a SelectorGroupChat so the planner runs first, writes spec.json, then the coder executes:

# main_structured.py
import asyncio
import json
from pathlib import Path
from autogen_agentchat.teams import SelectorGroupChat
from autogen_agentchat.conditions import MaxMessageTermination
from agents import coder
from planner import planner
from executor import user_proxy, workspace
from data_prep import prepare_sample_data

async def main():
    csv_path = prepare_sample_data(workspace)

    # Bootstrap df
    bootstrap = f"""
import pandas as pd
df = pd.read_csv("{csv_path.name}")
df["timestamp"] = pd.to_datetime(df["timestamp"])
"""
    await user_proxy.run_code(bootstrap)

    team = SelectorGroupChat(
        participants=[planner, coder, user_proxy],
        model_client=planner.model_client,
        termination_condition=MaxMessageTermination(max_messages=8),
        selector_prompt=(
            "Choose the next agent: planner (first turn), coder (after spec written), "
            "executor (after code written)."
        ),
    )

    task = (
        "Dual-axis line chart of temperature and humidity over time, grouped by location, "
        "showing mean per hour. Title: 'Hourly Environmental Averages by Location'."
    )

    async for msg in team.run_stream(task=task):
        print(f"[{msg.source}] {msg.content[:300]}")

    chart_path = workspace / "chart.png"
    if chart_path.exists():
        print(f"Success: {chart_path}")
    else:
        print("Failed")

    shutil.rmtree(workspace, ignore_errors=True)

if __name__ == "__main__":
    asyncio.run(main())

The planner emits JSON like:

{
  "chart_type": "dual_axis_line",
  "x": "timestamp",
  "y": ["temperature"],
  "y2": ["humidity"],
  "title": "Hourly Environmental Averages by Location",
  "group_by": "location",
  "agg": "mean"
}

The coder reads this, writes the appropriate groupby().resample().agg() logic, and produces a faceted or multi-line chart without further prompt engineering.

Step 7: Add caching and cache-control for repeated renders

If you’re generating the same chart spec repeatedly — say, a dashboard that refreshes every minute — re-running the LLM and container is wasteful. Hash the spec and cache the PNG:

# cache.py
import hashlib
import json
from pathlib import Path

CACHE_DIR = Path("/tmp/autogen_chart_cache")
CACHE_DIR.mkdir(exist_ok=True)

def cache_key(spec: dict) -> str:
    return hashlib.sha256(json.dumps(spec, sort_keys=True).encode()).hexdigest()[:16]

def get_cached(spec: dict) -> Path | None:
    key = cache_key(spec)
    cached = CACHE_DIR / f"{key}.png"
    return cached if cached.exists() else None

def store_cache(spec: dict, src: Path) -> Path:
    key = cache_key(spec)
    dst = CACHE_DIR / f"{key}.png"
    shutil.copy2(src, dst)
    return dst

Wrap the team run:

# main_cached.py
async def run_with_cache(spec: dict, task: str):
    cached = get_cached(spec)
    if cached:
        print(f"Cache hit: {cached}")
        return cached

    # ... run team as before ...
    chart_path = workspace / "chart.png"
    if chart_path.exists():
        return store_cache(spec, chart_path)
    return None

When you deploy behind a gateway like n4n.ai, you can forward the Cache-Control header from the provider response so downstream clients respect the same TTL — useful if the chart data changes on a known schedule.

Step 8: Verify success in CI and production

Add a pytest that runs the full pipeline headless:

# test_chart_generation.py
import pytest
from main_structured import main
from pathlib import Path

@pytest.mark.asyncio
async def test_dual_axis_chart_generates_png(tmp_path):
    # Monkey-patch workspace to tmp_path
    import main_structured as mod
    original_workspace = mod.workspace
    mod.workspace = tmp_path
    try:
        await mod.main()
        chart = tmp_path / "chart.png"
        assert chart.exists()
        assert chart.stat().st_size > 1000  # non-trivial PNG
    finally:
        mod.workspace = original_workspace

Run in CI with docker build -t autogen-test . && docker run --rm autogen-test pytest -xvs. The containerized executor means the test environment matches production exactly.

For production observability, emit structured logs from the executor:

# observability.py
import structlog
import time

logger = structlog.get_logger()

async def execute_with_metrics(executor, code: str):
    start = time.perf_counter()
    result = await executor.execute(code)
    duration = time.perf_counter() - start
    logger.info(
        "code_execution",
        success=result.exit_code == 0,
        duration_ms=int(duration * 1000),
        output_lines=len(result.output.splitlines()),
        error=result.output if result.exit_code != 0 else None,
    )
    return result

This gives you latency percentiles, error rates, and a trail to correlate with provider-side metrics if you’re routing through an inference gateway.

Step 9: Extend to multi-chart reports

A single PNG is rarely the end product. Compose multiple charts into a PDF report using matplotlib.backends.backend_pdf:

# report_generator.py
from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.pyplot as plt

def build_report(chart_specs: list[dict], df: pd.DataFrame, output: Path):
    with PdfPages(output) as pdf:
        for spec in chart_specs:
            fig, ax = plt.subplots(figsize=(10, 5))
            # ... render logic per spec ...
            pdf.savefig(fig, bbox_inches="tight")
            plt.close(fig)

The planner can now emit a ReportSpec containing a list of ChartSpec objects. The coder writes a loop that iterates the list and calls pdf.savefig(). Same agent topology, richer output.

Verification checklist

• docker run --rm autogen-matplotlib:latest python -c "import matplotlib; matplotlib.use('Agg'); print('OK')" prints OK
• python main.py produces chart.png > 1 KB in the workspace directory
• pytest test_chart_generation.py -xvs passes
• Changing the task string yields a different valid chart without code changes
• Cache hit returns the same PNG path without invoking the LLM or container

You now have a reproducible, containerized, cacheable pipeline where an LLM writes matplotlib code, executes it in isolation, self-corrects on errors, and hands you a verified artifact. The same pattern scales to dashboards, scheduled reports, and user-facing chart APIs.