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='')[source]¶

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

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").

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 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.