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.

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.

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.

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.