Build a data analysis agent with AutoGen code execution

Building an autogen data analysis agent code execution loop is the fastest way to turn a raw CSV into actionable insights without hand-writing the pandas boilerplate. This guide walks through a hardened setup where the agent writes Python, runs it in a sandbox, and iterates on errors until the analysis is correct.

Step 1: Install AutoGen and create a sample dataset

Install the framework and its dependencies. Use a virtual environment to avoid polluting your system Python.

python -m venv .venv
source .venv/bin/activate
pip install pyautogen pandas matplotlib

Create a small CSV so the agent has something to work with. The snippet below writes sales.csv with three columns: date, region, and revenue.

import csv
from datetime import date, timedelta

rows = []
base = date(2024, 1, 1)
for i in range(90):
    d = base + timedelta(days=i)
    region = "east" if i % 2 else "west"
    revenue = 1000 + (i * 37) % 500
    rows.append((d.isoformat(), region, revenue))

with open("sales.csv", "w", newline="") as f:
    w = csv.writer(f)
    w.writerow(["date", "region", "revenue"])
    w.writerows(rows)

Run the script. You now have sales.csv in your working directory.

Step 2: Configure the LLM client

AutoGen expects an OpenAI-style chat completion endpoint. You can use OpenAI directly, but for a single endpoint that fronts 240+ models with automatic fallback when a provider is rate-limited or degraded, point AutoGen’s base_url at n4n.ai’s OpenAI-compatible endpoint and supply your key.

Create config.json:

[
  {
    "model": "gpt-4o-mini",
    "api_key": "YOUR_KEY",
    "base_url": "https://api.n4n.ai/v1",
    "api_type": "openai"
  }
]

Load it in Python with config_list_from_json. If you prefer environment variables, set OPENAI_API_BASE and reference config_list_from_env.

from autogen import config_list_from_json

config_list = config_list_from_json("config.json")

The model field accepts any model name the endpoint supports. Swap to claude-3-5-sonnet or a local model without changing agent code. The gateway forwards provider cache-control hints, so repeated analysis of the same dataset benefits from prompt caching if the upstream provider supports it.

Step 3: Define the code-executing user proxy

The UserProxyAgent is the component that actually runs generated code. Set human_input_mode="NEVER" to make it autonomous, and configure code_execution_config to control where code runs and what it can do.

from autogen import UserProxyAgent

code_exec_cfg = {
    "work_dir": "coding",
    "use_docker": False,
    "timeout": 60,
    "allowed_imports": ["pandas", "matplotlib", "numpy"],
}

user_proxy = UserProxyAgent(
    name="executor",
    human_input_mode="NEVER",
    code_execution_config=code_exec_cfg,
    max_consecutive_auto_reply=10,
)

use_docker=False runs code on your local machine. For anything beyond a trusted laptop, set use_docker=True and pre-build the autogen image. allowed_imports blocks arbitrary import os or subprocess calls, which is a minimum bar for safe autogen data analysis agent code execution.

How the executor processes a reply

When the assistant emits a code block, the proxy extracts it, writes it to coding/agentchat_xxx.py, and runs it with python. Stdout and stderr are captured. If the exit code is non-zero, the text is sent back to the assistant as a new message, triggering a self-healing retry. This loop continues until the script passes or max_consecutive_auto_reply is exhausted.

Step 4: Define the analysis assistant

The assistant writes the Python. Give it a tight system message so it returns runnable scripts rather than prose.

from autogen import AssistantAgent

assistant = AssistantAgent(
    name="analyst",
    llm_config={"config_list": config_list, "cache_seed": 42},
    system_message=(
        "You are a data analyst. Write concise Python using pandas to answer "
        "the user's question. Save plots to files in the working directory. "
        "Do not explain the code; output only a single code block."
    ),
)

The cache_seed enables deterministic responses during development. Remove it in production to let the model vary its approach. Keep the system message strict: any markdown outside the code fence forces the proxy to ignore it, but a loose instruction will cause the model to chat instead of compute.

Step 5: Run the first analysis task

Initiate the chat from the user proxy. The task should be specific enough that the agent knows which file to load.

task = (
    "Load sales.csv. Compute total revenue per region. "
    "Then plot monthly revenue as a line chart and save it to revenue.png."
)

user_proxy.initiate_chat(assistant, message=task)

AutoGen will stream the assistant’s proposed code to the proxy, which writes it to coding/xxx.py and executes it. If the script throws, the traceback is fed back to the assistant for a fix. This retry loop is the core of autogen data analysis agent code execution.

Example of a generated script

A successful first pass typically looks like this:

import pandas as pd
import matplotlib.pyplot as plt

df = pd.read_csv("sales.csv")
df["date"] = pd.to_datetime(df["date"])
df["month"] = df["date"].dt.to_period("M").astype(str)

monthly = df.groupby("month")["revenue"].sum()
monthly.plot(kind="line", title="Monthly Revenue")
plt.savefig("revenue.png")

If the agent forgot to convert date to datetime, the proxy would return a TypeError and the assistant would revise.

Step 6: Verify success

Success has three observable signals:

1. The process exits without an unhandled exception.
2. coding/ contains at least one .py file and a corresponding .txt stdout log.
3. revenue.png exists in the working directory and is non-zero size.

Check from a shell:

ls -la coding/
file revenue.png

Open revenue.png. You should see two lines (east/west) or aggregated monthly totals. If the image is missing, inspect coding/last_run.txt for the error the agent could not resolve.

Reading the stdout log

The log captures print output from the script. For the task above, you might see:

east    48000
west    47500

That confirms the groupby executed. If you see a ModuleNotFoundError, your allowed_imports list is too strict.

Step 7: Harden the execution environment

Local execution is fine for a demo, but a real autogen data analysis agent code execution deployment needs constraints.

• Docker isolation: set "use_docker": True. AutoGen will mount work_dir and run code inside autogen/code-executor.
• Timeouts: keep timeout at 30–60s to prevent runaway loops.
• Import allowlist: expand allowed_imports only when a new library is required.
• Resource limits: if using Docker, add --memory and --cpus flags via a custom container recipe.

Example hardened config:

code_exec_cfg = {
    "work_dir": "coding",
    "use_docker": True,
    "timeout": 30,
    "allowed_imports": ["pandas", "matplotlib", "numpy", "sklearn"],
}

Docker prerequisites

Build the image once:

docker build -t autogen-exec -f autogen/code-executor/Dockerfile .

Then set use_docker="autogen-exec" (string form) to force that specific image. This avoids pulling at runtime.

Step 8: Iterate with follow-up queries

The same user_proxy keeps conversation state. Ask a refinement without restarting:

user_proxy.send(
    "Now show a bar chart of average revenue by region, sorted descending.",
    request_reply=True,
)

Because max_consecutive_auto_reply is set, the agent will again write, run, and self-correct. You can chain ten such requests; each runs as a fresh script in the same work directory. The autogen data analysis agent code execution pattern handles schema drift—if you add a product column tomorrow, just ask a new question.

Step 9: Capture per-token cost and route explicitly

When you use a gateway endpoint, you can pass provider routing hints via the extra_headers field in llm_config. This works transparently with the OpenAI client AutoGen uses.

llm_config = {
    "config_list": config_list,
    "extra_headers": {"x-provider-preference": "openai,anthropic"},
}

If you are on n4n.ai, per-token usage metering is reported back on each response, so you can log response.usage to track spend without building your own middleware. That is the only plumbing most teams need to productionize the loop.

What a production loop looks like

A minimal service wraps steps 2–8 in a function:

def analyze(csv_path: str, question: str) -> dict:
    proxy = UserProxyAgent(...)  # as above
    agent = AssistantAgent(...)  # as above
    proxy.initiate_chat(agent, message=f"File: {csv_path}. Task: {question}")
    return {"artifacts": os.listdir("coding")}

Run it behind a queue. The agent’s ability to execute code removes the need for a separate query layer—your SQL or pandas logic is generated on demand.

The autogen data analysis agent code execution pattern shines when schemas change daily. You never rewrite the parser; you just ask a new question.