airflow.providers.common.ai.toolsets.hook

Generic adapter that exposes Airflow Hook methods as pydantic-ai tools.

Classes

HookToolset

Expose selected methods of an Airflow Hook as pydantic-ai tools.

Module Contents

class airflow.providers.common.ai.toolsets.hook.HookToolset(hook, *, allowed_methods, tool_name_prefix='', pinned_arguments=None, max_retries=None)[source]

Bases: airflow.providers.common.ai.utils.toolset_base.AirflowToolset

Expose selected methods of an Airflow Hook as pydantic-ai tools.

This adapter introspects the method signatures and docstrings of the given hook to build ToolDefinition objects that an LLM agent can call.

Parameters:

  • hook (airflow.providers.common.compat.sdk.BaseHook) – An instantiated Airflow Hook. Its connection ID – the attribute the hook’s conn_name_attr names, such as postgres_conn_id – is templated when the toolset is passed to AgentOperator / @task.agent, so HookToolset(PostgresHook(postgres_conn_id="tenant_{{ ... }}"), ...) reaches a different database per task instance. The hook in the Dag file is not modified; each task instance gets a copy.

  • allowed_methods (list[str]) – Method names to expose as tools. Required — auto-discovery is intentionally not supported for safety.

  • tool_name_prefix (str) – Optional prefix prepended to each tool name (e.g. "s3_" → "s3_list_keys").

  • pinned_arguments (dict[str, Any] | None) – Experimental. Arguments the Dag author fixes, such as the bucket a storage hook may use: {"bucket_name": "reports"}. Each is left out of the arguments the model sees, refused if the model supplies it anyway, and passed to every allowed method as it is written here, not rendered as a template. Every allowed method must take each pinned argument as a named parameter: one that does not, such as a method taking bucket or only **kwargs, raises ValueError, because the model could still choose the value through it. Expose such a method from a second HookToolset.

  • max_retries (int | None) – How many times the model may correct a call with invalid arguments, or one that supplies a pinned argument, before the run fails. An exception from the hook itself fails the run straight away. None (the default) uses the agent’s tool retry budget, its retries, as pydantic-ai’s own toolsets do.

agent_template_fields: collections.abc.Sequence[str] = ('conn_id',)[source]
property conn_id: str | None[source]

The hook’s connection ID, or None when the hook keeps it under neither attribute.

property id: str[source]

An ID for the toolset that is unique among all toolsets registered with the same agent.

If you’re implementing a concrete implementation that users can instantiate more than once, you should let them optionally pass a custom ID to the constructor and return that here.

A toolset needs to have an ID in order to be used in a durable execution environment like Temporal, in which case the ID will be used to identify the toolset’s activities within the workflow.

IDs wrapped in angle brackets (‘<agent>’ for an agent’s own function toolset, ‘<output>’ for its output tools) name a role the framework fills on the user’s behalf rather than a registered toolset. Don’t return one from your own toolset.

async get_tools(ctx)[source]

The tools that are available in this toolset.

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

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.

Was this entry helpful?