How to give a LlamaIndex agent a code execution tool

LlamaIndex is built around retrieval, and retrieval agents have a characteristic failure: they find the right numbers and then get the arithmetic wrong. Ask one for the quarter-over-quarter growth rate across four retrieved documents and it will confidently produce a number that isn't right, because it computed it by generating tokens rather than by doing maths.

A code execution tool fixes that properly. The model writes Python, something else runs it, and the answer comes back from an interpreter instead of from a language model's impression of one. LlamaIndex has a clean extension point for this — you need somewhere safe to run the code, which is what I build PandaStack for, though the tool shape below works with any sandbox provider.

Why not just exec() it

Because the code is written by a language model, and the model is reading documents that may not be yours. A retrieval-augmented agent takes untrusted text as input and generates code as output — that's the exact shape of a prompt-injection-to-code-execution chain, and `exec()` inside your API process means anything that survives the chain runs with your process's credentials and network access.

You don't need to believe an attack is likely to want the boundary. A model that writes `os.remove` when it meant `os.path.exists` does the same damage as a hostile one.

The tool

LlamaIndex tools are plain callables wrapped in a FunctionTool. The docstring and type hints become the schema the model sees, so they're not documentation — they're prompt.

# pip install pandastack llama-index
from llama_index.core.tools import FunctionTool
from pandastack import Sandbox

# One sandbox, one kernel, for the life of the agent.
sandbox = Sandbox.create(template="code-interpreter", ttl_seconds=1800)
kernel = sandbox.create_code_context(language="python")

def run_python(code: str) -> str:
    """Execute Python in a secure sandbox and return its output.

    Variables, imports and loaded data persist between calls, so you can
    build up an analysis across several steps. numpy, pandas and matplotlib
    are available. Always print() the values you want to see.
    """
    result = kernel.run_code(code)
    if result.error:
        return f"ERROR:\n{result.error}"
    out = result.stdout.strip() or "(no output — did you print() the result?)"
    if result.stderr.strip():
        out += f"\n\nstderr:\n{result.stderr.strip()}"
    return out

code_tool = FunctionTool.from_defaults(
    fn=run_python,
    name="run_python",
    description=(
        "Execute Python code in a persistent sandbox. Use this for any "
        "calculation, aggregation, or data transformation. Never do "
        "arithmetic yourself — compute it here."
    ),
)

Two details in there matter more than they look. `create_code_context` gives you a persistent kernel rather than a fresh process per call — so the agent can load a dataframe in one step and query it in the next, which is how humans actually work with data. And the return value is a string, deliberately: the model reads it, so it should be readable rather than a serialised object.

Wiring it into an agent

from llama_index.core.agent.workflow import FunctionAgent
from llama_index.llms.openai import OpenAI

agent = FunctionAgent(
    tools=[code_tool, retriever_tool],   # your existing query engine tool
    llm=OpenAI(model="gpt-4o"),
    system_prompt=(
        "You answer questions about the indexed documents. "
        "Retrieve facts with the query tool. Do every calculation with "
        "run_python — never compute a number in your head. "
        "State the numbers you used before computing with them."
    ),
)

response = await agent.run(
    "What was the compound growth rate across the four quarters in the report?"
)

That system prompt is doing real work. Without an explicit instruction, models will happily answer a numeric question directly when they think it's easy — and "easy" is exactly where the plausible-but-wrong answers live.

Getting retrieved data into the kernel

The interesting pattern in a LlamaIndex agent isn't the code tool on its own — it's the handoff from retrieval to computation. Two approaches, and the second is much better as data grows.

The naive version has the model paste retrieved numbers into the code it writes. Works for a handful of values, and every value passes through the model's tokeniser on the way, which is precisely where transcription errors come from.

The better version puts the data in the sandbox directly and lets the code read it. The model then writes code that references a file it never had to reproduce:

import json

nodes = retriever.retrieve("quarterly revenue")
payload = [{"text": n.text, "metadata": n.metadata} for n in nodes]

sandbox.filesystem.write(
    "/workspace/retrieved.json",
    json.dumps(payload).encode(),
)

# The agent can now be told, in the system prompt:
#   "Retrieved context is written to /workspace/retrieved.json. Load it
#    with json.load() rather than retyping values."

This also means large result sets never enter the context window. A hundred retrieved chunks are a file the code reads, not a hundred chunks of prompt you're paying for on every subsequent turn.

Charts and rich output

A code-interpreter sandbox captures rich results, not just text — a matplotlib figure comes back as an image rather than as the string `<Figure size 640x480>`:

result = kernel.run_code("""
import matplotlib.pyplot as plt
plt.plot([1, 2, 3], [4, 9, 16])
plt.title("Growth")
plt.show()
""")

if result.png:
    with open("chart.png", "wb") as f:
        import base64
        f.write(base64.b64decode(result.png))

Keep the image out of the tool's return string. Return "chart saved" to the model and handle the bytes in your application — sending a base64 PNG back into the context window burns thousands of tokens for something the model can't see anyway.

Sandbox lifecycle

One sandbox per conversation is the right granularity, and it's worth being deliberate about it. A single global sandbox means two users' analyses share a kernel, which is a data leak with extra steps. A fresh sandbox per tool call throws away the persistent state that makes the tool useful.

def build_agent_for_session(session_id: str):
    sandbox = Sandbox.create(
        template="code-interpreter",
        ttl_seconds=1800,
        metadata={"session": session_id},
    )
    kernel = sandbox.create_code_context(language="python")
    ...
    return agent, sandbox

# When the conversation ends
sandbox.kill()

The TTL is the safety net for the case where your cleanup code doesn't run — a crashed worker, a user who closes the tab. Set it to something slightly longer than a plausible conversation and let the platform reap the rest.

Recap

1. Wrap a sandbox call in a FunctionTool; the docstring is prompt, so write it for the model.
2. Use a persistent kernel so state carries across tool calls.
3. Tell the agent explicitly, in the system prompt, never to compute numbers itself.
4. Write retrieved data into the sandbox as a file rather than routing it through the model.
5. Return stdout and errors as readable strings; keep images out of the context window.
6. One sandbox per session, with a TTL as the backstop, and kill it when the conversation ends.