How to Build an AI Data Analyst That Analyses CSV Files and Creates Charts

ai-data-analyst/
├── app/
│   ├── __init__.py
│   ├── cleaner_profiler.py
│   ├── sandbox_runner.py
│   └── chart_generator.py
├── tests/
│   └── test_data_analyst.py
└── requirements.txt

from typing import Dict, Any
import pandas as pd
import numpy as np

def clean_and_profile_dataframe(df: pd.DataFrame) -> tuple[pd.DataFrame, Dict[str, Any]]:
    """Cleans whitespace, coerces currency/numeric strings, and extracts a compact statistical profile."""
    cleaned = df.copy()
    cleaning_actions = []

    # 1. Strip column name whitespace
    cleaned.columns = [str(c).strip() for c in cleaned.columns]

    # 2. Drop exact duplicate rows
    dup_count = int(cleaned.duplicated().sum())
    if dup_count > 0:
        cleaned = cleaned.drop_duplicates().reset_index(drop=True)
        cleaning_actions.append(f"Dropped {dup_count} duplicate rows.")

    # 3. Coerce currency/percentage object columns where >80% parse as numeric
    for col in cleaned.select_dtypes(include=["object"]).columns:
        sample_non_null = cleaned[col].dropna().astype(str).str.strip()
        if len(sample_non_null) == 0:
            continue
        stripped = sample_non_null.str.replace(r"^[\$₹€£,]+|%$", "", regex=True).str.replace(",", "", regex=False)
        coerced = pd.to_numeric(stripped, errors="coerce")
        if coerced.notna().mean() >= 0.85:
            cleaned[col] = pd.to_numeric(
                cleaned[col].astype(str).str.replace(r"[\$₹€£,%]", "", regex=True).str.strip(),
                errors="coerce",
            )
            cleaning_actions.append(f"Converted column '{col}' from string to float.")

    # 4. Build compact profile for the LLM
    col_profiles = {}
    for col in cleaned.columns:
        series = cleaned[col]
        info: Dict[str, Any] = {
            "dtype": str(series.dtype),
            "null_count": int(series.isna().sum()),
            "unique_count": int(series.nunique(dropna=True)),
        }
        if pd.api.types.is_numeric_dtype(series):
            desc = series.describe()
            info["mean"] = round(float(desc.get("mean", 0.0)), 4) if not np.isnan(desc.get("mean", np.nan)) else None
            info["std"] = round(float(desc.get("std", 0.0)), 4) if not np.isnan(desc.get("std", np.nan)) else None
            info["min"] = float(desc.get("min", 0.0)) if not np.isnan(desc.get("min", np.nan)) else None
            info["max"] = float(desc.get("max", 0.0)) if not np.isnan(desc.get("max", np.nan)) else None
        else:
            info["top_values"] = series.dropna().astype(str).value_counts().head(5).to_dict()
        col_profiles[col] = info

    profile = {
        "row_count": int(len(cleaned)),
        "column_count": int(len(cleaned.columns)),
        "cleaning_actions": cleaning_actions,
        "columns": col_profiles,
    }
    return cleaned, profile

import ast
from typing import Dict, Any, Set
import pandas as pd
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

ALLOWED_MODULES: Set[str] = {"pandas", "numpy", "matplotlib", "matplotlib.pyplot", "math"}
FORBIDDEN_CALL_NAMES: Set[str] = {
    "eval", "exec", "open", "compile", "__import__", "getattr", "setattr",
    "delattr", "globals", "locals", "input", "breakpoint", "system", "popen",
}
FORBIDDEN_ATTRIBUTES: Set[str] = {
    "to_pickle", "read_pickle", "to_sql", "to_csv", "to_excel", "system",
    "__subclasses__", "__bases__", "__globals__", "__code__",
}

def validate_analysis_code_ast(code_str: str) -> None:
    """Inspects the Python AST to block file I/O, network access, and disallowed imports."""
    tree = ast.parse(code_str)
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                if alias.name not in ALLOWED_MODULES:
                    raise PermissionError(f"Disallowed import: {alias.name}")
        elif isinstance(node, ast.ImportFrom):
            if not node.module or node.module not in ALLOWED_MODULES:
                raise PermissionError(f"Disallowed import from: {node.module}")
        elif isinstance(node, ast.Call):
            if isinstance(node.func, ast.Name) and node.func.id in FORBIDDEN_CALL_NAMES:
                raise PermissionError(f"Disallowed function call: {node.func.id}")
        elif isinstance(node, ast.Attribute):
            if node.attr in FORBIDDEN_ATTRIBUTES or node.attr.startswith("__"):
                raise PermissionError(f"Disallowed attribute access: {node.attr}")

def execute_verified_analysis(
    df: pd.DataFrame,
    code_str: str,
    chart_output_path: str | None = None,
) -> Dict[str, Any]:
    """Validates `code_str`, runs it against a copy of `df`, and saves `fig` if generated."""
    validate_analysis_code_ast(code_str)
    plt.close("all")
    fig, ax = plt.subplots(figsize=(8, 4.5), dpi=150)

    safe_locals: Dict[str, Any] = {
        "df": df.copy(deep=True),
        "pd": pd,
        "np": np,
        "plt": plt,
        "fig": fig,
        "ax": ax,
        "result": None,
    }
    safe_builtins = {
        "len": len, "round": round, "min": min, "max": max, "sum": sum,
        "abs": abs, "float": float, "int": int, "str": str, "dict": dict, "list": list,
    }
    exec(code_str, {"__builtins__": safe_builtins}, safe_locals)

    saved_chart = None
    if chart_output_path and len(ax.patches) + len(ax.lines) + len(ax.collections) > 0:
        fig.tight_layout()
        fig.savefig(chart_output_path, dpi=150)
        saved_chart = chart_output_path
    plt.close(fig)

    res = safe_locals.get("result")
    if isinstance(res, pd.DataFrame):
        res = res.to_dict(orient="records")
    elif isinstance(res, pd.Series):
        res = res.to_dict()

    return {"result": res, "chart_path": saved_chart}

import os
import pandas as pd
import pytest
from app.cleaner_profiler import clean_and_profile_dataframe
from app.sandbox_runner import validate_analysis_code_ast, execute_verified_analysis

def test_csv_cleaner_and_profiler():
    raw_df = pd.DataFrame({
        " Category ": ["Cloud", "Cloud", "Edge", "AI"],
        "Revenue_USD": ["$1,200.50", "$1,200.50", "$850.00", "$2,400.00"],
    })
    cleaned, profile = clean_and_profile_dataframe(raw_df)
    assert profile["row_count"] == 3  # Duplicate row dropped
    assert cleaned["Revenue_USD"].dtype == float
    assert profile["columns"]["Revenue_USD"]["max"] == 2400.0

def test_ast_sandbox_blocks_os_and_dunder_escapes():
    with pytest.raises(PermissionError, match="Disallowed import"):
        validate_analysis_code_ast("import os\nos.system('ls')")
    with pytest.raises(PermissionError, match="Disallowed function call"):
        validate_analysis_code_ast("open('/etc/passwd').read()")
    with pytest.raises(PermissionError, match="Disallowed attribute access"):
        validate_analysis_code_ast("x = df.__class__.__bases__")

def test_verified_pandas_aggregation_and_chart_export(tmp_path):
    df = pd.DataFrame({
        "Segment": ["Enterprise", "SMB", "Enterprise", "SMB"],
        "ARR": [150000, 40000, 210000, 60000],
    })
    code = """
grouped = df.groupby('Segment')['ARR'].sum()
grouped.plot(kind='bar', ax=ax, title='Total ARR by Segment')
result = grouped
"""
    chart_file = str(tmp_path / "arr_chart.png")
    out = execute_verified_analysis(df, code, chart_output_path=chart_file)
    assert out["result"] == {"Enterprise": 360000, "SMB": 100000}
    assert os.path.exists(chart_file)