Code mode

Note

Experimental: this can change or be removed in a minor release of this provider. See Stable and experimental features.

Pass pydantic-ai-harness’s CodeMode capability in capabilities= (see Capabilities and guardrails) to collapse the agent’s tools into a single run_code tool powered by the Monty sandbox. Instead of one model round-trip per tool call, the model writes a single Python snippet that calls the tools as functions – with loops, conditionals, and asyncio.gather – in one turn. For multi-tool workflows this cuts round-trips and token use.

The generated code runs in Monty’s deny-by-default sandbox: it cannot read the filesystem, the network, or environment variables. It can only call the tools you registered. Code mode therefore does not widen what the agent can reach – the tools it calls still run in the worker – it only changes how the model invokes them. See Securing agent tools for the tool boundary.

When to use it

Code mode pays off for orchestration-heavy, computation-light workflows: calling several tools, looping over their results, filtering, and combining them. Collapsing many sequential tool calls into one turn is where the round-trip and token savings come from – the example below answers a per-customer question in a single run_code block instead of one model round-trip per customer.

It is not a general-purpose code runtime. The generated code is only the glue between tool calls; every real capability must come from a tool. Monty runs a subset of Python and cannot import third-party libraries (pandas, numpy, requests, boto3, …) and has no filesystem or network access. If a task needs to crunch data inline with a library, you have two options, both better than code mode:

  • Push the work into a tool. Do the aggregation in SQL (SQLToolset), or expose a hook method that returns the processed result (HookToolset). The tool runs in the full worker environment with all its dependencies, and code mode just orchestrates it.

  • Give the agent a real environment, with SandboxToolset. It hands the model a shell and a filesystem in a disposable sandbox off the worker, so third-party packages, a real interpreter and installed binaries are all available. It costs a sandbox per run and a second or so to provision, against well under a millisecond for Monty, so reach for it when inline library code is genuinely required rather than by default. See Sandboxed execution for agents for the backends and their limitations.

The two are not exclusive: CodeMode and a SandboxToolset can be enabled together, and the file tools fold into run_code while run_command stays a tool of its own.

Requires the code-mode extra:

pip install "apache-airflow-providers-common-ai[code-mode]"

and the import from pydantic_ai_harness import CodeMode in the Dag file.

airflow/providers/common/ai/example_dags/example_agent.py[source]

if SQLToolset is not None and CodeMode is not None:

    @dag(tags=["example"])
    def example_agent_operator_code_mode():
        AgentOperator(
            task_id="code_mode_analyst",
            prompt="For the top 3 customers by order count, what was each one's total spend?",
            llm_conn_id="pydanticai_default",
            system_prompt="You are a SQL analyst. Write Python that calls the tools to answer.",
            toolsets=[SQLToolset(db_conn_id="postgres_default", allowed_tables=["customers", "orders"])],
            # Requires the `code-mode` extra:
            #   pip install "apache-airflow-providers-common-ai[code-mode]"
            capabilities=[CodeMode(max_tool_calls=200)],
        )

CodeMode takes its own arguments; see the pydantic-ai-harness code mode docs. The one you are most likely to need is max_tool_calls: each run_code snippet may make at most that many nested tool calls, 100 by default. A snippet that asks for more fails, the calls it already made are not undone, and the model has to split the work across several snippets. If the agent fans out over a list, for example one tool call per row or per file, set it above the largest list you expect.

Code mode cannot be combined with durable=True (see Capabilities and guardrails for which capabilities durable replay covers), and it turns off Approve an agent’s tool calls. A tool that needs approval and is called from run_code does not run: the model gets an error back and may carry on without it, so the task can still succeed. To keep such a tool out of run_code, list the other tools in CodeMode(tools=[...]); called directly, it fails the task.

Importing CodeMode in the Dag file loads pydantic-ai-harness and Monty each time the file is parsed. Keep code mode agents in their own Dag file if the rest of the file does not need them.

Note

pydantic-ai-harness is pre-1.0. The code-mode extra is opt-in so its dependency churn never affects the base provider install.

Was this entry helpful?