Build a Data Analyst Agent with Claude API

import anthropic
import pandas as pd
import io
import sys
import subprocess
import json
import tempfile
import os
from pathlib import Path

client = anthropic.Anthropic()


# ─── Data Loading ──────────────────────────────────────────────────────────────

def load_csv(file_path: str) -> tuple[pd.DataFrame, str]:
    """
    Load a CSV and return the DataFrame and a schema description for Claude.
    """
    df = pd.read_csv(file_path)
    
    # Build schema summary
    schema_lines = [
        f"File: {Path(file_path).name}",
        f"Rows: {len(df):,}",
        f"Columns: {len(df.columns)}",
        "",
        "Column schema:"
    ]
    
    for col in df.columns:
        dtype = str(df[col].dtype)
        nulls = df[col].isna().sum()
        
        if pd.api.types.is_numeric_dtype(df[col]):
            summary = f"min={df[col].min():.2f}, max={df[col].max():.2f}, mean={df[col].mean():.2f}"
        elif pd.api.types.is_datetime64_any_dtype(df[col]):
            summary = f"from {df[col].min()} to {df[col].max()}"
        else:
            n_unique = df[col].nunique()
            top_values = df[col].value_counts().head(3).to_dict()
            summary = f"{n_unique} unique values, top: {top_values}"
        
        null_note = f", {nulls} nulls" if nulls > 0 else ""
        schema_lines.append(f"  - {col} ({dtype}{null_note}): {summary}")
    
    # Sample rows
    schema_lines.append("")
    schema_lines.append("First 3 rows:")
    schema_lines.append(df.head(3).to_string())
    
    return df, "\n".join(schema_lines)


def save_csv_for_sandbox(df: pd.DataFrame, temp_dir: str) -> str:
    """Save DataFrame as CSV for the sandbox to read."""
    path = os.path.join(temp_dir, "data.csv")
    df.to_csv(path, index=False)
    return path


# ─── Sandboxed Code Executor ───────────────────────────────────────────────────

ALLOWED_IMPORTS = {
    "pandas", "numpy", "json", "math", "statistics",
    "collections", "itertools", "datetime", "re"
}


def validate_code(code: str) -> tuple[bool, str]:
    """
    Basic static validation of generated code.
    Blocks dangerous operations before execution.
    """
    forbidden_patterns = [
        "import os", "import sys", "import subprocess", "import socket",
        "import shutil", "__import__", "open(", "eval(", "exec(",
        "compile(", "globals()", "locals()", "getattr", "setattr",
        "delattr", "vars(", "dir(", "__class__", "__bases__",
        "import requests", "import urllib", "import http"
    ]
    
    code_lower = code.lower()
    for pattern in forbidden_patterns:
        if pattern.lower() in code_lower:
            return False, f"Forbidden pattern detected: '{pattern}'"
    
    return True, "OK"


def execute_code_sandbox(code: str, data_csv_path: str) -> dict:
    """
    Execute analysis code in a subprocess sandbox.
    
    The code receives a 'df' variable (the loaded DataFrame) and should
    print results to stdout in JSON format.
    """
    # Validate before execution
    is_safe, reason = validate_code(code)
    if not is_safe:
        return {
            "success": False,
            "error": f"Code validation failed: {reason}",
            "output": None
        }
    
    # Wrap the code with DataFrame loading preamble and output capture
    wrapped_code = f"""
import pandas as pd
import numpy as np
import json
import math
import statistics
from collections import Counter

# Load the data
df = pd.read_csv("{data_csv_path}")

# User code starts here
{code}
"""
    
    # Write to temp file and execute in subprocess
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False, 
                                      encoding="utf-8") as f:
        f.write(wrapped_code)
        script_path = f.name
    
    try:
        result = subprocess.run(
            [sys.executable, script_path],
            capture_output=True,
            text=True,
            timeout=30,  # 30 second limit
            env={
                "PATH": os.environ.get("PATH", ""),
                "PYTHONPATH": os.environ.get("PYTHONPATH", ""),
                "HOME": os.environ.get("HOME", "")
            }
        )
        
        if result.returncode == 0:
            return {
                "success": True,
                "output": result.stdout.strip(),
                "error": None
            }
        else:
            return {
                "success": False,
                "output": result.stdout.strip() or None,
                "error": result.stderr.strip()
            }
    
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "output": None,
            "error": "Code execution timed out after 30 seconds."
        }
    
    finally:
        # Always clean up the temp file
        if os.path.exists(script_path):
            os.unlink(script_path)


# ─── Tool Definitions ──────────────────────────────────────────────────────────

ANALYST_TOOLS = [
    {
        "name": "run_analysis_code",
        "description": """Execute Python code to analyse the dataset.
        
The code has access to a pandas DataFrame called 'df' containing the full dataset.
Use print() to output results. Output will be returned to you for interpretation.

IMPORTANT: 
- Always print your results, not just compute them
- For numbers, round to 2 decimal places
- For DataFrames, use .to_string() or .to_dict()
- Do not import any modules not in the standard data science stack
""",
        "input_schema": {
            "type": "object",
            "properties": {
                "code": {
                    "type": "string",
                    "description": "Valid Python code to execute against the DataFrame df"
                },
                "description": {
                    "type": "string",
                    "description": "One-line description of what this code does"
                }
            },
            "required": ["code", "description"]
        }
    }
]


# ─── Data Analyst Agent ────────────────────────────────────────────────────────

class DataAnalystAgent:
    """Conversational data analyst backed by Claude."""
    
    def __init__(self, csv_path: str, model: str = "claude-sonnet-4-6"):
        self.model = model
        self.df, self.schema = load_csv(csv_path)
        self.temp_dir = tempfile.mkdtemp()
        self.data_path = save_csv_for_sandbox(self.df, self.temp_dir)
        self.conversation_history = []
        
        print(f"Loaded dataset:")
        print(self.schema)
        print("\nData analyst agent ready. Ask me anything about this data.\n")
    
    def _build_system_prompt(self) -> str:
        return f"""You are an expert data analyst. You have access to a Python execution tool
that lets you run pandas and numpy code against a dataset.

DATASET SCHEMA:
{self.schema}

Approach:
1. Understand the user's question
2. Write precise Python code to answer it
3. Observe the output and interpret it clearly  
4. If the first code run is incomplete, run additional code to refine
5. Give a clear, jargon-free answer citing the numbers you found

Always explain what the numbers mean, not just what they are."""
    
    def ask(self, question: str) -> str:
        """
        Ask a question about the dataset.
        Returns Claude's final natural language answer.
        """
        self.conversation_history.append({
            "role": "user",
            "content": question
        })
        
        max_iterations = 8
        iteration = 0
        
        while iteration < max_iterations:
            iteration += 1
            
            response = client.messages.create(
                model=self.model,
                max_tokens=4096,
                system=self._build_system_prompt(),
                tools=ANALYST_TOOLS,
                messages=self.conversation_history
            )
            
            # Collect tool calls and text
            tool_calls = []
            text_blocks = []
            
            for block in response.content:
                if block.type == "tool_use":
                    tool_calls.append(block)
                elif block.type == "text":
                    text_blocks.append(block.text)
            
            # Add assistant response to history
            self.conversation_history.append({
                "role": "assistant",
                "content": response.content
            })
            
            # If no tool calls, Claude is done
            if response.stop_reason == "end_turn" or not tool_calls:
                answer = "\n".join(text_blocks)
                return answer
            
            # Execute all tool calls and collect results
            tool_results = []
            
            for tool_call in tool_calls:
                code = tool_call.input.get("code", "")
                description = tool_call.input.get("description", "")
                
                print(f"  Executing: {description}")
                
                result = execute_code_sandbox(code, self.data_path)
                
                if result["success"]:
                    content = result["output"] or "(No output produced)"
                else:
                    content = f"Error: {result['error']}"
                    if result["output"]:
                        content += f"\nPartial output: {result['output']}"
                
                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": tool_call.id,
                    "content": content
                })
            
            # Add tool results to history
            self.conversation_history.append({
                "role": "user",
                "content": tool_results
            })
        
        return "Analysis reached maximum iterations. Please try a more specific question."
    
    def run_interactive(self):
        """Run an interactive analysis session."""
        print("Type your question or 'quit' to exit.\n")
        
        while True:
            question = input("You: ").strip()
            if question.lower() in ("quit", "exit"):
                break
            if question:
                answer = self.ask(question)
                print(f"\nAnalyst: {answer}\n")

if __name__ == "__main__":
    # Point to your CSV file
    agent = DataAnalystAgent("sales_data.csv")
    
    # Single question mode
    print(agent.ask("Which product category generated the most revenue last quarter?"))
    print(agent.ask("What is the average order value and how does it vary by region?"))
    
    # Or run an interactive session
    # agent.run_interactive()

You: Which month had the highest total sales?

Executing: Calculate monthly sales totals and find the maximum
Executing: Format monthly breakdown for top 3 months

Analyst: March 2025 was your strongest month with £847,230 in total sales — 
about 23% above average monthly revenue of £689,500.

The top three months were:
- March 2025: £847,230
- October 2024: £812,450  
- November 2024: £798,100

There's a clear seasonal pattern: autumn/early spring tends to outperform summer months in this dataset.