Code-Executing Data Agent: Let an LLM Answer Questions With pandas, Safely
Hands-on lab · IDE in your browser

Code-Executing Data Agent: Let an LLM Answer Questions With pandas, Safely

Build an agent that answers business questions about a CSV by writing pandas code and running it.

Time
50 min
Checked steps
5
Level
Intermediate
Setup
None
Read step 1

Hands-on labs require Pro · $29.99/mo · cancel anytime

Lab cockpit50 min · 5 stepsSession running
4 / 5 steps passingPut the rules in code · step 5 of 5
prep.py▶ Run✓ Check
"""Turns the raw export into the table the agent analyses. It runs inside the sandbox before the model's code."""import pandas as pd  def prepare(raw):    """Apply DATA_NOTES.md in code: keep the first row of each order_id, parse order_date as a date, and add a    `revenue` column: units * unit_price * (1 - discount) for completed orders, 0.0 for the rest."""    
TerminalOutput

The job

Fernhill Outdoor's managers want answers from the order export without waiting for an analyst. You build the agent that writes and runs the pandas code for them: first a sandbox the code cannot escape, then the context and checks that make its answers right, and finally a guard against questions written to make it misbehave.

5 steps, each checked when you finish it

A check runs your work at the end of every step. Hints and the full solution are there if you get stuck.

  1. 1

    Run code in a sandbox

    Fernhill Outdoor's managers ask questions like "which region sold most in March?".

    You writerun_code()
  2. 2

    Show the model the data

    The agent loop in agent.py is written: the model calls run_python, the sandbox runs the code, the output goes back, and the loop repeats until the model replies with ANSWER: ....

    You writedescribe_data()
  3. 3

    Answers come from code

    The system prompt says "never guess a number you have not computed".

    You writegrounded()
  4. 4

    Refuse dangerous code

    The model writes whatever code the question asks for.

    You writecheck_code()
  5. 5

    Put the rules in code

    DATA_NOTES.md is in the system prompt: count each order once, revenue is units × price × (1 − discount), and only completed orders earn revenue.

    You writeprepare()

Step 1 as it appears in the lab

The lab’s own text. The hint and the solution stay inside the lab.

Step 1: Run code in a sandbox

Fernhill Outdoor's managers ask questions like "which region sold most in March?". An LLM cannot add up 1,400 rows in its head, but it can write pandas code that does. Your agent will run that code, so the first job is a place to run it where a mistake or a trick cannot hurt anything.

sandbox.py sets the limits: TIMEOUT seconds, MAX_OUTPUT characters back to the model, and limits(), which caps memory and CPU. PRELUDE loads the table as df before the model's code runs.

Do this

1. Write run_code(code): run PRELUDE + code in a new Python process inside a temporary folder that holds copies of orders.csv and prep.py. Give the process only a minimal set of environment variables. Return {"ok": True, "output": stdout}, cut to MAX_OUTPUT with a note. On an error, return ok False and the last lines of the error. On a timeout, return ok False and say it was stopped.

2. Run it. The script tries a normal query, an endless loop, a memory bomb, a typo, a huge print and a look at the environment variables.

sandbox.py, the file you edit59 lines
"""Runs Python code written by the model, away from the agent, with limits."""
import ast
import os
import shutil
import subprocess
import sys
import tempfile

DATA = "orders.csv"
TIMEOUT = 5           # seconds of wall time per run
MAX_OUTPUT = 2000     # characters of output returned to the model
MEMORY = 600 << 20    # bytes of address space for the code
PRELUDE = "import pandas as pd\nimport numpy as np\nfrom prep import prepare\ndf = prepare(pd.read_csv('orders.csv'))\n"


# ---------- Step 1: run code in a separate process ----------
def run_code(code):
    """Run `code` after PRELUDE in a fresh Python process, in a temporary folder holding copies of the data and
    prep.py.
    Returns {"ok": bool, "output": text}: stdout on success; on an error, the error's last lines."""
    # TODO (Step 1): in a tempfile.TemporaryDirectory(), copy DATA (as orders.csv) and prep.py, write PRELUDE + code
    # to job.py, and run [sys.executable, "job.py"] there with subprocess.run: cwd=the folder, capture_output=True,
    # text=True, timeout=TIMEOUT, a minimal env (PATH, HOME=the folder, OPENBLAS_NUM_THREADS=1, OMP_NUM_THREADS=1)
    # and preexec_fn=limits.
    # Timeout -> {"ok": False, "output": "Stopped: ..."}; non-zero exit -> {"ok": False, "output": last 3 stderr lines};
    # success -> {"ok": True, "output": stdout cut to MAX_OUTPUT with a note saying it was cut}.
    raise NotImplementedError("Step 1: write run_code()")


def limits():
    """Runs in the child before the code starts: cap its memory and CPU."""
    import resource
    resource.setrlimit(resource.RLIMIT_AS, (MEMORY, MEMORY))
    resource.setrlimit(resource.RLIMIT_CPU, (TIMEOUT, TIMEOUT))


# ---------- Step 4: refuse dangerous code before it runs ----------
ALLOWED_IMPORTS = {"pandas", "numpy", "math", "statistics", "datetime", "collections", "itertools", "re"}
BLOCKED_NAMES = {"open", "exec", "eval", "compile", "__import__", "input", "globals", "locals", "vars", "getattr",
                 "setattr", "delattr", "breakpoint", "exit", "quit"}


def check_code(code):
    """None if the code may run, else the reason it may not. Rules: it must parse; imports only from
    ALLOWED_IMPORTS; no use of a BLOCKED_NAMES name; no attribute starting with an underscore; no file access
    through pandas or numpy (attributes in FILE_ATTRS or starting with read_)."""
    return None


FILE_ATTRS = {"to_csv", "to_excel", "to_json", "to_parquet", "to_pickle", "tofile", "save", "savetxt", "to_sql", "to_hdf",
           "load", "loadtxt", "fromfile", "genfromtxt"}


def safe_run(code):
    """check_code(), then run_code()."""
    reason = check_code(code)
    if reason:
        return {"ok": False, "output": f"Refused before running: {reason}. Only analyse df with pandas and print results."}
    return run_code(code)
Provided for you:DATA_NOTES.mdagent.pyharness.pyorders.csvprep.pyquestions.jsontry_it.py

Frequently asked questions

How do you run LLM-generated Python code safely?

In a separate process with a timeout, memory and CPU limits, a temporary working folder holding a copy of the data, and no inherited environment variables, so API keys cannot leak. Check the code before it runs too, with an allow-list of imports and a ban on file and system access.

Why does a data agent give wrong totals when the code runs fine?

Usually because the model forgot a business rule, such as removing duplicate rows or excluding refunds. Rules that every answer depends on should be applied in code before the model sees the table, not only described in the prompt.

How do you stop a data agent from making up numbers?

Require a final ANSWER line and check that its number appears in the output of code the agent ran. When it does not, send the answer back and ask the model to compute and print it.

Building a data analysis agent that runs code

LLMs are bad at arithmetic over thousands of rows and good at writing the pandas code that does it. A code-executing data agent puts the two together: the model writes code, a sandbox runs it, and the output goes back until the model can answer. In this lab you build every part of that agent in Python: a subprocess sandbox with timeouts, memory caps, output limits and a clean environment; a data description for the prompt; a check that the final answer was printed by code; an AST allow-list that refuses file and system access; and a prepare step that turns the finance team's rules into code.