This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch adapters-wrapped-toolsets in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 6719e942236ac939e6c47d2c3467882228f69edb Author: Kaxil Naik <[email protected]> AuthorDate: Thu Nov 5 21:48:22 2026 +0000 Accept .filtered(), .prefixed() and .renamed() toolsets in the agent framework adapters AirflowTools for the OpenAI Agents SDK, Strands, ADK, the Anthropic tool runner and the Claude Agent SDK all build their tools with collect_tools(), which called source.airflow_tools(). A toolset wrapped with .filtered(), .prefixed() or .renamed() has no such method, so AirflowTools(SQLToolset(...).filtered(...)) failed with AttributeError. collect_tools() now reads through those three wrappers, nested in any order. The tools still come from the innermost toolset's own airflow_tools(); the outermost wrapper's get_tools(), given the same inert run context airflow_tools_from_toolset() uses, decides which of them are offered and under what name, so a filter runs once, when the adapter is built. Any other wrapper, such as .approval_required(), raises a TypeError that names it and says what to use instead. The adapters take a ToolSource type that lists the accepted wrappers. --- providers/common/ai/docs/frameworks/index.rst | 3 +- .../common/ai/docs/frameworks/openai_agents.rst | 28 ++- providers/common/ai/docs/toolsets/index.rst | 4 + .../airflow/providers/common/ai/tools/__init__.py | 41 +++-- .../providers/common/ai/tools/_from_toolset.py | 99 ++++++++++- .../src/airflow/providers/common/ai/tools/adk.py | 9 +- .../airflow/providers/common/ai/tools/anthropic.py | 11 +- .../providers/common/ai/tools/claude_agent_sdk.py | 9 +- .../providers/common/ai/tools/openai_agents.py | 9 +- .../airflow/providers/common/ai/tools/strands.py | 9 +- .../unit/common/ai/tools/test_openai_agents.py | 12 +- .../ai/tests/unit/common/ai/tools/test_tools.py | 193 ++++++++++++++++++++- 12 files changed, 380 insertions(+), 47 deletions(-) diff --git a/providers/common/ai/docs/frameworks/index.rst b/providers/common/ai/docs/frameworks/index.rst index 661e3ca98d1..cb6e86baf75 100644 --- a/providers/common/ai/docs/frameworks/index.rst +++ b/providers/common/ai/docs/frameworks/index.rst @@ -227,7 +227,8 @@ An adapter maps these onto the framework's own tool type and error status, and m sure a ``ToolCallError`` ends the run rather than reaching the model, as the Strands plugin does. :func:`~airflow.providers.common.ai.tools.collect_tools` turns the toolsets and tools an adapter is given into one list, as the Strands, ADK, Anthropic, OpenAI Agents and Claude Agent SDK -adapters do. +adapters do. It accepts toolsets wrapped with ``.filtered()``, ``.prefixed()`` and +``.renamed()``, and raises ``TypeError`` for any other wrapper. Blocking hook calls, such as a SQL query, run one at a time in the task's process, so an agent that calls two database tools at once gets its answers one after the other. diff --git a/providers/common/ai/docs/frameworks/openai_agents.rst b/providers/common/ai/docs/frameworks/openai_agents.rst index 8b55cdd03c0..11662a0daa8 100644 --- a/providers/common/ai/docs/frameworks/openai_agents.rst +++ b/providers/common/ai/docs/frameworks/openai_agents.rst @@ -244,18 +244,28 @@ it: ``list_keys`` and ``read_key`` and the bucket pinned. The same toolset works in ``AirflowTools(...)``. -``.filtered()`` and ``.approval_required()``, which pydantic-ai toolsets offer, do not -reach the SDK through ``AirflowTools``. They return wrappers that ``AirflowTools`` cannot -read, and building it fails: +To offer the model only some of a toolset's tools, wrap the toolset with ``.filtered()`` +before you pass it in. These tools are every ``SQLToolset`` tool except ``query``: + +.. code-block:: python + + tools = AirflowTools(SQLToolset(db_conn_id="warehouse").filtered(lambda ctx, tool: tool.name != "query")) + +``.prefixed()`` and ``.renamed()`` carry over the same way, nested in any order. The filter +runs once, when ``AirflowTools`` is built, with a placeholder run context whose ``deps`` +is ``None``, so a filter that reads an attribute of ``ctx.deps`` fails there. Whatever the +filter decides holds for every run of an agent given these tools, where a Pydantic AI agent +asks it again before each model turn. + +Any other wrapper fails with a ``TypeError`` that names it, such as for +``.approval_required()``: .. code-block:: text - AttributeError: 'FilteredToolset' object has no attribute 'airflow_tools' + TypeError: AirflowTools cannot carry ApprovalRequiredToolset (.approval_required()) to another framework's agent: only a Pydantic AI run pauses for the approval it asks for. Pass the toolset without it. To have a person approve tool calls, run the agent with AgentOperator; a framework's own tool approval, such as needs_approval in the OpenAI Agents SDK, stops the run and nothing in Airflow resumes it. -Choose the tools with ``allowed_tables``, ``allowed_methods`` and ``pinned_arguments`` -instead. There is no per-call approval on this path; see -:ref:`howto/frameworks:openai-agents-limits`. Each toolset's page covers the rest of its -options: :ref:`howto/toolsets`. +:ref:`howto/frameworks:openai-agents-limits` says what ``needs_approval`` does in a task. +Each toolset's page covers the rest of its options: :ref:`howto/toolsets`. Tools ----- @@ -335,7 +345,7 @@ The SDK routes a call by tool name, so two toolsets that define the same name ra .. code-block:: text - ValueError: More than one tool is named check_query, get_schema, list_tables, query. Give the toolsets different tool_prefix values, where they take one, so the model can call each of them. + ValueError: More than one tool is named check_query, get_schema, list_tables, query. Give the toolsets different tool_prefix values, where they take one, or wrap them with .prefixed(), so the model can call each of them. .. _howto/frameworks:openai-agents-failures: diff --git a/providers/common/ai/docs/toolsets/index.rst b/providers/common/ai/docs/toolsets/index.rst index 230b6e9c0fd..5428721ff22 100644 --- a/providers/common/ai/docs/toolsets/index.rst +++ b/providers/common/ai/docs/toolsets/index.rst @@ -75,6 +75,10 @@ own sandbox (:doc:`../tool_approval` lists every limit). ``.filtered()`` drops t from what the model is offered, as the :ref:`MCP guide <howto/toolset:mcp-filtered>` shows. ``.defer_loading()`` hides a toolset's tools from the model until it searches for them, which pays off on a toolset with many tools (see :doc:`progressive_disclosure`). +In an agent written for another framework, which gets the toolset through +``AirflowTools``, ``.filtered()``, ``.prefixed()`` and ``.renamed()`` carry over; +``.approval_required()`` and ``.defer_loading()`` raise a ``TypeError`` (see +:doc:`../frameworks/index`). :doc:`logging` and :doc:`langchain` do not reach a system of their own: one wraps a toolset to log its calls, the other converts a toolset for a LangChain agent. diff --git a/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py b/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py index a7b6e84e4b9..eb4f8af58cd 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py +++ b/providers/common/ai/src/airflow/providers/common/ai/tools/__init__.py @@ -28,7 +28,8 @@ The toolsets this provider ships, such as :class:`~airflow.providers.common.ai.toolsets.sql.SQLToolset` and :class:`~airflow.providers.common.ai.toolsets.hook.HookToolset`, implement :class:`ToolProvider` and expose their tools through -:meth:`ToolProvider.airflow_tools`. +:meth:`ToolProvider.airflow_tools`. Such a toolset wrapped with ``.filtered()``, +``.prefixed()`` or ``.renamed()`` is accepted too. """ from __future__ import annotations @@ -37,7 +38,11 @@ import logging import time from collections.abc import Awaitable, Callable, Iterable, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Protocol +from typing import TYPE_CHECKING, Any, Protocol, TypeAlias + +from pydantic_ai.toolsets.filtered import FilteredToolset +from pydantic_ai.toolsets.prefixed import PrefixedToolset +from pydantic_ai.toolsets.renamed import RenamedToolset from airflow.providers.common.ai.exceptions import MCPToolsetNotSupportedError from airflow.providers.common.ai.utils.masking import mask_secrets @@ -48,7 +53,7 @@ if TYPE_CHECKING: log = logging.getLogger(__name__) -__all__ = ["AirflowTool", "ToolCallError", "ToolProvider", "ToolResult", "collect_tools"] +__all__ = ["AirflowTool", "ToolCallError", "ToolProvider", "ToolResult", "ToolSource", "collect_tools"] # The first sentence of the error for an MCPToolset handed to an adapter; the second says what # to use instead. @@ -151,35 +156,50 @@ class ToolProvider(Protocol): ... -def collect_tools(sources: Iterable[ToolProvider | AirflowTool]) -> list[AirflowTool]: +ToolSource: TypeAlias = ( + ToolProvider | AirflowTool | FilteredToolset[Any] | PrefixedToolset[Any] | RenamedToolset[Any] +) +"""What a framework adapter accepts: a tool, or a toolset that ``.filtered()``, ``.prefixed()`` and ``.renamed()`` may wrap.""" + + +def collect_tools(sources: Iterable[ToolSource]) -> list[AirflowTool]: """ Flatten toolsets and individual tools, as a framework adapter receives them, into tools. + A toolset wrapped with ``.filtered()``, ``.prefixed()`` or ``.renamed()`` contributes the + tools the wrappers offer, under the names they give them. The ``.filtered()`` function runs + once, here, with an inert run context, so a filter that reads per-run state does not + re-evaluate it before each model turn. + + :raises TypeError: when a source is another wrapper, such as ``.approval_required()``, or + has no ``airflow_tools()`` method. :raises MCPToolsetNotSupportedError: when a source is an ``MCPToolset``. Connect the framework's own MCP client to the server instead. :raises ValueError: when two tools share a name. Frameworks route a call by name, so one of them would silently never run; give the toolsets different ``tool_prefix`` values, - where they take one. + where they take one, or wrap them with ``.prefixed()``. Also when a ``.renamed()`` map + names a tool the toolset does not have. """ return _collect_tools(sources) -def _collect_tools( - sources: Iterable[ToolProvider | AirflowTool], *, mcp_hint: str | None = None -) -> list[AirflowTool]: +def _collect_tools(sources: Iterable[ToolSource], *, mcp_hint: str | None = None) -> list[AirflowTool]: """ :func:`collect_tools` for the adapters this provider ships. :param mcp_hint: How the adapter's framework connects to an MCP server, as a sentence. It replaces the generic advice in the error raised for an ``MCPToolset``. """ + # Circular import: _from_toolset imports AirflowTool from this module. + from airflow.providers.common.ai.tools._from_toolset import airflow_tools_of + tools: list[AirflowTool] = [] for source in sources: if isinstance(source, AirflowTool): tools.append(source) continue try: - tools.extend(source.airflow_tools()) + tools.extend(airflow_tools_of(source)) except MCPToolsetNotSupportedError: if mcp_hint is None: raise @@ -188,6 +208,7 @@ def _collect_tools( if duplicates := sorted({name for name in names if names.count(name) > 1}): raise ValueError( f"More than one tool is named {', '.join(duplicates)}. " - "Give the toolsets different tool_prefix values, where they take one, so the model can call each of them." + "Give the toolsets different tool_prefix values, where they take one, or wrap them with " + ".prefixed(), so the model can call each of them." ) return tools diff --git a/providers/common/ai/src/airflow/providers/common/ai/tools/_from_toolset.py b/providers/common/ai/src/airflow/providers/common/ai/tools/_from_toolset.py index 21bd028f338..6a59920d122 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/tools/_from_toolset.py +++ b/providers/common/ai/src/airflow/providers/common/ai/tools/_from_toolset.py @@ -27,12 +27,17 @@ from collections.abc import Hashable, Iterator from contextlib import contextmanager from contextvars import ContextVar from dataclasses import replace -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, TypeGuard from pydantic import ValidationError from pydantic_ai import RunContext from pydantic_ai.exceptions import ModelRetry, ToolFailed from pydantic_ai.models.test import TestModel +from pydantic_ai.toolsets.approval_required import ApprovalRequiredToolset +from pydantic_ai.toolsets.filtered import FilteredToolset +from pydantic_ai.toolsets.prefixed import PrefixedToolset +from pydantic_ai.toolsets.renamed import RenamedToolset +from pydantic_ai.toolsets.wrapper import WrapperToolset from pydantic_ai.usage import RunUsage from airflow.providers.common.ai.tools import AirflowTool, ToolCallError, ToolResult @@ -88,12 +93,102 @@ def airflow_tools_from_toolset(toolset: AbstractToolset[Any], *, deps: Any = Non :param toolset: The pydantic-ai toolset to expose. :param deps: Exposed to the toolset as ``ctx.deps``. """ - ctx: RunContext[Any] = RunContext(deps=deps, model=TestModel(), usage=RunUsage(), max_retries=1) + ctx = _inert_context(deps) toolset_tools = run_coroutine_sync(toolset.get_tools(ctx)) in_order = _InOrder() return [_as_airflow_tool(toolset, name, tool, ctx, in_order) for name, tool in toolset_tools.items()] +def _inert_context(deps: Any = None) -> RunContext[Any]: + return RunContext(deps=deps, model=TestModel(), usage=RunUsage(), max_retries=1) + + +# The wrappers that only choose which tools the model is offered, and under what name. +_NAMING_WRAPPERS = (FilteredToolset, PrefixedToolset, RenamedToolset) + + +def _is_naming_wrapper( + toolset: object, +) -> TypeGuard[FilteredToolset[Any] | PrefixedToolset[Any] | RenamedToolset[Any]]: + # Exact types: the tools reached around a wrapper skip its call_tool, which a subclass may override. + return type(toolset) in _NAMING_WRAPPERS + + +def airflow_tools_of(source: object) -> list[AirflowTool]: + """ + Return the tools of a toolset a framework adapter was given. + + A source with an ``airflow_tools()`` method is asked for them. A toolset wrapped with + ``.filtered()``, ``.prefixed()`` or ``.renamed()``, nested in any order, gives the innermost + toolset's own ``airflow_tools()``, so they behave as they do without the wrappers. The + wrappers' own ``get_tools`` decides which of them the model is offered and under what name, + given the same inert ``RunContext`` :func:`airflow_tools_from_toolset` passes. A ``.filtered()`` + function therefore runs once, when the adapter is built, rather than before every model turn + as in a pydantic-ai run. + + :raises TypeError: when ``source`` is any other wrapper, such as ``.approval_required()``, + or there is no ``airflow_tools()`` method under its wrappers. + :raises ValueError: when a ``.renamed()`` map names a tool the toolset does not have. + """ + # Looked up, not checked against ToolProvider: since Python 3.12 a Protocol check misses a + # method that a mock or proxy only provides through __getattr__. + own = getattr(source, "airflow_tools", None) + if callable(own): + return list(own()) + if not _is_naming_wrapper(source): + raise TypeError(_cannot_take(source)) + renames: list[PrefixedToolset[Any] | RenamedToolset[Any]] = [] + inner: AbstractToolset[Any] = source + while _is_naming_wrapper(inner): + if not isinstance(inner, FilteredToolset): + renames.append(inner) + inner = inner.wrapped + provide = getattr(inner, "airflow_tools", None) + if not callable(provide): + raise TypeError(_cannot_take(inner)) + by_name = {tool.name: tool for tool in provide()} + offered = run_coroutine_sync(source.get_tools(_inert_context())) + tools = [] + for name in offered: + # Map the offered name back through each wrapper, as its ``call_tool`` does. + original = name + for wrapper in renames: + if isinstance(wrapper, PrefixedToolset): + original = original.removeprefix(f"{wrapper.prefix}_") + else: + original = wrapper.name_map.get(original, original) + if original not in by_name: + raise ValueError( + f"The tool offered as {name!r} maps back to {original!r}, which " + f"{type(inner).__name__} does not have. .renamed() maps each new name to the " + "toolset's own name, as in .renamed({'run_sql': 'query'})." + ) + tools.append(replace(by_name[original], name=name)) + return tools + + +def _cannot_take(toolset: object) -> str: + kind = type(toolset).__name__ + if isinstance(toolset, ApprovalRequiredToolset): + return ( + "AirflowTools cannot carry ApprovalRequiredToolset (.approval_required()) to another " + "framework's agent: only a Pydantic AI run pauses for the approval it asks for. Pass the " + "toolset without it. To have a person approve tool calls, run the agent with " + "AgentOperator; a framework's own tool approval, such as needs_approval in the OpenAI " + "Agents SDK, stops the run and nothing in Airflow resumes it." + ) + if isinstance(toolset, WrapperToolset): + return ( + f"AirflowTools cannot read through {kind}. Only .filtered(), .prefixed() and .renamed() " + "carry over to another framework's agent; pass the toolset without it." + ) + return ( + f"AirflowTools cannot take a {kind}: it has no airflow_tools() method. Pass AirflowTool " + "objects or an Airflow toolset such as SQLToolset, which .filtered(), .prefixed() and " + ".renamed() may wrap." + ) + + class _InOrder: """ Run calls one at a time, in the order they were made, across threads and event loops. diff --git a/providers/common/ai/src/airflow/providers/common/ai/tools/adk.py b/providers/common/ai/src/airflow/providers/common/ai/tools/adk.py index a876e17195e..7a2716a11da 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/tools/adk.py +++ b/providers/common/ai/src/airflow/providers/common/ai/tools/adk.py @@ -42,7 +42,7 @@ if TYPE_CHECKING: from google.adk.agents.readonly_context import ReadonlyContext from google.adk.tools.tool_context import ToolContext - from airflow.providers.common.ai.tools import ToolProvider + from airflow.providers.common.ai.tools import ToolSource __all__ = ["AirflowTools"] @@ -55,7 +55,8 @@ class AirflowTools(BaseToolset): Pass toolsets that implement :class:`~airflow.providers.common.ai.tools.ToolProvider`, such as ``SQLToolset`` - and ``HookToolset``, or individual + and ``HookToolset``, optionally wrapped with ``.filtered()``, + ``.prefixed()`` or ``.renamed()``, or individual :class:`~airflow.providers.common.ai.tools.AirflowTool` objects. The agent, its model and its runner stay plain ADK. @@ -69,7 +70,7 @@ class AirflowTools(BaseToolset): warehouse = SQLToolset("warehouse", allowed_tables=["orders"]) agent = LlmAgent(name="analyst", model=model, tools=[AirflowTools(warehouse)]) - Each tool keeps the source tool's name, description and argument schema, and every + Each tool keeps its source's name, description and argument schema, and every result passes through Airflow's secret masker first. A result reaches the model as ``{"result": ...}``; a failure the model can correct, such as a query naming a missing column, as ``{"error": ...}``, ADK's own convention for a failed tool. Such @@ -82,7 +83,7 @@ class AirflowTools(BaseToolset): :param sources: Toolsets and tools to add, in order. """ - def __init__(self, *sources: ToolProvider | AirflowTool) -> None: + def __init__(self, *sources: ToolSource) -> None: super().__init__() self._adk_tools: list[BaseTool] = [ _AirflowAdkTool(tool) for tool in _collect_tools(sources, mcp_hint=_MCP_HINT) diff --git a/providers/common/ai/src/airflow/providers/common/ai/tools/anthropic.py b/providers/common/ai/src/airflow/providers/common/ai/tools/anthropic.py index e1f79d2d393..d313e388e72 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/tools/anthropic.py +++ b/providers/common/ai/src/airflow/providers/common/ai/tools/anthropic.py @@ -55,7 +55,7 @@ if TYPE_CHECKING: from anthropic.types.beta import BetaMessage from anthropic.types.beta.parsed_beta_message import ParsedBetaMessage - from airflow.providers.common.ai.tools import ToolProvider + from airflow.providers.common.ai.tools import ToolSource log = logging.getLogger(__name__) @@ -82,7 +82,8 @@ class AirflowTools: Pass toolsets that implement :class:`~airflow.providers.common.ai.tools.ToolProvider`, such as ``SQLToolset`` - and ``HookToolset``, or individual + and ``HookToolset``, optionally wrapped with ``.filtered()``, + ``.prefixed()`` or ``.renamed()``, or individual :class:`~airflow.providers.common.ai.tools.AirflowTool` objects. Give :attr:`tools` to ``client.beta.messages.tool_runner`` and drive the runner with :meth:`run`: @@ -103,7 +104,7 @@ class AirflowTools: ) message = tools.run(runner) - Each tool keeps the source tool's name, description and argument schema, and every + Each tool keeps its source's name, description and argument schema, and every result passes through Airflow's secret masker first. A failure the model can correct, such as a query naming a missing column, reaches it as a ``tool_result`` with ``is_error: true``, counted against the tool's retry limit once per model @@ -123,7 +124,7 @@ class AirflowTools: :param sources: Toolsets and tools to add, in order. """ - def __init__(self, *sources: ToolProvider | AirflowTool) -> None: + def __init__(self, *sources: ToolSource) -> None: self.tools: list[BetaFunctionTool[Callable[..., str]]] = [ _sync_tool(tool) for tool in _collect_tools(sources, mcp_hint=_MCP_HINT) ] @@ -164,7 +165,7 @@ class AsyncAirflowTools: :param sources: Toolsets and tools to add, in order. """ - def __init__(self, *sources: ToolProvider | AirflowTool) -> None: + def __init__(self, *sources: ToolSource) -> None: self.tools: list[BetaAsyncFunctionTool[Callable[..., Coroutine[Any, Any, str]]]] = [ _async_tool(tool) for tool in _collect_tools(sources, mcp_hint=_ASYNC_MCP_HINT) ] diff --git a/providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py b/providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py index 4a9da23d161..ab96a0c640c 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py +++ b/providers/common/ai/src/airflow/providers/common/ai/tools/claude_agent_sdk.py @@ -50,7 +50,7 @@ from airflow.providers.common.ai.utils.tool_metrics import calling_framework if TYPE_CHECKING: from collections.abc import AsyncIterable, Awaitable, Callable - from airflow.providers.common.ai.tools import ToolProvider + from airflow.providers.common.ai.tools import ToolSource log = logging.getLogger(__name__) @@ -152,7 +152,8 @@ class AirflowTools: Pass toolsets that implement :class:`~airflow.providers.common.ai.tools.ToolProvider`, such as ``SQLToolset`` - and ``HookToolset``, or individual + and ``HookToolset``, optionally wrapped with ``.filtered()``, + ``.prefixed()`` or ``.renamed()``, or individual :class:`~airflow.providers.common.ai.tools.AirflowTool` objects. They are served through an in-process MCP server; add it and :attr:`allowed_tools` to ``ClaudeAgentOptions``, or build them with :meth:`build_options`, and drive the SDK's @@ -169,7 +170,7 @@ class AirflowTools: options = tools.build_options(model="claude-sonnet-5", max_turns=10) result = await tools.run(query(prompt="...", options=options)) - Each tool keeps the source tool's name, description and argument schema, and every + Each tool keeps its source's name, description and argument schema, and every result passes through Airflow's secret masker first. A failure the model can correct, such as a query naming a missing column, reaches it as an ``is_error`` tool result. A failure the model cannot fix, such as a hook raising, makes @@ -189,7 +190,7 @@ class AirflowTools: (``mcp__<server_name>__<tool name>``). """ - def __init__(self, *sources: ToolProvider | AirflowTool, server_name: str = "airflow") -> None: + def __init__(self, *sources: ToolSource, server_name: str = "airflow") -> None: self.server_name = server_name airflow_tools = _collect_tools(sources, mcp_hint=_MCP_HINT) self.allowed_tools: list[str] = [f"mcp__{server_name}__{t.name}" for t in airflow_tools] diff --git a/providers/common/ai/src/airflow/providers/common/ai/tools/openai_agents.py b/providers/common/ai/src/airflow/providers/common/ai/tools/openai_agents.py index 8d855171f82..ea149442d14 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/tools/openai_agents.py +++ b/providers/common/ai/src/airflow/providers/common/ai/tools/openai_agents.py @@ -67,7 +67,7 @@ if TYPE_CHECKING: from opentelemetry.context import Context from opentelemetry.trace import Span as OtelSpan, TracerProvider - from airflow.providers.common.ai.tools import ToolProvider + from airflow.providers.common.ai.tools import ToolSource __all__ = ["AirflowTools", "AirflowTraceProcessor"] @@ -83,7 +83,8 @@ class AirflowTools(list[Tool]): Pass toolsets that implement :class:`~airflow.providers.common.ai.tools.ToolProvider`, such as ``SQLToolset`` - and ``HookToolset``, or individual + and ``HookToolset``, optionally wrapped with ``.filtered()``, + ``.prefixed()`` or ``.renamed()``, or individual :class:`~airflow.providers.common.ai.tools.AirflowTool` objects. The agent, its model and its runner stay plain OpenAI Agents SDK. @@ -101,7 +102,7 @@ class AirflowTools(list[Tool]): It is a list of ``FunctionTool`` objects typed as ``Agent.tools`` is, so add your own tools with ``[*AirflowTools(warehouse), my_tool]``. - Each tool keeps the source tool's name, description and argument schema, and every + Each tool keeps its source's name, description and argument schema, and every result passes through Airflow's secret masker first. The schema is passed as it is, with ``strict_json_schema=False``: the SDK's strict mode would make every argument required. A failure the model can correct, such as a query naming a missing column, @@ -120,7 +121,7 @@ class AirflowTools(list[Tool]): :param sources: Toolsets and tools to add, in order. """ - def __init__(self, *sources: ToolProvider | AirflowTool) -> None: + def __init__(self, *sources: ToolSource) -> None: super().__init__(_to_function_tool(tool) for tool in _collect_tools(sources, mcp_hint=_MCP_HINT)) diff --git a/providers/common/ai/src/airflow/providers/common/ai/tools/strands.py b/providers/common/ai/src/airflow/providers/common/ai/tools/strands.py index a927c7e8429..b1238d9e8ed 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/tools/strands.py +++ b/providers/common/ai/src/airflow/providers/common/ai/tools/strands.py @@ -52,7 +52,7 @@ if TYPE_CHECKING: ToolUse, ) - from airflow.providers.common.ai.tools import ToolProvider + from airflow.providers.common.ai.tools import ToolSource __all__ = ["AirflowTools"] @@ -68,7 +68,8 @@ class AirflowTools(Plugin): Pass toolsets that implement :class:`~airflow.providers.common.ai.tools.ToolProvider`, such as ``SQLToolset`` - and ``HookToolset``, or individual + and ``HookToolset``, optionally wrapped with ``.filtered()``, + ``.prefixed()`` or ``.renamed()``, or individual :class:`~airflow.providers.common.ai.tools.AirflowTool` objects. The agent, its model and its loop stay plain Strands. @@ -85,7 +86,7 @@ class AirflowTools(Plugin): Pass it in ``plugins=``, not ``tools=``: Strands ignores a plugin given as a tool, and handing over only its ``tools`` leaves out the hook that ends the run on a failure. - Each tool keeps the source tool's name, description and argument schema, and every + Each tool keeps its source's name, description and argument schema, and every result passes through Airflow's secret masker first. A failure the model can correct, such as a query naming a missing column, reaches it as a Strands result with ``status="error"``, counted against the tool's retry limit once per model turn @@ -99,7 +100,7 @@ class AirflowTools(Plugin): :param sources: Toolsets and tools to add, in order. """ - def __init__(self, *sources: ToolProvider | AirflowTool) -> None: + def __init__(self, *sources: ToolSource) -> None: # Before Plugin.__init__, which reads every attribute, ``name`` and ``tools`` included. self._name = f"airflow-tools-{next(_instance_numbers)}" self._run = object() diff --git a/providers/common/ai/tests/unit/common/ai/tools/test_openai_agents.py b/providers/common/ai/tests/unit/common/ai/tools/test_openai_agents.py index afb82eee50c..9c7b482229c 100644 --- a/providers/common/ai/tests/unit/common/ai/tools/test_openai_agents.py +++ b/providers/common/ai/tests/unit/common/ai/tools/test_openai_agents.py @@ -18,7 +18,7 @@ from __future__ import annotations import copy import json -from typing import Any +from typing import TYPE_CHECKING, Any from unittest.mock import MagicMock, patch import pytest @@ -68,6 +68,9 @@ from tests_common.test_utils.config import conf_vars from unit.common.ai.tools.test_tracing import TI from unit.common.ai.toolsets.test_sql import _make_mock_db_hook +if TYPE_CHECKING: + from airflow.providers.common.ai.tools import ToolSource + TRACING = "airflow.providers.common.ai.tools.tracing" MODULE = "airflow.providers.common.ai.tools.openai_agents" _OK = ToolResult("ok") @@ -133,7 +136,7 @@ def _tool(result: ToolResult | Exception) -> AirflowTool: ) -def _function_tools(*sources: AirflowTool | SQLToolset) -> list[FunctionTool]: +def _function_tools(*sources: ToolSource) -> list[FunctionTool]: tools = [] for tool in AirflowTools(*sources): assert isinstance(tool, FunctionTool) @@ -177,6 +180,11 @@ class TestAirflowTools: assert tool.strict_json_schema is False assert tool.params_json_schema == source.parameters + def test_a_filtered_toolset_offers_every_tool_the_filter_keeps(self): + tools = _function_tools(SQLToolset("pg_default").filtered(lambda ctx, tool: tool.name != "query")) + + assert [t.name for t in tools] == ["list_tables", "get_schema", "check_query"] + def test_accepts_individual_tools_alongside_toolsets(self): tools = AirflowTools(_tool(ToolResult("a")), SQLToolset("pg_default")) diff --git a/providers/common/ai/tests/unit/common/ai/tools/test_tools.py b/providers/common/ai/tests/unit/common/ai/tools/test_tools.py index be80d3b3fc7..fcc754e92ca 100644 --- a/providers/common/ai/tests/unit/common/ai/tools/test_tools.py +++ b/providers/common/ai/tests/unit/common/ai/tools/test_tools.py @@ -17,11 +17,16 @@ from __future__ import annotations import asyncio +import json import traceback from functools import partial -from typing import Any +from typing import TYPE_CHECKING, Any +from unittest.mock import MagicMock import pytest +from pydantic_ai.toolsets.filtered import FilteredToolset +from pydantic_ai.toolsets.function import FunctionToolset +from pydantic_ai.toolsets.wrapper import WrapperToolset from airflow.providers.common.ai.exceptions import MCPToolsetNotSupportedError from airflow.providers.common.ai.tools import ( @@ -32,6 +37,13 @@ from airflow.providers.common.ai.tools import ( collect_tools, ) from airflow.providers.common.ai.toolsets.mcp import MCPToolset +from airflow.providers.common.ai.toolsets.sql import SQLToolset + +from unit.common.ai.toolsets.test_sql import _make_mock_db_hook + +if TYPE_CHECKING: + from pydantic_ai import RunContext + from pydantic_ai.tools import ToolDefinition def _tool(function) -> AirflowTool: @@ -117,7 +129,7 @@ class TestCollectTools: """Frameworks route a call by name, so one of the two would never run.""" with pytest.raises( ValueError, - match="More than one tool is named query. .*tool_prefix values, where they take one, so", + match="More than one tool is named query. .*where they take one, or wrap them with .prefixed\\(\\), so", ): collect_tools([_Provider(_named("query")), _Provider(_named("query"))]) @@ -153,6 +165,183 @@ class TestCollectTools: with pytest.raises(NotImplementedError, match="^not yet$"): _collect_tools([_Unfinished()], mcp_hint="Connect Acme's McpClient instead.") + def test_a_mock_toolset_is_asked_for_its_tools(self): + """The mock has ``airflow_tools`` only through ``__getattr__``, which a Protocol check misses on Python 3.12+.""" + source = MagicMock(spec=["airflow_tools"]) + source.airflow_tools.return_value = [_named("a")] + + assert [tool.name for tool in collect_tools([source])] == ["a"] + + +def _sql_toolset(name: str = "Ada") -> SQLToolset: + ts = SQLToolset("pg_default") + ts._hook = _make_mock_db_hook(records=[(1, name)], last_description=[("id",), ("name",)]) + return ts + + +def _names(tools: list[AirflowTool]) -> list[str]: + return [tool.name for tool in tools] + + +def _not_query(ctx: RunContext[Any], tool: ToolDefinition) -> bool: + return tool.name != "query" + + +async def _not_query_async(ctx: RunContext[Any], tool: ToolDefinition) -> bool: + return tool.name != "query" + + +class TestCollectToolsThroughWrappers: + @pytest.mark.parametrize("keep", [_not_query, _not_query_async], ids=["sync", "async"]) + def test_a_filter_drops_the_tools_it_rejects(self, keep): + tools = collect_tools([_sql_toolset().filtered(keep)]) + + assert _names(tools) == ["list_tables", "get_schema", "check_query"] + + def test_the_filter_runs_once_with_each_tools_definition(self): + seen: list[tuple[RunContext[Any], ToolDefinition]] = [] + + def keep_all(ctx: RunContext[Any], tool: ToolDefinition) -> bool: + seen.append((ctx, tool)) + return True + + tools = collect_tools([_sql_toolset().filtered(keep_all)]) + asyncio.run({tool.name: tool for tool in tools}["query"].call({"sql": "SELECT 1"})) + + assert [tool.name for _, tool in seen] == ["list_tables", "get_schema", "query", "check_query"] + # The toolset's own definitions, not ones rebuilt from the AirflowTool. + assert all(tool.sequential for _, tool in seen) + assert all(ctx.deps is None for ctx, _ in seen) + + @pytest.mark.parametrize( + ("wrap", "offered", "query_tool"), + [ + pytest.param( + lambda ts: ts.prefixed("wh"), + ["wh_list_tables", "wh_get_schema", "wh_query", "wh_check_query"], + "wh_query", + id="prefixed", + ), + pytest.param( + lambda ts: ts.renamed({"run_sql": "query"}), + ["list_tables", "get_schema", "run_sql", "check_query"], + "run_sql", + id="renamed", + ), + ], + ) + def test_a_renamed_tool_runs_the_inner_tool(self, wrap, offered, query_tool): + tools = {tool.name: tool for tool in collect_tools([wrap(_sql_toolset())])} + + result = asyncio.run(tools[query_tool].call({"sql": "SELECT id, name FROM users"})) + + assert list(tools) == offered + assert json.loads(result.content)["rows"] == [[1, "Ada"]] + + def test_nested_wrappers_apply_from_the_inside_out(self): + toolset = ( + _sql_toolset() + .prefixed("wh") + .filtered(lambda ctx, tool: tool.name != "wh_get_schema") + .renamed({"run_sql": "wh_query"}) + .filtered(lambda ctx, tool: tool.name != "wh_check_query") + ) + + tools = {tool.name: tool for tool in collect_tools([toolset])} + result = asyncio.run(tools["run_sql"].call({"sql": "SELECT id, name FROM users"})) + + assert list(tools) == ["wh_list_tables", "run_sql"] + assert json.loads(result.content)["rows"] == [[1, "Ada"]] + + def test_two_prefixes_keep_two_toolsets_on_the_same_tools_apart(self): + tools = collect_tools([_sql_toolset("Ada").prefixed("a"), _sql_toolset("Grace").prefixed("b")]) + by_name = {tool.name: tool for tool in tools} + + rows = { + name: json.loads(asyncio.run(by_name[name].call({"sql": "SELECT id, name FROM users"})).content)[ + "rows" + ] + for name in ("a_query", "b_query") + } + + assert _names(tools) == [ + *(f"a_{name}" for name in ("list_tables", "get_schema", "query", "check_query")), + *(f"b_{name}" for name in ("list_tables", "get_schema", "query", "check_query")), + ] + assert rows == {"a_query": [[1, "Ada"]], "b_query": [[1, "Grace"]]} + + def test_a_rename_map_written_backwards_says_which_way_it_goes(self): + with pytest.raises( + ValueError, + match=r"The tool offered as 'query' maps back to 'run_sql'.*\.renamed\(\{'run_sql': 'query'\}\)", + ): + collect_tools([_sql_toolset().renamed({"query": "run_sql"})]) + + def test_a_wrapper_that_supplies_its_own_tools_is_asked_for_them(self): + class _OwnTools(WrapperToolset[Any]): + def airflow_tools(self) -> list[AirflowTool]: + return [_named("own")] + + assert _names(collect_tools([_OwnTools(_sql_toolset())])) == ["own"] + + def test_a_subclass_of_a_naming_wrapper_is_refused(self): + """Its call_tool could add behaviour that tools reached around it would skip.""" + + class _AuditedFilter(FilteredToolset[Any]): + pass + + with pytest.raises(TypeError, match="cannot read through _AuditedFilter"): + collect_tools([_AuditedFilter(_sql_toolset(), lambda ctx, tool: True)]) + + def test_an_mcp_toolset_behind_a_filter_says_what_to_connect_instead(self): + with pytest.raises(MCPToolsetNotSupportedError, match="Connect Acme's McpClient instead."): + _collect_tools( + [MCPToolset("mcp_default").filtered(lambda ctx, tool: True)], + mcp_hint="Connect Acme's McpClient instead.", + ) + + @pytest.mark.parametrize( + "wrap", + [ + pytest.param(lambda ts: ts.approval_required(), id="outermost"), + pytest.param(lambda ts: ts.approval_required().filtered(lambda ctx, tool: True), id="nested"), + ], + ) + def test_approval_is_refused_and_points_at_the_frameworks_own(self, wrap): + with pytest.raises(TypeError) as exc_info: + collect_tools([wrap(_sql_toolset())]) + + assert str(exc_info.value) == ( + "AirflowTools cannot carry ApprovalRequiredToolset (.approval_required()) to another " + "framework's agent: only a Pydantic AI run pauses for the approval it asks for. Pass the " + "toolset without it. To have a person approve tool calls, run the agent with " + "AgentOperator; a framework's own tool approval, such as needs_approval in the OpenAI " + "Agents SDK, stops the run and nothing in Airflow resumes it." + ) + + def test_another_wrapper_is_refused_by_name(self): + async def unchanged(ctx: RunContext[Any], tools: list[ToolDefinition]) -> list[ToolDefinition]: + return tools + + with pytest.raises(TypeError) as exc_info: + collect_tools([_sql_toolset().prepared(unchanged)]) + + assert str(exc_info.value) == ( + "AirflowTools cannot read through PreparedToolset. Only .filtered(), .prefixed() and " + ".renamed() carry over to another framework's agent; pass the toolset without it." + ) + + @pytest.mark.parametrize( + "wrap", + [ + pytest.param(lambda ts: ts, id="bare"), + pytest.param(lambda ts: ts.filtered(lambda ctx, tool: True), id="filtered"), + ], + ) + def test_a_toolset_without_airflow_tools_is_refused(self, wrap): + with pytest.raises(TypeError, match="cannot take a FunctionToolset: it has no airflow_tools"): + collect_tools([wrap(FunctionToolset())]) + @pytest.mark.enable_redact class TestAirflowToolMasking:
