airflow.providers.common.ai.utils.toolset_base¶

Behaviour shared by the toolsets this provider ships.

Attributes¶

`log`
`P`
`R`

Classes¶

`AirflowToolset`	A toolset whose tool results are safe to hand to a model.
`MaskingToolset`	Apply the same masking as `AirflowToolset` to any toolset.
`MaskingCapability`	Mask what every tool hands the model, however the tool reached the agent.

Functions¶

`validate_max_retries`(max_retries)	Return `max_retries` unchanged, or raise `ValueError` if it is negative.
`ensure_masked`(toolset)	Return `toolset` wrapped in `MaskingToolset`, unless it already masks its own output.

Module Contents¶

airflow.providers.common.ai.utils.toolset_base.log[source]¶

airflow.providers.common.ai.utils.toolset_base.P[source]¶

airflow.providers.common.ai.utils.toolset_base.R[source]¶

airflow.providers.common.ai.utils.toolset_base.validate_max_retries(max_retries)[source]¶

Return max_retries unchanged, or raise ValueError if it is negative.

class airflow.providers.common.ai.utils.toolset_base.AirflowToolset[source]¶

Bases: pydantic_ai.toolsets.abstract.AbstractToolset[Any]

A toolset whose tool results are safe to hand to a model.

Two methods look alike and have different jobs. execute_tool() is the one a subclass writes: it runs the tool and returns the result as the tool produced it. call_tool() is pydantic-ai’s entry point, implemented here once: it runs execute_tool and passes what it returns, and any exception it raises, through Airflow’s secret masker, so a connection password that ends up in a database error or a hook’s return value is replaced with *** before the model, the model provider or a trace sees it. A subclass that overrides call_tool instead skips that masking, which is why ensure_masked() wraps such a toolset again.

The two signatures differ on purpose. call_tool keeps the positional shape pydantic-ai invokes it with. execute_tool takes ctx and tool keyword-only, so arguments can be added to it later without breaking subclasses.

async call_tool(name, tool_args, ctx, tool)[source]¶

Call a tool with the given arguments.

Args:: name: The name of the tool to call. tool_args: The arguments to pass to the tool. ctx: The run context. tool: The tool definition returned by [get_tools][pydantic_ai.toolsets.AbstractToolset.get_tools] that was called.

replayable: bool = True[source]¶

abstractmethod execute_tool(name, tool_args, *, ctx, tool)[source]¶

Async:

Run tool name with validated tool_args and return its result unmasked.

This is the method a subclass implements; call_tool() runs it and masks what it returns. ctx and tool are keyword-only so that arguments can be added here later without breaking subclasses.

airflow_tools()[source]¶

Return this toolset’s tools for an agent framework other than Pydantic AI.

See airflow.providers.common.ai.tools for the adapters that take them.

static run_blocking(fn, /, *args, **kwargs)[source]¶

Async:

Run a blocking hook call in a worker thread, keeping the event loop free.

Calls made through this method are serialized across the process.

class airflow.providers.common.ai.utils.toolset_base.MaskingToolset[source]¶

Bases: pydantic_ai.toolsets.wrapper.WrapperToolset[Any]

Apply the same masking as AirflowToolset to any toolset.

AgentOperator wraps every toolset it runs in one, so a toolset the Dag author wrote gets masked output too.

async call_tool(name, tool_args, ctx, tool)[source]¶

Call a tool with the given arguments.

Args:: name: The name of the tool to call. tool_args: The arguments to pass to the tool. ctx: The run context. tool: The tool definition returned by [get_tools][pydantic_ai.toolsets.AbstractToolset.get_tools] that was called.

class airflow.providers.common.ai.utils.toolset_base.MaskingCapability[source]¶

Bases: pydantic_ai.capabilities.AbstractCapability[Any]

Mask what every tool hands the model, however the tool reached the agent.

MaskingToolset only covers a toolset the operator can find and wrap. A capability can supply tools the operator never sees as a toolset: MCP, a Toolset built per run, or a Toolset inside PrefixTools or a CombinedCapability. This masks around every tool call instead.

create_agent() adds it to every agent it builds, so add it yourself only to an agent you build with Agent(...), last in its capabilities:

agent = Agent(model, toolsets=[my_toolset], capabilities=[*my_capabilities, MaskingCapability()])

It declares the innermost position, and pydantic-ai breaks ties between capabilities that declare it by list order, the last one innermost, so the agent’s other capabilities hooking tool execution see a result only after it is masked. Two kinds still see it first: a capability passed to run(capabilities=...), which wraps inside the agent’s own, and a toolset wrapper a capability adds with get_wrapper_toolset, which runs below every capability’s hooks.

get_ordering()[source]¶

Join the innermost capabilities; list order decides which of them is closest to the tool.

async wrap_tool_execute(ctx, *, call, tool_def, args, handler)[source]¶

Run the tool and mask its result, or the error it raised, before anything else sees it.

airflow.providers.common.ai.utils.toolset_base.ensure_masked(toolset)[source]¶

Return toolset wrapped in MaskingToolset, unless it already masks its own output.

A function that builds a toolset for each run, which pydantic-ai also accepts, is wrapped too. So is an AirflowToolset whose call_tool is overridden, since the override can bypass the masking.