airflow.providers.common.ai.toolsets.mcp¶

MCP server toolset that resolves configuration from an Airflow connection.

Classes¶

MCPToolset

Toolset that connects to an MCP server configured via an Airflow connection.

Module Contents¶

class airflow.providers.common.ai.toolsets.mcp.MCPToolset(mcp_conn_id, *, tool_prefix=None)[source]¶

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

Toolset that connects to an MCP server configured via an Airflow connection.

Reads MCP server transport type, URL, command, and credentials from the connection via MCPHook and creates the appropriate PydanticAI MCP server instance. All AbstractToolset methods delegate to the underlying MCP server.

This is the recommended way to use MCP servers in Airflow — it stores server configuration in Airflow connections (and secret backends) rather than hard-coding URLs and credentials in DAG code.

If you prefer full PydanticAI control, you can pass MCP server instances directly to AgentOperator(toolsets=[...]), since MCPServerStreamableHTTP, MCPServerSSE, and MCPServerStdio all implement AbstractToolset.

Parameters:

mcp_conn_id (str) – Airflow connection ID for the MCP server.
tool_prefix (str | None) – Optional prefix prepended to tool names (e.g. "weather" → "weather_get_forecast").

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.

async __aenter__()[source]¶

Enter the toolset context.

This is where you can set up network connections in a concrete implementation.

async __aexit__(*args)[source]¶

Exit the toolset context.

This is where you can tear down network connections in a concrete implementation.

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.