How to give a Pydantic AI agent a code execution tool

Pydantic AI's distinguishing idea is that an agent is a typed object: typed dependencies in, a validated result out, tools whose schemas come from real Python signatures. That structure happens to be an unusually good fit for a code execution tool, because the thing a code tool needs most is a well-managed handle to a sandbox — and dependency injection is exactly the mechanism for that.

Below is the whole setup: the dependency, the tool, the retry path, and the lifecycle. I build PandaStack, so that's the sandbox in the examples; the structure is the same with any provider.

Put the sandbox in the dependencies

The temptation is a module-level sandbox that every tool call reaches for. Don't — it means every concurrent conversation shares one kernel, and one user's variables become another's. Pydantic AI's `deps_type` exists to avoid precisely this.

# pip install pandastack pydantic-ai
from dataclasses import dataclass
from pydantic_ai import Agent, RunContext, ModelRetry
from pandastack import Sandbox, CodeContext

@dataclass
class Deps:
    sandbox: Sandbox
    kernel: CodeContext

agent = Agent(
    "openai:gpt-4o",
    deps_type=Deps,
    system_prompt=(
        "You are a data analyst. Answer questions by writing and running "
        "Python — never compute a number yourself. State your assumptions, "
        "then compute. Always print() the values you want to report."
    ),
)

The tool

A tool is a decorated function whose first parameter is the `RunContext`. Everything else — the parameter names, the type hints, the docstring — becomes the schema the model sees, which is why the docstring is worth writing carefully.

@agent.tool
def run_python(ctx: RunContext[Deps], code: str) -> str:
    """Execute Python in a secure sandbox and return its printed output.

    State persists between calls, so you can load data in one step and
    analyse it in the next. pandas, numpy and matplotlib are available.
    Print anything you want to see — return values are not captured.

    Args:
        code: The Python source to execute.
    """
    result = ctx.deps.kernel.run_code(code)

    if result.error:
        # Hand the traceback back to the model as a retry rather than
        # failing the run — models fix their own NameErrors reliably.
        raise ModelRetry(f"The code raised an error:\n{result.error}")

    output = result.stdout.strip()
    if not output:
        return "(the code ran successfully but printed nothing — use print())"
    if result.stderr.strip():
        output += f"\n\nstderr:\n{result.stderr.strip()}"
    return output

`ModelRetry` is the piece worth pausing on. It's Pydantic AI's mechanism for telling the model "that didn't work, here's why, try again" — the message goes back into the conversation and the model gets another turn. For code execution that's exactly right: a traceback is the most actionable feedback a model can receive, and models correct their own syntax and name errors reliably.

Running it

import asyncio

async def analyse(question: str, csv_path: str) -> str:
    sandbox = Sandbox.create(template="code-interpreter", ttl_seconds=1800)
    try:
        sandbox.filesystem.upload(csv_path, "/workspace/data.csv")
        kernel = sandbox.create_code_context(language="python")
        deps = Deps(sandbox=sandbox, kernel=kernel)

        result = await agent.run(
            f"{question} The data is at /workspace/data.csv.",
            deps=deps,
        )
        return result.output
    finally:
        sandbox.kill()

print(asyncio.run(analyse("Which region grew fastest year over year?", "./sales.csv")))

The `try/finally` is not decoration. A sandbox that outlives its conversation is a resource leak that costs money, and the TTL is a backstop for the case where your process dies before the `finally` runs — not a substitute for cleaning up.

Typed output, which is the actual point of Pydantic AI

The reason to use this framework over a raw SDK is that the agent's answer can be a validated model rather than a paragraph. Combined with a code tool, you get something quite strong: the numbers are computed by an interpreter, and the structure is enforced by pydantic.

from pydantic import BaseModel, Field

class Analysis(BaseModel):
    answer: str = Field(description="One-sentence answer to the question.")
    value: float = Field(description="The computed figure.")
    unit: str = Field(description="Unit or currency of the figure.")
    code_used: str = Field(description="The Python that produced the figure.")

analyst = Agent(
    "openai:gpt-4o",
    deps_type=Deps,
    output_type=Analysis,
    retries=3,
    system_prompt=(
        "Compute every figure with run_python. Put the exact code that "
        "produced your final number in code_used, so it can be re-run."
    ),
)

That `code_used` field is a small idea with a large payoff: it makes the answer auditable. Someone can paste the code into the same sandbox and check that it produces the number the agent reported, which is a much stronger guarantee than "the model said so" and costs one extra field.

More than one tool over the same sandbox

Once the sandbox is in deps, additional tools are almost free — and a couple of small ones make the code tool substantially more effective. The most valuable is a file listing, because a model that can see what is on disk stops guessing at filenames:

@agent.tool
def list_workspace(ctx: RunContext[Deps]) -> str:
    """List the files available in /workspace, with sizes."""
    result = ctx.deps.sandbox.exec("ls -la /workspace")
    return result.stdout


@agent.tool
def read_file(ctx: RunContext[Deps], path: str, max_bytes: int = 4000) -> str:
    """Read the beginning of a text file in the sandbox.

    Use this to inspect a file's structure before writing code against it.
    """
    if not path.startswith("/workspace/"):
        raise ModelRetry("Only paths under /workspace can be read.")
    data = ctx.deps.sandbox.filesystem.read(path)
    return data[:max_bytes].decode("utf-8", errors="replace")

Note the path check raising `ModelRetry` rather than returning an error string: the model gets told why the call was rejected and can correct it, and the constraint is enforced in your code rather than hoped for in the prompt. That pattern — validate in the tool, explain in the retry — is the one to reach for whenever a tool has rules.

Resist adding much more than this. Every tool widens the schema the model has to reason over, and a code tool already subsumes most of what extra tools would do — the agent can shell out, parse, and transform inside the sandbox without you designing an interface for it.

Why the sandbox has to be a real boundary

The code being executed was written by a language model, working from data you may not control. That's the whole argument. `exec()` in your process means model-generated code runs with your process's environment variables, database connections and network access — and a model doesn't have to be attacked to do damage, it just has to be wrong about a filename.

A microVM sandbox puts a hardware virtualisation boundary in the way: separate kernel, separate filesystem, separate network namespace, and a TTL that reaps it whether or not your code remembers to. `subprocess` with a timeout is not the same thing — it shares the kernel, the filesystem and the credentials.

In a web application

For anything longer-lived than a script, tie the sandbox to the session rather than to a single request, so the kernel's state survives between turns of a conversation:

from fastapi import FastAPI

app = FastAPI()
_sessions: dict[str, Deps] = {}

def deps_for(session_id: str) -> Deps:
    if session_id not in _sessions:
        sb = Sandbox.create(
            template="code-interpreter",
            ttl_seconds=1800,
            metadata={"session": session_id},
        )
        _sessions[session_id] = Deps(sandbox=sb, kernel=sb.create_code_context())
    return _sessions[session_id]

@app.post("/ask")
async def ask(session_id: str, question: str):
    result = await agent.run(question, deps=deps_for(session_id))
    return {"answer": result.output}

In production, replace the dictionary with something that survives a restart and expires entries — the sandbox TTL will clean up the far side regardless, but a stale handle in your process will produce confusing 404s until it's evicted.

Recap

1. Put the sandbox and kernel in deps_type, never in a module-level global.
2. Write the docstring for the model — it becomes the tool schema.
3. Raise ModelRetry on execution errors so the traceback becomes the next attempt.
4. Bound retries so an impossible task can't loop forever.
5. Use output_type to get a validated result, and include the code that produced the number.
6. Clean up in a finally block, with a TTL as the backstop.