Writing your own toolset¶

Note

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

When an agent needs a system that no shipped toolset covers, write a toolset for it by subclassing AirflowToolset, the class every toolset in this provider is built on. The subclass gets what the shipped toolsets have: its results and errors pass through Airflow’s secret masker even outside AgentOperator, its tools work in other agent frameworks through airflow_tools(), and its blocking hook calls run off the event loop. AgentOperator masks the output of every toolset it runs, a FunctionToolset included. If a method of an existing hook already does what the agent needs, pass the hook to HookToolset instead. If the tool needs no connection, a plain function in pydantic-ai’s FunctionToolset is less code.

This toolset looks packages up on PyPI through an HTTP connection. A project’s response from PyPI’s JSON API lists every release of the project, so the toolset returns only the fields the model needs:

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

class PackageArgs(TypedDict):
    """The arguments of get_package."""

    # A pattern keeps the model's input to a valid name, so it cannot reach another path on the host.
    name: Annotated[
        str,
        Field(
            description="The distribution name, such as apache-airflow.",
            pattern=r"^[A-Za-z0-9]([A-Za-z0-9._-]*[A-Za-z0-9])?$",
        ),
    ]


class ReleaseArgs(PackageArgs):
    """The arguments of get_release."""

    version: Annotated[
        str, Field(description="An exact version, such as 3.1.0.", pattern=r"^[A-Za-z0-9][A-Za-z0-9.!+_-]*$")
    ]


# Each tool's description, and the adapter that gives its arguments' JSON schema and validator.
TOOLS: dict[str, tuple[str, TypeAdapter[Any]]] = {
    "get_package": (
        "Look up a package on PyPI: its latest version, summary and supported Python versions.",
        TypeAdapter(PackageArgs),
    ),
    "get_release": (
        "Look up one release of a package on PyPI: what it depends on and when it was uploaded.",
        TypeAdapter(ReleaseArgs),
    ),
}


class PyPIToolset(AirflowToolset):
    """Look packages up on PyPI, or a mirror of its JSON API, through an HTTP connection."""

    # AgentOperator renders these attributes with Jinja for each task instance.
    agent_template_fields: Sequence[str] = ("http_conn_id",)

    def __init__(self, http_conn_id: str = "pypi") -> None:
        self.http_conn_id = http_conn_id

    @property
    def id(self) -> str:
        return f"pypi-{self.http_conn_id}"

    async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]:
        return {
            name: ToolsetTool(
                toolset=self,
                tool_def=ToolDefinition(
                    name=name, description=description, parameters_json_schema=args.json_schema()
                ),
                max_retries=ctx.max_retries,
                args_validator=SchemaValidator(args.core_schema),
            )
            for name, (description, args) in TOOLS.items()
        }

    async def execute_tool(
        self,
        name: str,
        tool_args: dict[str, Any],
        *,
        ctx: RunContext[Any],
        tool: ToolsetTool[Any],
    ) -> dict[str, Any]:
        if name == "get_package":
            path = tool_args["name"]
        elif name == "get_release":
            path = f"{tool_args['name']}/{tool_args['version']}"
        else:
            raise ValueError(f"Unknown tool: {name!r}")

        hook = HttpHook(method="GET", http_conn_id=self.http_conn_id)
        # HttpHook.run blocks. run_blocking runs it in a thread, one hook call at a time, so
        # give it a timeout: a stalled request would hold up every other blocking call.
        response = await self.run_blocking(
            hook.run, f"pypi/{path}/json", extra_options={"check_response": False, "timeout": 30}
        )
        if response.status_code == 404:
            # The model named something that does not exist; it reads this and can try again.
            raise ModelRetry(f"PyPI has no {path!r}. Check the name and version.")
        # Any other failure is not the model's to fix, so it fails the task.
        response.raise_for_status()

        # The full document lists every file of every release. Return what the model needs.
        document = response.json()
        info = document["info"]
        if name == "get_package":
            return {key: info[key] for key in ("name", "version", "summary", "requires_python")}
        return {
            "version": info["version"],
            "requires_python": info["requires_python"],
            "requires_dist": info["requires_dist"],
            "yanked": info["yanked"],
            "uploaded": document["urls"][0]["upload_time_iso_8601"] if document["urls"] else None,
        }

Pass it to an agent like any other toolset:

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

@dag(tags=["example"])
def example_custom_toolset():
    @task.agent(
        llm_conn_id="pydanticai_default",
        system_prompt="You answer questions about Python packages. Look them up; do not guess.",
        toolsets=[PyPIToolset(http_conn_id="pypi")],
    )
    def check_package(package: str) -> str:
        return f"What is the latest release of {package}, and which Python versions does it support?"

    check_package("apache-airflow-providers-common-ai")

The example needs apache-airflow-providers-http, a model connection named pydanticai_default, and an HTTP connection named pypi whose host is https://pypi.org, for example:

export AIRFLOW_CONN_PYPI='{"conn_type": "http", "host": "https://pypi.org"}'

When the model calls get_release for version 0.10.0 of apache-airflow-providers-common-ai, it receives the following, with requires_dist cut from 32 entries to three here:

{"version": "0.10.0", "requires_python": ">=3.10",
 "requires_dist": ["apache-airflow-providers-common-compat>=1.15.0",
                   "apache-airflow-providers-standard>=1.20.0",
                   "apache-airflow>=3.0.0"],
 "yanked": false, "uploaded": "2026-09-29T06:39:53.490917Z"}

What a subclass implements¶

id: A name for the toolset. With durable=True, each call is recorded under it, and a retry finds the calls an earlier attempt recorded by it. If the id changes between attempts, the retry stops matching the earlier attempt at the first call to this toolset, and that call and every step after it, model requests included, run again. Deriving it from the connection ID, as the example does, gives each connection’s toolset its own.
get_tools: The tools the model is offered, keyed by name. Each ToolDefinition carries the description and the JSON schema of the arguments, which are all the model knows about the tool. args_validator checks the model’s arguments before execute_tool runs, and arguments that fail it go back to the model to correct. The example’s pattern constraints keep a name or version from changing the URL the hook calls. A pydantic TypeAdapter over a TypedDict gives the schema, and a pydantic-core SchemaValidator built from its core_schema is the validator. Import TypedDict from typing_extensions. On Python before 3.12, building a TypeAdapter over a typing.TypedDict raises PydanticUserError: Please use `typing_extensions.TypedDict` instead of `typing.TypedDict` on Python < 3.12., and the Dag file fails to import.
execute_tool: Runs one call. name says which tool, and tool_args holds the arguments args_validator returned. Return the result as the tool produced it: the base class’s call_tool masks it on the way out. Implement execute_tool rather than overriding call_tool, since an override skips that masking and AgentOperator then wraps the toolset in a masking layer of its own. With durable=True a result has to be JSON-serializable, or the task fails with DurableJournalError and is not retried.

Two class attributes are optional:

agent_template_fields: The attributes AgentOperator renders with Jinja for each task instance, so PyPIToolset(http_conn_id="pypi_{{ var.value.environment }}") picks a connection per environment. Rendering sets the attributes on a copy of the toolset and does not run __init__ again, so derive nothing from them in __init__: build hooks inside execute_tool, as the example does. A hook built in __init__ keeps the connection ID as written in the Dag file. Only AgentOperator and @task.agent render these attributes.
replayable: True by default: with durable=True, a retry reuses the result an earlier attempt recorded instead of calling the system again. Set it to False when every attempt must make the call, as ManagedAgentToolset does because a managed agent acts on systems Airflow cannot observe.

Errors the model sees and errors that fail the task¶

What execute_tool raises decides who deals with it:

ModelRetry(message) sends the message to the model, which can call the tool again. Each one counts against the tool’s max_retries, which in the example is ctx.max_retries, the agent’s retries. Once they are used up, the run fails with UnexpectedModelBehavior, whose message starts Tool 'get_package' exceeded max retries count of 1 under the default retries. The example raises ModelRetry for a package that does not exist, and the model reads:
```
PyPI has no 'apache-airflow-providers-commn-ai'. Check the name and version.

Fix the errors and try again.
```
ToolFailed(message), from pydantic_ai.exceptions, tells the model the call failed for good, without counting against max_retries. Use it when calling again with other arguments would not help, such as a record the connection is not allowed to read.
Any other exception fails the agent run, so the task fails and its Airflow retries apply. The full exception, with its cause, is logged to the task log. What leaves the toolset usually keeps its type, but its message is masked and its cause is dropped.

Blocking calls¶

execute_tool runs on the agent’s event loop, and most hook methods block. Call them through self.run_blocking(fn, *args, **kwargs), which runs the call in a worker thread. Calls made through it run one at a time across the whole process, because agent frameworks run tool calls concurrently and hooks are not thread-safe in general. One slow call therefore holds up the blocking calls of every other toolset in the same worker process.

Using the toolset in another agent framework¶

The AirflowTools adapters for Strands, ADK, the Anthropic tool runner and the OpenAI Agents SDK take the toolset as they take a shipped one, for example AirflowTools(PyPIToolset()) (see Agent frameworks). They read its tools through airflow_tools(), which returns them as framework-neutral AirflowTool objects. Outside a pydantic-ai run there is no real RunContext: ctx.deps is None and ctx.model is a stand-in, and only ctx.retry and ctx.max_retries are kept per tool, with max_retries set to 1. A toolset meant for other frameworks should read nothing else from ctx.