What Is a Data Analysis Agent in AutoGen
A Data Analysis Agent built with AutoGen is an AI-powered conversational assistant that can autonomously load, clean, transform, visualize, and interpret datasets. Instead of writing manual scripts for every analysis task, you define an agent with specific capabilities — such as executing Python code, reading files, and generating charts — and then interact with it using natural language. The agent handles the technical execution while you focus on asking questions and making decisions.
AutoGen, developed by Microsoft Research, is an open-source framework for building multi-agent conversational systems. It allows you to create specialized agents that can talk to each other, use tools, and execute code in sandboxed environments. A Data Analysis Agent is one of the most practical applications of this framework — it combines a language model's reasoning with actual code execution to produce reliable, reproducible analytical results.
Why a Data Analysis Agent Matters
Traditional data analysis workflows require constant context-switching between querying data, writing code, interpreting outputs, and documenting findings. This fragmentation slows down exploration and increases the chance of errors. A Data Analysis Agent addresses these pain points in several key ways:
- Natural language interface: You describe what you want in plain English, and the agent translates it into executable code.
- Code execution and verification: The agent doesn't just suggest code — it runs it, shows the output, and self-corrects when errors occur.
- Iterative refinement: You can ask follow-up questions, request different visualizations, or drill deeper without starting over.
- Reproducibility: All code and results are captured in the conversation history, creating an audit trail for every analytical step.
- Accessibility: Team members without deep programming skills can perform sophisticated analyses by simply describing their intent.
Setting Up Your Environment
Before building the agent, install the required packages. AutoGen works with Python 3.8+ and requires an API key from a supported LLM provider (OpenAI, Azure OpenAI, or compatible endpoints).
# Create and activate a virtual environment (recommended)
python -m venv autogen_env
source autogen_env/bin/activate # On Windows: autogen_env\Scripts\activate
# Install AutoGen and data analysis dependencies
pip install pyautogen pandas matplotlib seaborn numpy openai
Set your API key as an environment variable to keep it secure:
# Linux / macOS
export OPENAI_API_KEY="your-api-key-here"
# Windows PowerShell
$env:OPENAI_API_KEY="your-api-key-here"
# Or configure in a .env file with python-dotenv
# pip install python-dotenv
# Then in your script: from dotenv import load_dotenv; load_dotenv()
Creating Your First Data Analysis Agent
The core pattern in AutoGen is defining an AssistantAgent that uses a language model to generate responses, paired with a UserProxyAgent that executes code on behalf of the assistant. For data analysis, you configure the assistant to write Python code and the user proxy to run it in a local execution environment.
Basic Two-Agent Setup
import autogen
import pandas as pd
import os
# Configure the LLM
llm_config = {
"config_list": [
{
"model": "gpt-4",
"api_key": os.environ.get("OPENAI_API_KEY"),
}
],
"temperature": 0.1, # Lower temperature for more deterministic code generation
"timeout": 120,
}
# Create the assistant agent that will write analysis code
data_analyst = autogen.AssistantAgent(
name="DataAnalyst",
system_message="""You are an expert data analyst. Your role is to:
1. Write Python code to analyze datasets using pandas, matplotlib, seaborn, and numpy.
2. Load data from CSV, JSON, or Excel files when provided with a file path.
3. Perform data cleaning, transformation, aggregation, and statistical analysis.
4. Create clear, well-labeled visualizations.
5. Explain your findings in plain language alongside the code.
6. If an error occurs, read the error message, fix the code, and try again.
7. Save generated charts to disk with descriptive filenames using plt.savefig().
Always wrap your code in python ... blocks so the executor can run it.
Never mention the code block markers in your explanation text — just provide the code block and then explain results.""",
llm_config=llm_config,
)
# Create the user proxy agent that executes code
user_proxy = autogen.UserProxyAgent(
name="UserProxy",
human_input_mode="NEVER", # Fully automated execution
max_consecutive_auto_reply=10,
code_execution_config={
"work_dir": "analysis_output",
"use_docker": False, # Set to True if you want Docker sandboxing
"timeout": 300,
},
system_message="Executor. Run the Python code provided by the DataAnalyst and report results.",
)
Understanding the Configuration
The system_message in the assistant agent is critical — it defines the agent's personality, capabilities, and output format. For data analysis, you want the agent to always provide executable code blocks. The human_input_mode="NEVER" setting on the user proxy means the agent runs fully autonomously without waiting for human approval at each step. The max_consecutive_auto_reply prevents infinite loops by limiting how many times the agents can talk to each other without human intervention.
The code_execution_config dictionary specifies where code runs. Setting work_dir creates a dedicated folder for generated files. The use_docker flag determines whether code executes in a Docker container (safer but requires Docker installation) or directly on your machine (convenient for local development with trusted data).
Running Your First Analysis
With the agents defined, you initiate a conversation by sending a message from the user proxy to the data analyst. Here's how to load a CSV file and ask for an initial exploration:
# Create sample data for demonstration
sample_data = pd.DataFrame({
"date": pd.date_range("2023-01-01", periods=200, freq="D"),
"sales": np.random.normal(5000, 800, 200) + np.linspace(0, 3000, 200),
"region": np.random.choice(["North", "South", "East", "West"], 200),
"product_category": np.random.choice(["Electronics", "Clothing", "Food"], 200),
"units_sold": np.random.randint(10, 100, 200),
"customer_satisfaction": np.random.uniform(3.0, 5.0, 200).round(2),
})
sample_data.to_csv("sales_data.csv", index=False)
# Initiate the analysis conversation
user_proxy.initiate_chat(
data_analyst,
message="""I have a sales dataset in 'sales_data.csv'. Please perform the following analysis:
1. Load and inspect the data — show the shape, column types, and first 5 rows.
2. Check for missing values and report findings.
3. Compute basic summary statistics for numeric columns.
4. Create a line chart showing daily sales trend over time. Save it as 'daily_sales_trend.png'.
5. Create a bar chart comparing average sales by region. Save it as 'sales_by_region.png'.
6. Identify the top 3 performing product categories by total units sold.
After completing all tasks, provide a brief executive summary of the key insights.""",
clear_history=True,
)
Understanding What Happens Under the Hood
When you call initiate_chat, AutoGen orchestrates a multi-turn conversation:
- Turn 1: The user proxy sends your message to the DataAnalyst agent.
- Turn 2: The DataAnalyst receives the message, consults its system prompt, and generates a response containing Python code wrapped in markdown code blocks. It may also include explanatory text.
- Turn 3: The user proxy detects the code block, extracts it, and executes it in the configured environment. Any output (printed values, chart files saved, errors) is captured.
- Turn 4 onward: If execution succeeded, the user proxy feeds the output back to the DataAnalyst, which then interprets results and may generate additional code. If an error occurred, the DataAnalyst reads the traceback and attempts to fix the code — this self-correction loop continues until success or the
max_consecutive_auto_replylimit is reached.
This loop is what makes the agent powerful: it doesn't just generate code and stop — it verifies execution, learns from errors, and iterates toward a correct solution.
Building a More Advanced Multi-Agent Team
For complex data projects, a single assistant agent can become overloaded. AutoGen supports multi-agent teams where specialized agents collaborate. Here's a production-grade setup with three specialized agents and a group chat manager:
import autogen
import os
import json
llm_config = {
"config_list": [
{
"model": "gpt-4",
"api_key": os.environ.get("OPENAI_API_KEY"),
}
],
"temperature": 0.1,
"timeout": 120,
}
# Agent 1: Data Engineer — handles loading, cleaning, transformation
data_engineer = autogen.AssistantAgent(
name="DataEngineer",
system_message="""You are a data engineer. Your responsibilities:
- Load data from files (CSV, JSON, Excel, Parquet) or databases.
- Handle missing values, outliers, and data type conversions.
- Perform merges, joins, pivots, and feature engineering.
- Output clean, analysis-ready DataFrames and save them as CSV when appropriate.
- Document all cleaning steps and transformations performed.
Always provide executable Python code in python ... blocks.""",
llm_config=llm_config,
)
# Agent 2: Data Analyst — handles statistics and visualization
data_analyst = autogen.AssistantAgent(
name="DataAnalyst",
system_message="""You are a data analyst. Your responsibilities:
- Perform statistical analysis: distributions, correlations, hypothesis tests.
- Create professional visualizations with matplotlib and seaborn.
- Generate clear insights from data patterns.
- Save all charts with descriptive filenames.
- Reference data prepared by the DataEngineer.
Always provide executable Python code in python ... blocks.
When reading data, check for 'clean_data.csv' or the filename mentioned by the DataEngineer.""",
llm_config=llm_config,
)
# Agent 3: Report Writer — synthesizes findings into narratives
report_writer = autogen.AssistantAgent(
name="ReportWriter",
system_message="""You are a report writer. Your responsibilities:
- Read the conversation history to understand what analysis was performed.
- Synthesize findings from the DataEngineer and DataAnalyst into a coherent narrative.
- Structure your report with: Executive Summary, Methodology, Key Findings, Recommendations.
- Write in clear, business-friendly language.
- Save the final report as 'analysis_report.md' using Python file I/O.
- Do NOT perform data analysis — focus solely on writing the report based on others' findings.
Use Python code blocks to save your report to disk.""",
llm_config=llm_config,
)
# User proxy with code execution
user_proxy = autogen.UserProxyAgent(
name="UserProxy",
human_input_mode="NEVER",
max_consecutive_auto_reply=15,
code_execution_config={
"work_dir": "advanced_analysis_output",
"use_docker": False,
"timeout": 300,
},
)
# Create a group chat with all agents
groupchat = autogen.GroupChat(
agents=[user_proxy, data_engineer, data_analyst, report_writer],
messages=[],
max_round=20,
speaker_selection_method="auto", # Let the LLM decide who speaks next
)
# Create a group chat manager
manager = autogen.GroupChatManager(
groupchat=groupchat,
llm_config=llm_config,
system_message="""You are a group chat manager for a data analysis team.
Route messages to the appropriate agent based on the current task:
- Data loading/cleaning → DataEngineer
- Statistical analysis/visualization → DataAnalyst
- Report writing → ReportWriter
- UserProxy handles all code execution.
Guide the team through a logical workflow: Engineer → Analyst → Writer.
Ensure each agent completes its task before moving to the next phase.""",
)
Orchestrating the Multi-Agent Workflow
# Initiate the group chat with a comprehensive task
user_proxy.initiate_chat(
manager,
message="""Team, please perform a complete analysis of the sales dataset at 'sales_data.csv':
**DataEngineer tasks:**
- Load the data and handle any quality issues.
- Create a 'month' column from the date field.
- Calculate revenue (sales * units_sold) as a new column.
- Save the cleaned dataset as 'clean_sales_data.csv'.
**DataAnalyst tasks:**
- Load the cleaned data.
- Analyze sales trends by month — create a monthly sales line chart saved as 'monthly_sales.png'.
- Analyze regional performance — create a grouped bar chart of revenue by region and product category, saved as 'region_category_revenue.png'.
- Calculate correlation between customer satisfaction and revenue.
- Identify any seasonal patterns.
**ReportWriter tasks:**
- After both Engineer and Analyst complete their work, write a comprehensive report.
- Include all key metrics, charts referenced by filename, and actionable recommendations.
- Save as 'final_sales_report.md'.
Please proceed through the workflow sequentially.""",
clear_history=True,
)
Adding Custom Tools and Skills
AutoGen agents can be enhanced with custom functions registered as tools. This is particularly useful for data analysis when you have proprietary algorithms, database connections, or specialized statistical methods.
import autogen
from typing import Dict, List, Optional
import pandas as pd
import numpy as np
# Define custom tools as regular Python functions
def detect_anomalies_iqr(dataframe: pd.DataFrame, column: str, multiplier: float = 1.5) -> Dict:
"""Detect anomalies in a column using the IQR method."""
q1 = dataframe[column].quantile(0.25)
q3 = dataframe[column].quantile(0.75)
iqr = q3 - q1
lower_bound = q1 - multiplier * iqr
upper_bound = q3 + multiplier * iqr
anomalies = dataframe[(dataframe[column] < lower_bound) | (dataframe[column] > upper_bound)]
return {
"column": column,
"lower_bound": lower_bound,
"upper_bound": upper_bound,
"anomaly_count": len(anomalies),
"anomaly_percentage": round(len(anomalies) / len(dataframe) * 100, 2),
"anomalous_rows": anomalies.head(10).to_dict("records"),
}
def forecast_linear_trend(dataframe: pd.DataFrame, date_column: str, value_column: str, periods: int = 30) -> Dict:
"""Generate a simple linear trend forecast."""
from sklearn.linear_model import LinearRegression
df = dataframe.copy()
df["ordinal"] = (pd.to_datetime(df[date_column]) - pd.to_datetime(df[date_column]).min()).dt.days
X = df[["ordinal"]].values
y = df[value_column].values
model = LinearRegression()
model.fit(X, y)
future_days = np.arange(X[-1, 0] + 1, X[-1, 0] + periods + 1).reshape(-1, 1)
forecast = model.predict(future_days)
return {
"r_squared": round(model.score(X, y), 3),
"slope": round(model.coef_[0], 4),
"intercept": round(model.intercept_, 2),
"forecast_values": forecast.tolist(),
"forecast_days": [int(d[0]) for d in future_days],
}
# Register tools with the assistant agent
llm_config_with_tools = {
"config_list": [
{
"model": "gpt-4",
"api_key": os.environ.get("OPENAI_API_KEY"),
}
],
"temperature": 0.1,
"timeout": 120,
"tools": [
{
"type": "function",
"function": {
"name": "detect_anomalies_iqr",
"description": "Detect statistical anomalies in a DataFrame column using the Interquartile Range method. Returns bounds and anomalous rows.",
"parameters": {
"type": "object",
"properties": {
"dataframe": {"type": "object", "description": "The pandas DataFrame to analyze."},
"column": {"type": "string", "description": "Column name to check for anomalies."},
"multiplier": {"type": "number", "description": "IQR multiplier (default 1.5)."},
},
"required": ["dataframe", "column"],
},
},
},
{
"type": "function",
"function": {
"name": "forecast_linear_trend",
"description": "Fit a linear regression trend and forecast future values. Returns R², slope, and forecasted values.",
"parameters": {
"type": "object",
"properties": {
"dataframe": {"type": "object", "description": "DataFrame with date and value columns."},
"date_column": {"type": "string", "description": "Name of the date column."},
"value_column": {"type": "string", "description": "Name of the numeric column to forecast."},
"periods": {"type": "integer", "description": "Number of future periods to forecast."},
},
"required": ["dataframe", "date_column", "value_column"],
},
},
},
],
}
# Create the enhanced agent with tool access
enhanced_analyst = autogen.AssistantAgent(
name="EnhancedAnalyst",
system_message="""You are a senior data analyst with access to specialized tools:
- Use detect_anomalies_iqr() to find statistical outliers in any numeric column.
- Use forecast_linear_trend() to project future values based on historical trends.
When a user asks about anomalies or forecasting, invoke these tools via function calling.
Combine tool outputs with your own analysis code for comprehensive answers.""",
llm_config=llm_config_with_tools,
)
# Register the actual function implementations with the user proxy
user_proxy_with_tools = autogen.UserProxyAgent(
name="UserProxy",
human_input_mode="NEVER",
max_consecutive_auto_reply=10,
code_execution_config={
"work_dir": "tool_analysis_output",
"use_docker": False,
},
function_map={
"detect_anomalies_iqr": detect_anomalies_iqr,
"forecast_linear_trend": forecast_linear_trend,
},
)
# Run analysis with tool-using agent
user_proxy_with_tools.initiate_chat(
enhanced_analyst,
message="""Load 'sales_data.csv' and perform this analysis:
1. Use detect_anomalies_iqr on the 'sales' column and report any outliers found.
2. Use forecast_linear_trend to predict sales for the next 30 days.
3. Create a visualization showing historical sales plus the forecasted trend with a confidence-style shading.
4. Save the combined chart as 'sales_forecast.png'.""",
)
The function_map in the user proxy is where you wire up the actual Python functions to the tool descriptions that the LLM sees. When the assistant decides to call a tool, AutoGen routes that call to the corresponding function in the map, executes it, and returns the result to the conversation.
Handling Large Datasets Efficiently
When working with large datasets (hundreds of thousands or millions of rows), you need strategies to prevent the agent from loading entire datasets into the conversation context or running out of memory during code execution.
# Strategy 1: Provide metadata instead of raw data in prompts
import pandas as pd
large_df = pd.read_parquet("large_sales_data.parquet")
# Extract metadata for the agent's context
metadata_summary = {
"shape": large_df.shape,
"columns": list(large_df.columns),
"dtypes": {col: str(dtype) for col, dtype in large_df.dtypes.items()},
"memory_usage_mb": round(large_df.memory_usage(deep=True).sum() / 1024 / 1024, 2),
"sample_rows": large_df.head(3).to_dict("records"),
"numeric_summary": large_df.describe().to_dict(),
"missing_values": large_df.isnull().sum().to_dict(),
}
# Pass metadata in the message, not the full dataset
user_proxy.initiate_chat(
data_analyst,
message=f"""I have a large dataset with these characteristics:
{json.dumps(metadata_summary, indent=2, default=str)}
The full file is at 'large_sales_data.parquet'. Please:
1. Write code that uses chunked reading (use chunksize parameter) to process the file.
2. Compute monthly revenue aggregates using efficient groupby operations.
3. Identify the top 100 customers by total spend — use memory-efficient methods.
4. Save aggregated results to 'monthly_revenue_summary.csv'.
Important: Do NOT load the entire dataset into memory at once. Use chunked processing.""",
)
# Strategy 2: Configure execution limits
user_proxy_large_data = autogen.UserProxyAgent(
name="UserProxy",
human_input_mode="NEVER",
max_consecutive_auto_reply=8,
code_execution_config={
"work_dir": "large_data_output",
"use_docker": True, # Docker provides memory isolation
"timeout": 600, # Longer timeout for large data processing
"last_n_messages": 2, # Only keep last 2 messages in context to save tokens
},
)
Best Practices for Data Analysis Agents
1. Craft Detailed System Messages
The system message is the most important configuration element. Be explicit about what the agent should do, how it should format output, what libraries are available, and how to handle errors. A vague system message leads to inconsistent results.
# Good system message example
system_message = """You are a data analyst with these specific capabilities:
- Write Python code using pandas (as pd), matplotlib.pyplot (as plt), seaborn (as sns), numpy (as np).
- Always load data with pd.read_csv() or appropriate reader.
- When creating charts: set figsize=(10, 6), use plt.title(), plt.xlabel(), plt.ylabel(), and plt.legend().
- Save ALL charts with plt.savefig('descriptive_name.png', dpi=150, bbox_inches='tight') followed by plt.close().
- If you encounter an error: read the full traceback, identify the root cause, fix the code, and retry.
- Never use plt.show() — always use plt.savefig().
- After completing analysis tasks, provide a concise summary of findings.
- Wrap ALL code in python ... blocks."""
# Poor system message example (avoid this)
bad_system_message = """You are a helpful assistant that can analyze data."""
2. Start Small and Iterate
For complex analysis tasks, break them into multiple smaller conversations rather than one massive prompt. This gives you more control and lets you inspect intermediate results.
# Phase 1: Data loading and cleaning
user_proxy.initiate_chat(
data_analyst,
message="First, load 'raw_sales.csv' and clean it. Handle missing values and save as 'clean_sales.csv'.",
clear_history=True,
)
# Phase 2: Exploratory analysis (inspect clean_sales.csv manually first if needed)
user_proxy.initiate_chat(
data_analyst,
message="Now load 'clean_sales.csv' and create summary statistics and distributions.",
clear_history=False, # Keep previous context for continuity
)
# Phase 3: Deep dive
user_proxy.initiate_chat(
data_analyst,
message="Great. Now focus on regional performance — create detailed breakdowns.",
clear_history=False,
)
3. Implement Guardrails and Validation
Agents can sometimes produce unexpected code. Implement validation to catch issues early:
# Custom executor with validation
import re
import autogen
class ValidatedUserProxy(autogen.UserProxyAgent):
"""User proxy that validates code before execution."""
def __init__(self, *args, blocked_patterns=None, **kwargs):
super().__init__(*args, **kwargs)
self.blocked_patterns = blocked_patterns or []
def execute_code(self, code: str):
# Check for potentially harmful operations
dangerous_patterns = [
r"os\.system",
r"subprocess\.",
r"shutil\.rmtree",
r"os\.remove",
r"__import__\s*\(\s*['\"]os['\"]",
]
for pattern in dangerous_patterns + self.blocked_patterns:
if re.search(pattern, code):
return f"⚠️ Blocked: Code contains potentially dangerous operation matching pattern '{pattern}'"
# Check file write destinations
write_pattern = r"(?:pd\.\w+\.to_\w+|plt\.savefig|open)\s*\(\s*['\"]([^'\"]+)['\"]"
for match in re.finditer(write_pattern, code):
filepath = match.group(1)
if filepath.startswith("/") or ".." in filepath:
return f"⚠️ Blocked: File write to path '{filepath}' is not allowed. Use relative paths within the work directory."
return super().execute_code(code)
# Use the validated proxy
safe_proxy = ValidatedUserProxy(
name="SafeProxy",
human_input_mode="NEVER",
max_consecutive_auto_reply=10,
code_execution_config={"work_dir": "safe_output", "use_docker": True},
)
4. Manage Context Window and Costs
Long conversations with many code executions can consume significant tokens. Use these strategies to control costs:
llm_config_cost_optimized = {
"config_list": [
{
"model": "gpt-4",
"api_key": os.environ.get("OPENAI_API_KEY"),
}
],
"temperature": 0.1,
"timeout": 120,
# Control token usage
"max_tokens": 2000, # Limit response length
"preserve_history": False, # Don't send full history on every call (use with caution)
}
# Alternatively, use a cheaper model for simpler tasks
llm_config_mixed = {
"config_list": [
{
"model": "gpt-4",
"api_key": os.environ.get("OPENAI_API_KEY"),
},
{
"model": "gpt-3.5-turbo",
"api_key": os.environ.get("OPENAI_API_KEY"),
}
],
"temperature": 0.1,
# AutoGen will try models in order — GPT-4 first, fallback to GPT-3.5
}
5. Log and Persist Results
Always save conversation logs and generated artifacts. AutoGen provides built-in logging:
import logging
import autogen
# Enable AutoGen's built-in logging
autogen.ChatAgent.set_logger(logging.getLogger("autogen"))
# Configure file logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
handlers=[
logging.FileHandler("autogen_analysis.log"),
logging.StreamHandler(),
],
)
# Also save conversation history programmatically
def save_conversation(agent, filename="conversation_history.json"):
"""Save agent conversation history to disk."""
history = agent.chat_history
with open(filename, "w") as f:
json.dump([msg.to_dict() if hasattr(msg, "to_dict") else str(msg) for msg in history], f, indent=2)
# Call after analysis completes
save_conversation(user_proxy, "analysis_conversation.json")
6. Use Docker for Production Safety
When building agents that will analyze data from untrusted sources or be used by multiple team members, always enable Docker sandboxing:
code_execution_config_production = {
"work_dir": "production_output",
"use_docker": True,
"docker_image": "python:3.11-slim", # Specify exact image
"timeout": 600,
# Mount a volume for persistent data access
# Note: Docker must be installed and running on the host machine
}
# Pre-build a Docker image with required packages
# Create a Dockerfile:
# FROM python:3.11-slim
# RUN pip install pandas matplotlib seaborn numpy scikit-learn openpyxl
# Build: docker build -t autogen-data-analysis:latest .
# Then use: "docker_image": "autogen-data-analysis:latest"
Common Pitfalls and Troubleshooting
- Agent generates code but doesn't execute it: Ensure the code is wrapped in
pythonblocks. The user proxy only detects and executes fenced code blocks. - Infinite loops of error/correction: Lower
max_consecutive_auto_replyto a smaller number (e.g., 5). Some errors may be fundamentally unsolvable within the sandbox. - Missing file errors: The agent's working directory starts fresh each session. Either provide absolute paths or ensure files are in the
work_dirbefore the conversation starts. - Charts not saving: The agent may call
plt.show()instead ofplt.savefig(). Explicitly instruct it to save files in the system message. - Token limit exceeded: For very long conversations, implement a conversation summarization step or periodically clear history and reload essential context from saved files.
- Agent hallucinates data: Always require the agent to read actual data files rather than inventing sample data. Include explicit instructions to "load the provided file" rather than "create sample data."
Complete End-to-End Example
Here is a complete, production-ready script that combines all the concepts discussed:
#!/usr/bin/env python3
"""
Complete Data Analysis Agent System with AutoGen
Features: multi-agent team, custom tools, validation, logging, and report generation.
"""
import autogen
import pandas as pd
import numpy as np
import os
import json
import logging
from datetime import datetime
from typing import Dict
# ── Configuration ────────────────────────────────────────────
# Setup logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
handlers=[
logging.FileHandler("autogen_analysis.log"),
logging.StreamHandler(),
],
)
logger = logging.getLogger(__name__)
# LLM configuration
llm_config = {
"config_list": [
{
"model": "gpt-4",
"api_key": os.environ.get("OPENAI_API_KEY"),
}
],
"temperature": 0.1,
"timeout": 120,
}
# Output directory
OUTPUT_DIR = f"analysis_output_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
os.makedirs(OUTPUT_DIR, exist_ok=True)
# ── Custom Tools ─────────────────────────────────────────────
def calculate_metrics(dataframe: pd.DataFrame, column: str) -> Dict:
"""Calculate comprehensive metrics for a numeric column."""
series = dataframe[column].dropna()
return {
"mean": round(series.mean(), 2),
"median": round(series.median(), 2),
"std": round(series.std(), 2),
"skewness": round(series.skew(), 2),
"kurtosis": round(series.kurtosis(), 2),
"percentile_25": round(series.quantile(0.25), 2),
"percentile_75": round(series.quantile(0.75), 2),
"percentile_95": round(series.quantile(0.95), 2),
"coefficient_of_variation": round(series.std() / series.mean() * 100, 2),
}
# ── Agent Definitions ────────────────────────────────────────
data_engineer = autogen.AssistantAgent(
name="DataEngineer",
system_message=f"""Data Engineer agent. Responsibilities:
- Load data files and perform cleaning/transformation.
- Handle missing values, type conversions, and feature engineering.
- Save cleaned datasets to '{OUTPUT_DIR}/clean_data.csv'.
- Document all steps taken.
Always output code in python ... blocks.""",
llm_config=llm_config,
)
data_analyst = autogen.AssistantAgent(
name="DataAnalyst",
system_message=f"""Data Analyst agent. Responsibilities:
- Statistical analysis and visualization.
- Use the calculate_metrics tool for numeric columns.
- Create professional charts and save them to '{OUTPUT_DIR}/'.
- Always use plt.savefig('{OUTPUT_DIR}/chart_name.png', dpi=150, bbox_inches='tight') then plt.close().
- Never use plt.show().
Always output code in python ... blocks.""",
llm_config={
**llm_config,
"tools": [
{
"type": "function",
"function": {
"name": "calculate_metrics",
"description": "Calculate comprehensive statistical metrics for a numeric DataFrame column.",
"parameters": {
"type": "object",
"properties": {
"dataframe": {"type": "object", "description": "Pandas DataFrame."},
"column": {"type": "string", "description": "Numeric column name."},
},
"required": ["dataframe", "column"],
},
},
}
],
},
)
report_writer = autogen.AssistantAgent(
name="ReportWriter",
system_message=f"""Report Writer agent. Responsibilities:
- Synthesize