Building a Data Analysis Agent

The Data Analysis Agent Vision

The Agent's Overall Architecture

DATA_AGENT_SYSTEM_PROMPT = '''
You are a data analysis agent. Given a CSV file and a question, answer it through iterative Python code.

Phases to follow:
1. EXPLORE: Load the data, print shape, dtypes, head(), describe(), and check for nulls.
2. CLEAN: Handle missing values, fix dtypes, remove outliers.
3. ANALYZE: Compute statistics, groupbys, correlations - whatever answers the question.
4. REPORT: Generate a matplotlib chart saved to /workspace/chart.png and write a text summary.

Rules:
- Write code in ```python ... ``` blocks.
- Save intermediate results to /workspace/ for use in later steps.
- Print key findings after each computation so you can observe them.
- When all phases are done, say TASK COMPLETE and summarize the answer.
'''

Phase 1: Data Exploration Code

import pandas as pd
import numpy as np

# Load dataset
df = pd.read_csv('/workspace/data.csv')

# Basic exploration
print('=== SHAPE ===' )
print(df.shape)

print('\n=== DTYPES ===')
print(df.dtypes)

print('\n=== HEAD ===')
print(df.head())

print('\n=== DESCRIBE ===')
print(df.describe(include='all'))

print('\n=== NULL COUNTS ===')
print(df.isnull().sum())

print('\n=== UNIQUE VALUES (categorical) ===')
for col in df.select_dtypes(include='object').columns:
    print(f'{col}: {df[col].nunique()} unique values: {df[col].unique()[:5]}')

# Save for next iteration
df.to_parquet('/workspace/raw.parquet')
print('\nData saved to workspace.')

Phase 2: Data Cleaning

import pandas as pd
import numpy as np

# Load from previous iteration
df = pd.read_parquet('/workspace/raw.parquet')

# Fill numeric nulls with median (robust to outliers)
for col in df.select_dtypes(include='number').columns:
    if df[col].isnull().sum() > 0:
        median_val = df[col].median()
        df[col] = df[col].fillna(median_val)
        print(f'Filled nulls in {col} with median {median_val:.2f}')

# Drop rows where categorical key columns are null
df = df.dropna(subset=['category', 'date'])

# Fix date column type
df['date'] = pd.to_datetime(df['date'], errors='coerce')
df = df.dropna(subset=['date'])  # drop unparseable dates

# Remove clear outliers: revenue > 3 std deviations from mean
if 'revenue' in df.columns:
    z_scores = (df['revenue'] - df['revenue'].mean()) / df['revenue'].std()
    df = df[z_scores.abs() < 3]
    print(f'Removed outliers, rows remaining: {len(df)}')

df.to_parquet('/workspace/cleaned.parquet')
print('Cleaning complete:', df.shape)

Phase 3: Answering the Question

import pandas as pd
import matplotlib
matplotlib.use('Agg')  # non-interactive backend for server use
import matplotlib.pyplot as plt

df = pd.read_parquet('/workspace/cleaned.parquet')

# Analysis: revenue by category for last quarter
df['date'] = pd.to_datetime(df['date'])
last_quarter = df[df['date'] >= '2024-10-01']

revenue_by_category = (
    last_quarter
    .groupby('category')['revenue']
    .sum()
    .sort_values(ascending=False)
)
print('Revenue by category (last quarter):')
print(revenue_by_category)
print(f'\nTop category: {revenue_by_category.index[0]} (${revenue_by_category.iloc[0]:,.0f})')

# Save analysis result
revenue_by_category.to_json('/workspace/analysis_result.json')

Phase 4: Generating Charts

import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import json

with open('/workspace/analysis_result.json') as f:
    data = json.load(f)

categories = list(data.keys())
values = [data[k] / 1e6 for k in categories]  # convert to millions

fig, ax = plt.subplots(figsize=(10, 6))
bars = ax.barh(categories, values, color='steelblue', edgecolor='white')

# Add value labels on bars
for bar, val in zip(bars, values):
    ax.text(bar.get_width() + 0.05, bar.get_y() + bar.get_height()/2,
            f'${val:.1f}M', va='center', fontsize=11)

ax.set_xlabel('Revenue ($ Millions)', fontsize=12)
ax.set_title('Revenue by Product Category - Q4 2024', fontsize=14, fontweight='bold')
ax.invert_yaxis()  # highest on top
plt.tight_layout()
plt.savefig('/workspace/chart.png', dpi=150, bbox_inches='tight')
print('Chart saved to /workspace/chart.png')

Generating the Written Report

import json

with open('/workspace/analysis_result.json') as f:
    revenue_data = json.load(f)

top_category = max(revenue_data, key=revenue_data.get)
top_revenue = revenue_data[top_category]
total_revenue = sum(revenue_data.values())
share = top_revenue / total_revenue * 100

report = f'''# Q4 2024 Revenue Analysis

## Answer
The **{top_category}** category generated the highest revenue in Q4 2024:
**${top_revenue:,.0f}** ({share:.1f}% of total Q4 revenue).

## Key Findings
- Total Q4 revenue: ${total_revenue:,.0f}
- Top 3 categories: {list(revenue_data.items())[:3]}
- The top category outperformed the average by {top_revenue / (total_revenue / len(revenue_data)):.1f}x

## Chart
See chart.png for the full breakdown by category.

## Recommendation
Consider reallocating marketing budget toward {top_category} to
capitalize on its demonstrated strong performance.
'''

with open('/workspace/report.md', 'w') as f:
    f.write(report)
print('TASK COMPLETE')
print(report)

Assembling the Agent Orchestrator

import shutil
from pathlib import Path

def run_data_analysis_agent(csv_path: str, question: str) -> dict:
    workspace = setup_workspace()
    shutil.copy(csv_path, workspace / 'data.csv')
    
    messages = [
        {'role': 'system', 'content': DATA_AGENT_SYSTEM_PROMPT},
        {'role': 'user', 'content': f'File: /workspace/data.csv\nQuestion: {question}'}
    ]
    
    state = ExecutionState(task=question)
    
    for iteration in range(15):  # max 15 iterations
        response = llm.complete(messages)
        code = extract_code_block(response)
        
        if not code or 'TASK COMPLETE' in response:
            break
        
        stdout, stderr = execute_in_docker(code, workspace=workspace)
        observation = format_observation(stdout, stderr)
        messages += [
            {'role': 'assistant', 'content': response},
            {'role': 'user', 'content': observation}
        ]
        state.after_execution(stdout, f'iteration_{iteration}')
    
    return {
        'report': (workspace / 'report.md').read_text() if (workspace / 'report.md').exists() else '',
        'chart_path': str(workspace / 'chart.png'),
        'iterations': state.iteration
    }

Handling Ambiguous Questions

def handle_ambiguous_question(question: str, df_summary: dict) -> str:
    clarification_prompt = f'''The user asked: '{question}'

Available data:
- Columns: {df_summary['columns']}
- Date range: {df_summary['date_range']}
- Metrics available: {df_summary['numeric_columns']}

Is the question clear enough to answer definitively?
If yes, restate the specific interpretation you will use.
If no, list the 2-3 clarifications needed to proceed.'''
    
    response = llm.complete([{'role': 'user', 'content': clarification_prompt}])
    return response

Quality Checks on Agent Output

def validate_analysis_output(workspace: str, question: str, report: str) -> dict:
    validation_prompt = f'''Original question: {question}

Agent report:
{report[:2000]}

Rate this analysis on a scale of 1-5 for:
1. Does it directly answer the question? (1=No, 5=Yes)
2. Are specific numbers cited? (1=No, 5=Yes)
3. Is the conclusion clearly stated? (1=No, 5=Yes)

Return JSON: {{"question_answered": N, "numbers_cited": N, "clear_conclusion": N, "overall": N}}'''
    
    result = llm.complete([{'role': 'user', 'content': validation_prompt}],
                          response_format={'type': 'json_object'})
    scores = json.loads(result)
    passed = scores['overall'] >= 4
    return {'scores': scores, 'passed': passed}

Extending the Agent with Domain Tools

Quick Check

Lesson Recap