This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch durable-toolsets-without-id in repository https://gitbox.apache.org/repos/asf/airflow.git
commit e0a3755474701ea75969154d161c07ca6d118b53 Author: Kaxil Naik <[email protected]> AuthorDate: Tue Nov 10 20:42:41 2026 +0000 Let durable agent tasks use toolsets without an id AgentOperator(durable=True) and @task.agent(durable=True) no longer refuse a FunctionToolset without an id or a function that builds a toolset per run. The operator gives an id-less FunctionToolset a copy numbered in the order the toolsets are passed, and a function its module and qualified name (without the path hash Airflow puts in a Dag file's module name). Function tool steps are named after the tool alone; a toolset id the author chose keeps its place in step names, so journals recorded before this change still replay. A pydantic-ai MCPToolset without an id is now refused when the Dag loads instead of on the first run, and so is a callable without a qualified name or two functions sharing one. --- providers/common/ai/docs/capabilities.rst | 8 +- providers/common/ai/docs/durable_execution.rst | 66 ++++++-- providers/common/ai/docs/operators/agent.rst | 9 +- providers/common/ai/docs/troubleshooting.rst | 21 ++- .../providers/common/ai/durable/capability.py | 138 +++++++++++++++- .../ai/example_dags/example_agent_durable.py | 39 ++++- .../airflow/providers/common/ai/operators/agent.py | 100 +++++++----- .../unit/common/ai/durable/test_capability.py | 177 ++++++++++++++++++++- .../tests/unit/common/ai/operators/test_agent.py | 144 +++++++++++++++-- 9 files changed, 614 insertions(+), 88 deletions(-) diff --git a/providers/common/ai/docs/capabilities.rst b/providers/common/ai/docs/capabilities.rst index 08fc2aaf135..41b8bcc71c6 100644 --- a/providers/common/ai/docs/capabilities.rst +++ b/providers/common/ai/docs/capabilities.rst @@ -110,9 +110,11 @@ again. Whether a capability's work is replayed depends on where it runs: * - ``WebSearch``, ``WebFetch`` or ``ImageGeneration`` falling back to a local tool, for a provider without the native one; ``Toolset``, ``MCP``, ``PrefixTools``, ``CombinedCapability``, and capabilities from an agent spec file - - Tool results are replayed. Each toolset they contribute needs an ``id``; pydantic-ai - names one without it when the agent is built. An ``id`` on a ``Toolset`` capability - does not reach the toolset inside it, so set it on that toolset. + - Tool results are replayed. In ``AgentOperator``, a ``FunctionToolset`` held directly by + a ``Toolset`` capability needs no ``id`` (:ref:`durable-toolset-ids`); any other toolset + they contribute needs one, and pydantic-ai names one without it when the agent is + built. An ``id`` on a ``Toolset`` capability does not reach the toolset inside it, so + set it on that toolset. * - pydantic-ai-harness ``CodeMode`` - Each ``run_code`` snippet runs again, and the tool calls it makes are replayed; see :ref:`durable-code-mode`. diff --git a/providers/common/ai/docs/durable_execution.rst b/providers/common/ai/docs/durable_execution.rst index 57cda143168..1e485079fa2 100644 --- a/providers/common/ai/docs/durable_execution.rst +++ b/providers/common/ai/docs/durable_execution.rst @@ -65,30 +65,64 @@ The capability needs a name for the agent's steps: the agent's ``name``, or ``agent_params`` sets one, and the task id otherwise. Outside a running Airflow task the capability does nothing, so the same agent runs normally in a test or a notebook. -Give every toolset an id ------------------------- +.. _durable-toolset-ids: -pydantic-ai records each call of a ``FunctionToolset``, ``MCPToolset`` or -``DynamicToolset`` under the toolset's ``id``, and refuses to run one without it. Set -the id on the toolset itself: +Toolset ids +----------- + +pydantic-ai names the recorded steps of a ``FunctionToolset``, ``MCPToolset`` or +``DynamicToolset`` after the toolset's ``id``. ``AgentOperator`` and ``@task.agent`` fill +in the ids you leave out of a ``FunctionToolset`` and of a function that builds a toolset +for each run passed in ``toolsets`` or ``agent_params["toolsets"]``, and of a +``FunctionToolset`` held directly by a ``Toolset`` capability, so this agent needs none: + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py + :language: python + :start-after: [START howto_agent_durable_toolsets_without_id] + :end-before: [END howto_agent_durable_toolsets_without_id] + +A tool call of a ``FunctionToolset`` without an id is recorded under the tool's name +alone, ``durable_order_status__function_toolset.call_tool:get_order`` above, so a retry +of the same agent finds the calls the failed attempt recorded. A function that builds a +toolset, such as ``carrier_tools``, is named by its module and qualified name +(``example_agent_durable.carrier_tools``). A toolset you give an ``id`` keeps it in its +step names. Adding a tool to a toolset, reordering the toolsets, or renaming a function +or an id between attempts changes what the model is offered, so the retry runs again +from the first model request (:ref:`durable-replay`). A function that builds a toolset +per run is the exception for its own tools: its tool list is a recorded step, so the +steps a retry replays, and the request that failed, keep the tools the failed attempt +saw, and later steps get the new list. A clear that deletes the journal +(:ref:`durable-clear`) starts over with the new list. + +An ``MCPToolset`` from pydantic-ai still needs an id from you. Its calls are named after +the toolset alone, and no id is derived from the server's URL or command, which can carry +credentials. The Dag fails to load with: -.. code-block:: python +.. code-block:: text - orders = FunctionToolset(tools=[get_order], id="orders") + ValueError: durable=True needs a unique id on every MCPToolset, such as MCPToolset(..., id='weather'): durable replay records each tool call under it, and none is derived from the server's URL or command, which can carry credentials. Set it on the toolset itself; the id of a Toolset capability around a toolset does not reach it. -``AgentOperator(durable=True)`` checks this when the Dag is parsed and fails with: +So does a callable that has no ``__qualname__`` to name it by, such as a +``functools.partial``; pass it as ``DynamicToolset(function, id=...)``: .. code-block:: text - ValueError: durable=True needs a unique id on every FunctionToolset, such as FunctionToolset(..., id='orders'): durable replay records each tool call under it. Set it on the toolset itself, and pass a function that builds a toolset as DynamicToolset(function, id=...); the id of a Toolset capability around a toolset does not reach it. + ValueError: durable=True names the steps of a toolset built per run after the function that builds it, but functools.partial(<function carrier_tools at 0xffff89c904a0>, carrier='dhl') has no __qualname__. Pass it as DynamicToolset(function, id=...). + +Two different toolsets with the same id are refused, since their calls would be filed +under one name, and so are two functions with the same qualified name, such as two +lambdas in one Dag file: -It also refuses two different toolsets with the same id, which would file their calls -under one name; passing one toolset object in two places is fine. +.. code-block:: text -Keep an id stable once tasks run with it: a retry finds its recorded tool calls by it. -The provider's own toolsets (``SQLToolset``, ``HookToolset`` and the rest) and tools -passed as ``agent_params={"tools": [...]}`` need nothing. A capability that contributes -``@durable_operation`` methods needs an ``id`` for the same reason; pydantic-ai raises + ValueError: durable=True needs a unique id on each toolset, but more than one toolset has the id 'orders_dag_lambdas.<lambda>': durable replay records each tool call under its toolset's id, so a retry could not tell their calls apart. + +Passing one toolset object in two places is fine. The provider's own toolsets +(``SQLToolset``, ``HookToolset``, its ``MCPToolset`` and the rest) and tools passed as +``agent_params={"tools": [...]}`` need nothing. ``AirflowDurability`` in a plain ``@task`` +does not fill in ids: there pydantic-ai raises ``UserError`` for a ``FunctionToolset``, +``MCPToolset`` or ``DynamicToolset`` without one. A capability that contributes +``@durable_operation`` methods needs an ``id`` for the same reason, and pydantic-ai raises ``UserError`` naming it when it has none. pydantic-ai-harness ``SpendLimits`` has one by default. @@ -136,6 +170,8 @@ Without the option, the task fails with: ValueError: Durable execution on Airflow before 3.3 needs [common.ai] durable_cache_path to be set, to say where its journal is kept. Example: durable_cache_path = file:///tmp/airflow_durable_cache +.. _durable-replay: + How a retry replays ------------------- diff --git a/providers/common/ai/docs/operators/agent.rst b/providers/common/ai/docs/operators/agent.rst index f0506d1a494..6b3df531085 100644 --- a/providers/common/ai/docs/operators/agent.rst +++ b/providers/common/ai/docs/operators/agent.rst @@ -431,9 +431,12 @@ Parameters - ``durable``: When ``True``, records each step the agent completes (model requests, tool calls, other capabilities' ``@durable_operation`` methods) in a journal kept for the task instance, and replays the recorded steps on retry - instead of running them again. Every ``FunctionToolset``, ``MCPToolset`` and - ``DynamicToolset`` needs an ``id``. On Airflow >= 3.3 the journal uses the task - state store (no configuration needed); on older Airflow versions it requires the + instead of running them again. A pydantic-ai ``MCPToolset`` needs an ``id``. A + ``FunctionToolset`` or a function that builds a toolset, passed in ``toolsets`` or + ``agent_params["toolsets"]``, needs none, and neither does a ``FunctionToolset`` held + directly by a ``Toolset`` capability (:ref:`durable-toolset-ids`). On Airflow >= 3.3 + the journal uses the task state store (no configuration needed); on older Airflow + versions it requires the ``[common.ai] durable_cache_path`` config option to be set. Default ``False``. A replayed step adds nothing to the usage counted against ``usage_limits`` or reported in the ``usage`` XCom -- not its request, tokens, cost, or tool calls -- so diff --git a/providers/common/ai/docs/troubleshooting.rst b/providers/common/ai/docs/troubleshooting.rst index e5867ec9532..52ff66d09ec 100644 --- a/providers/common/ai/docs/troubleshooting.rst +++ b/providers/common/ai/docs/troubleshooting.rst @@ -121,15 +121,22 @@ as a task failure. attempt, which a human review loop's regenerations do not keep. Pick one. See :doc:`durable_execution`. -``durable=True needs a unique id on every FunctionToolset, such as FunctionToolset(..., id='orders')`` - Durable replay records each tool call under its toolset's id, so a - ``FunctionToolset``, ``MCPToolset`` or ``DynamicToolset`` without one is refused when - the Dag loads. Set ``id=`` on the toolset itself; an id on a ``Toolset`` capability - around it does not reach it. See :doc:`durable_execution`. +``durable=True needs a unique id on every MCPToolset, such as MCPToolset(..., id='weather')`` + Durable replay records each call of a pydantic-ai ``MCPToolset`` under the toolset's + id, and none is derived from the server's URL or command. Set ``id=`` on the toolset + itself; an id on a ``Toolset`` capability around it does not reach it. See + :ref:`durable-toolset-ids`. + +``durable=True names the steps of a toolset built per run after the function that builds it, but ... has no __qualname__`` + A function that builds a toolset is named by its module and qualified name, and a + ``functools.partial`` or a callable object has none. Pass it as + ``DynamicToolset(function, id=...)``. See :ref:`durable-toolset-ids`. ``durable=True needs a unique id on each toolset, but more than one toolset has the id 'orders'`` - Two different toolset objects share an id, so durable replay could not tell their - recorded calls apart. Give each its own ``id=``. See :doc:`durable_execution`. + Two different toolset objects share an id, or two functions that build toolsets share + a module and qualified name (two lambdas in one Dag file), so durable replay could not + tell their recorded calls apart. Give each its own ``id=``, or wrap a function in + ``DynamicToolset(function, id=...)``. See :ref:`durable-toolset-ids`. ``durable=True attaches AirflowDurability itself; pass durable=True or the capability, not both.`` ``AgentOperator(durable=True)`` adds the capability to the agent, so a second one in diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py b/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py index 1a8d73b8732..f3adf01aa08 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py +++ b/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py @@ -19,22 +19,25 @@ from __future__ import annotations import base64 +import copy import dataclasses +import re from collections.abc import Awaitable, Callable, Iterator, Mapping from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass, field from functools import partial -from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, cast +from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypeGuard, cast from pydantic_ai.durable_exec import ( JSON_CODEC, BaseDurabilityCapability, + CallableOperationBackend, CapabilityOperationId, DurabilityCodec, DurabilityEngineSpec, DurableOperationId, - JournalCallableOperationBackend, + JournalOperationNamer, ModelRequestId, RoleBasedOperationConfig, ToolsetCallToolId, @@ -49,7 +52,7 @@ from pydantic_ai.exceptions import ( ) from pydantic_ai.messages import ModelResponse, ToolReturn from pydantic_ai.tools import AgentDepsT -from pydantic_ai.toolsets import DynamicToolset, FunctionToolset, ToolsetTool +from pydantic_ai.toolsets import AbstractToolset, DynamicToolset, FunctionToolset, ToolsetTool from pydantic_ai.toolsets.wrapper import WrapperToolset from pydantic_core import PydanticSerializationError @@ -61,6 +64,7 @@ from airflow.providers.common.ai.exceptions import DurableJournalError, DurableT from airflow.providers.common.ai.utils.masking import mask_secrets from airflow.providers.common.ai.utils.tool_metrics import record_tool_call from airflow.providers.common.ai.utils.toolset_base import AirflowToolset +from airflow.providers.common.ai.utils.toolsets import iter_toolsets if TYPE_CHECKING: from pydantic_ai.agent import EventStreamHandler @@ -74,7 +78,7 @@ if TYPE_CHECKING: from pydantic_ai.run import AgentRunResult from pydantic_ai.settings import ModelSettings from pydantic_ai.tools import RunContext, ToolDefinition - from pydantic_ai.toolsets import AbstractToolset + from pydantic_ai.toolsets import ToolsetFunc from airflow.providers.common.ai.durable.journal import DurableRun, JournalStep, StepKind @@ -96,6 +100,100 @@ else: DURABLE_UNIT_TOOLSETS = (FunctionToolset, DynamicToolset, MCPToolset) """Leaf toolsets pydantic-ai's durable backend runs as durable units; each needs a unique ``id``.""" +# Airflow imports a Dag file as a module named ``unusual_prefix_<sha1 of the file's path>_<file name>`` +# (``get_unique_dag_module_name`` in the shared ``module_loading`` library, which providers cannot import). +_DAG_MODULE_PREFIX = re.compile(r"^unusual_prefix_[0-9a-f]{40}_") + + +def callable_toolset_id(function: object) -> str | None: + """ + Return the id of a toolset that ``function`` builds per run: its module and qualified name. + + The hash of the file's path that Airflow puts in a Dag file's module name is left out, + because the path can differ between workers and bundle versions. ``None`` for a callable + without a qualified name, such as a ``functools.partial``. + """ + module = getattr(function, "__module__", None) + qualname = getattr(function, "__qualname__", None) + if not isinstance(module, str) or not isinstance(qualname, str): + return None + return f"{_DAG_MODULE_PREFIX.sub('', module)}.{qualname}" + + +def chosen_toolset_id(leaf: AbstractToolset[Any]) -> str | None: + """ + Return the id pydantic-ai names a leaf toolset's durable steps after, or ``None`` when it has none. + + An ``MCPToolset`` without an id derives one from its server's hostname, but durable steps + are only named after an id that was chosen for the toolset. + """ + # Before ``_explicit_id`` (pydantic-ai 2.62) no toolset derived its id, so ``id`` is the chosen one. + return getattr(leaf, "_explicit_id", leaf.id) + + +# The id ``give_durable_ids`` gives the n-th ``FunctionToolset`` without one. The same agent gets the +# same ids on every attempt, so its requests to the model match the recorded ones. Angle brackets +# mark an id the framework chose, as in pydantic-ai's own ``<agent>``; there is no colon, which +# pydantic-ai rejects in a toolset id because it builds instruction ids from it. +_ASSIGNED_ID_PREFIX = "<airflow-function_toolset-" + + +def _assigned_toolset_id(number: int) -> str: + return f"{_ASSIGNED_ID_PREFIX}{number}>" + + +# Each leaf ``give_durable_ids`` copied, keyed by ``id()``, with its copy. The leaf is kept so +# that its ``id()`` cannot be reused by another object while the map is in use. +GivenIds: TypeAlias = "dict[int, tuple[AbstractToolset[Any], AbstractToolset[Any]]]" + + +def give_durable_ids( + toolset: AbstractToolset[Any] | ToolsetFunc[Any], given: GivenIds +) -> AbstractToolset[Any]: + """ + Return ``toolset`` with an id on each leaf pydantic-ai runs as durable steps, where one can be given. + + pydantic-ai runs a ``FunctionToolset`` or ``DynamicToolset`` as durable steps only when + it has an id. An id-less ``FunctionToolset`` gets one numbered in the order the toolsets + are given, which its step names leave out (a function tool's step is named after the + tool); a function that builds a toolset per run, or an id-less ``DynamicToolset``, gets + :func:`callable_toolset_id`. A leaf that already has an id, and an ``MCPToolset``, are + left as they are. Leaves are copied, not changed, and ``given`` maps each leaf to its + copy so a toolset listed twice stays one toolset to pydantic-ai. + """ + if not isinstance(toolset, AbstractToolset): + return DynamicToolset(toolset, id=callable_toolset_id(toolset)) + if not any(_lacks_id(leaf) for leaf in iter_toolsets(toolset)): + # Rebuilding wrappers via visit_and_replace breaks wrapper subclasses with their own __init__. + return toolset + + def give(leaf: AbstractToolset[Any]) -> AbstractToolset[Any]: + if not _lacks_id(leaf): + return leaf + if (entry := given.get(id(leaf))) is None: + copied: AbstractToolset[Any] + if isinstance(leaf, DynamicToolset): + copied = DynamicToolset( + leaf.toolset_func, + per_run_step=leaf.per_run_step, + id=callable_toolset_id(leaf.toolset_func), + ) + else: + numbered = sum(isinstance(other, FunctionToolset) for _, other in given.values()) + function_toolset = copy.copy(leaf) + # FunctionToolset takes its id only in the constructor and keeps it here. + function_toolset._id = _assigned_toolset_id(numbered) + copied = function_toolset + entry = given[id(leaf)] = (leaf, copied) + return entry[1] + + return toolset.visit_and_replace(give) + + +def _lacks_id(leaf: AbstractToolset[Any]) -> TypeGuard[FunctionToolset[Any] | DynamicToolset[Any]]: + return isinstance(leaf, FunctionToolset | DynamicToolset) and leaf.id is None + + # Why a recorded step that matched is run live anyway: it was written by a version whose # types have changed since, so replaying it would fail on every retry. _UNLOADABLE = "the recorded result no longer loads, so it was written by another version" @@ -294,6 +392,10 @@ class AirflowDurability(BaseDurabilityCapability[AgentDepsT]): Toolsets the agent calls through pydantic-ai's durable backend need a unique ``id`` (``FunctionToolset(id=...)``, ``MCPToolset(..., id=...)``), as do capabilities that contribute ``@durable_operation`` methods; the ids name the steps in the journal. + ``AgentOperator`` and ``@task.agent`` fill in the ids of a ``FunctionToolset`` or a + function that builds a toolset passed in ``toolsets`` or ``agent_params["toolsets"]``, + and of a ``FunctionToolset`` held directly by a ``Toolset`` capability + (:ref:`durable-toolset-ids`); in a plain ``@task`` you set them. :param models: Extra models the run may switch to, keyed by id; see pydantic-ai's ``BaseDurabilityCapability``. @@ -399,12 +501,36 @@ class AirflowDurability(BaseDurabilityCapability[AgentDepsT]): return model -class _AirflowOperationBackend(JournalCallableOperationBackend[_StepMode]): +class _AirflowOperationNamer(JournalOperationNamer): + """ + pydantic-ai's journal step names, without an id Airflow gave a function toolset. + + A function tool's step name ends in the tool's name, so the id Airflow numbered the + toolset by adds nothing to it and would only make the name harder to read in the task + log. An id the Dag author chose stays in the name, so steps recorded under it before + this naming existed replay as they did. + """ + + def operation_name(self, operation_id: DurableOperationId) -> str: + name = super().operation_name(operation_id) + if ( + isinstance(operation_id, ToolsetCallToolId | ToolsetValidateToolArgumentsId) + and operation_id.toolset_kind == "function" + and operation_id.toolset_id.startswith(_ASSIGNED_ID_PREFIX) + ): + return name.replace(f"__{operation_id.toolset_id}", "", 1) + return name + + +class _AirflowOperationBackend(CallableOperationBackend[_StepMode]): """Runs each durable operation through the task's durable journal.""" def __init__(self, durability: AirflowDurability[Any]) -> None: super().__init__( - agent_name=durability.name, default_model_id=durability.default_model_id, config=_STEP_MODES + namer=_AirflowOperationNamer( + durability.name, default_model_id=durability.default_model_id or "default" + ), + config=_STEP_MODES, ) self._durability = durability # How pydantic-ai decodes each operation's result, by operation. diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py index 2ce8d0a84ec..68d10c10cc4 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py +++ b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py @@ -20,7 +20,7 @@ from __future__ import annotations from datetime import UTC, datetime, timedelta -from pydantic_ai import Tool +from pydantic_ai import RunContext, Tool from pydantic_ai.capabilities import SetToolMetadata from pydantic_ai.toolsets import FunctionToolset @@ -156,3 +156,40 @@ if SQLToolset is not None: # [END howto_agent_durable_tool_opt_out] example_agent_durable_tool_opt_out() + + +# --------------------------------------------------------------------------- +# 5. Toolsets without an id +# --------------------------------------------------------------------------- + + +# [START howto_agent_durable_toolsets_without_id] +def get_order(order_id: int) -> dict[str, str]: + """Look up an order's status and carrier.""" + return {"status": "shipped", "carrier": "dhl"} + + +def track_parcel(carrier: str, order_id: int) -> str: + """Return where a carrier's parcel for an order is.""" + return f"{carrier} parcel for order {order_id}: out for delivery" + + +def carrier_tools(ctx: RunContext) -> FunctionToolset: + """Build the carrier tools for each run.""" + return FunctionToolset(tools=[track_parcel]) + + +@dag(default_args={"retries": 3, "retry_delay": timedelta(seconds=30)}, tags=["example"]) +def example_agent_durable_toolsets_without_id(): + AgentOperator( + task_id="durable_order_status", + prompt="Where is order 4411?", + llm_conn_id="pydanticai_default", + durable=True, + toolsets=[FunctionToolset(tools=[get_order]), carrier_tools], + ) + + +# [END howto_agent_durable_toolsets_without_id] + +example_agent_durable_toolsets_without_id() diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py index 1c85bf7a83f..583c179e283 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py @@ -32,10 +32,16 @@ from pydantic import TypeAdapter from pydantic_ai import DeferredToolRequests, DeferredToolResults, ToolDenied from pydantic_ai.capabilities import AbstractCapability, Toolset, WrapperCapability from pydantic_ai.messages import ModelMessagesTypeAdapter, ModelResponse, ToolCallPart +from pydantic_ai.toolsets import DynamicToolset from pydantic_ai.toolsets.abstract import AbstractToolset from pydantic_ai.usage import RunUsage -from airflow.providers.common.ai.durable.capability import DURABLE_UNIT_TOOLSETS, AirflowDurability +from airflow.providers.common.ai.durable.capability import ( + DURABLE_UNIT_TOOLSETS, + AirflowDurability, + chosen_toolset_id, + give_durable_ids, +) from airflow.providers.common.ai.durable.journal import DurableJournal, build_task_storage, journal_scope from airflow.providers.common.ai.exceptions import ( ToolApprovalAlreadyRequestedError, @@ -106,9 +112,11 @@ if TYPE_CHECKING: from pydantic_ai.capabilities import AgentCapability from pydantic_ai.messages import ModelMessage from pydantic_ai.output import OutputSpec + from pydantic_ai.toolsets import ToolsetFunc from pydantic_ai.usage import UsageLimits from airflow.providers.common.ai.durable.base import DurableStorageProtocol + from airflow.providers.common.ai.durable.capability import GivenIds from airflow.providers.common.compat.sdk import TaskInstanceKey from airflow.sdk import Context from airflow.sdk.execution_time.context import TaskStateStoreAccessor @@ -378,9 +386,11 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa On Airflow >= 3.3 the journal is kept in the AIP-103 task state store, so no extra configuration is needed. On older Airflow versions it is persisted to ObjectStorage and requires ``[common.ai] durable_cache_path`` to be set. - Toolsets pydantic-ai runs as durable steps (``FunctionToolset``, - ``MCPToolset``, ``DynamicToolset``) need a unique ``id``, and so do - capabilities with ``@durable_operation`` methods. A tool with the metadata + A pydantic-ai ``MCPToolset`` needs a unique ``id``, and so do capabilities + with ``@durable_operation`` methods. A ``FunctionToolset`` or a function that + builds a toolset, passed in ``toolsets`` or ``agent_params["toolsets"]``, needs + none, and neither does a ``FunctionToolset`` held directly by a ``Toolset`` + capability (:ref:`durable-toolset-ids`). A tool with the metadata ``{"airflow": False}`` runs again on every attempt instead of replaying (:ref:`durable-tool-opt-out`). Cannot be combined with a ``SandboxToolset`` (raises), attached or @@ -635,42 +645,42 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa """ Refuse a toolset whose steps durable replay would have no name for. - pydantic-ai records each call of a ``FunctionToolset``, ``MCPToolset`` or - ``DynamicToolset`` under the toolset's ``id``, and refuses to run one without it. - The id is what a retry matches recorded calls by, so it has to come from the Dag - author: one derived from the toolset's position would point a retry at another - toolset's results as soon as the list is reordered. Checked here so the Dag fails - to load rather than at the first run. + pydantic-ai names the durable steps of a ``FunctionToolset``, ``MCPToolset`` or + ``DynamicToolset`` after the toolset's ``id``, and refuses to run one without it. The + toolsets ``_build_agent`` hands the agent are given ids here as it gives them + (``give_durable_ids``), and a leaf still without one is refused: an ``MCPToolset``, whose + server URL or command can carry credentials, or a function without a qualified name. + Checked here so the Dag fails to load rather than at the first run. """ - missing = { - type(leaf).__name__ - for toolset in self._declared_toolsets() - for leaf in iter_toolsets(toolset) - if isinstance(leaf, DURABLE_UNIT_TOOLSETS) and leaf.id is None - } - # A function that builds a toolset per run becomes a DynamicToolset with no id. - if any( - not isinstance(toolset, AbstractToolset) - for toolset in (*(self.toolsets or []), *(self.agent_params.get("toolsets") or [])) - ): - missing.add("DynamicToolset") - if kinds := sorted(missing): - raise ValueError( - f"durable=True needs a unique id on every {' and '.join(kinds)}, such as " - f"{kinds[0]}(..., id='orders'): durable replay records each tool call under it. " - "Set it on the toolset itself, and pass a function that builds a toolset as " - "DynamicToolset(function, id=...); the id of a Toolset capability around a toolset " - "does not reach it." - ) + given: GivenIds = {} + leaves = [ + leaf + for toolset in self._running_toolsets() + for leaf in iter_toolsets(give_durable_ids(toolset, given)) + if isinstance(leaf, DURABLE_UNIT_TOOLSETS) + ] # One instance listed twice is one toolset to pydantic-ai; two instances sharing an id - # would have their recorded calls filed under one name. + # would have their recorded calls filed under one name. Each listing of a function + # becomes a DynamicToolset of its own, so two lambdas in one module share an id. first_by_id: dict[str, AbstractToolset[Any]] = {} duplicates: list[str] = [] - for toolset in self._running_toolsets(): - for leaf in iter_toolsets(toolset): - if isinstance(leaf, DURABLE_UNIT_TOOLSETS) and leaf.id is not None: - if first_by_id.setdefault(leaf.id, leaf) is not leaf and leaf.id not in duplicates: - duplicates.append(leaf.id) + for leaf in leaves: + if (toolset_id := chosen_toolset_id(leaf)) is None and isinstance(leaf, DynamicToolset): + raise ValueError( + "durable=True names the steps of a toolset built per run after the function that " + f"builds it, but {leaf.toolset_func!r} has no __qualname__. Pass it as " + "DynamicToolset(function, id=...)." + ) + if toolset_id is None: + raise ValueError( + "durable=True needs a unique id on every MCPToolset, such as " + "MCPToolset(..., id='weather'): durable replay records each tool call under it, and " + "none is derived from the server's URL or command, which can carry credentials. Set " + "it on the toolset itself; the id of a Toolset capability around a toolset does not " + "reach it." + ) + if first_by_id.setdefault(toolset_id, leaf) is not leaf and toolset_id not in duplicates: + duplicates.append(toolset_id) if duplicates: raise ValueError( f"durable=True needs a unique id on each toolset, but more than one toolset has the id " @@ -826,13 +836,18 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa # masked here as well because toolset wrappers a capability adds (get_wrapper_toolset, # as ToolLoggingCapability does) run below the capability hooks and would otherwise see # the unmasked result. The durable journal masks what it records on its own. + given_ids: GivenIds = {} + + def prepare(toolset: AbstractToolset[Any] | ToolsetFunc[Any]) -> AbstractToolset[Any]: + return ensure_masked(give_durable_ids(toolset, given_ids) if self.durable else toolset) + if self.toolsets: - toolsets: list[AbstractToolset] = [ensure_masked(ts) for ts in self.toolsets] + toolsets: list[AbstractToolset] = [prepare(ts) for ts in self.toolsets] extra_kwargs["toolsets"] = toolsets elif extra_kwargs.get("toolsets"): - extra_kwargs["toolsets"] = [ensure_masked(ts) for ts in extra_kwargs["toolsets"]] + extra_kwargs["toolsets"] = [prepare(ts) for ts in extra_kwargs["toolsets"]] capabilities = [ - _replace_capability_toolset(capability, ensure_masked) + _replace_capability_toolset(capability, prepare) for capability in self.capabilities or passed_through or [] ] if self.durable: @@ -898,15 +913,14 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa candidates.append(toolset) return candidates - def _running_toolsets(self) -> list[AbstractToolset[Any]]: + def _running_toolsets(self) -> list[AbstractToolset[Any] | ToolsetFunc[Any]]: """ - Return the toolsets ``_build_agent`` hands the agent. + Return the toolsets ``_build_agent`` hands the agent, including functions that build one per run. Unlike ``_declared_toolsets``, ``agent_params["toolsets"]`` counts only when ``toolsets=`` is empty, because ``_build_agent`` ignores it otherwise. """ - listed = self.toolsets or self.agent_params.get("toolsets") or [] - candidates = [toolset for toolset in listed if isinstance(toolset, AbstractToolset)] + candidates = list(self.toolsets or self.agent_params.get("toolsets") or []) for capability in self._declared_capabilities: if (toolset := _resolve_capability_toolset(capability)) is not None: candidates.append(toolset) diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_capability.py b/providers/common/ai/tests/unit/common/ai/durable/test_capability.py index c9f9b6022c8..e951c59bf43 100644 --- a/providers/common/ai/tests/unit/common/ai/durable/test_capability.py +++ b/providers/common/ai/tests/unit/common/ai/durable/test_capability.py @@ -26,6 +26,7 @@ from __future__ import annotations import asyncio import base64 import dataclasses +import functools from collections.abc import Iterable from typing import TYPE_CHECKING, Any from unittest import mock @@ -35,6 +36,7 @@ from pydantic import BaseModel from pydantic_ai import Agent, CancellationToken, RunContext from pydantic_ai.capabilities import AbstractCapability, SetToolMetadata, durable_operation from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior +from pydantic_ai.mcp import MCPToolset from pydantic_ai.messages import ( ModelMessage, ModelResponse, @@ -45,11 +47,17 @@ from pydantic_ai.messages import ( ) from pydantic_ai.models.function import AgentInfo, FunctionModel from pydantic_ai.toolsets import DynamicToolset, FunctionToolset +from pydantic_ai.toolsets.wrapper import WrapperToolset from pydantic_ai.usage import RequestUsage, RunUsage, UsageLimits from airflow.providers.common.ai.durable import AirflowDurability from airflow.providers.common.ai.durable.base import RUNS_KEY -from airflow.providers.common.ai.durable.capability import _JournalCodec +from airflow.providers.common.ai.durable.capability import ( + GivenIds, + _JournalCodec, + callable_toolset_id, + give_durable_ids, +) from airflow.providers.common.ai.durable.journal import DurableJournal, Unrecordable, journal_scope from airflow.providers.common.ai.exceptions import DurableJournalError from airflow.providers.common.ai.utils.toolset_base import AirflowToolset, ensure_masked @@ -1086,6 +1094,173 @@ class TestToolOptOut: assert calls["query"] == 1 +def _query_tools(calls: Calls, *extra: Any, instructions: str | None = None) -> FunctionToolset[None]: + """A ``FunctionToolset`` without an id, with a ``query`` tool and any ``extra`` tools.""" + + def query(sql: str) -> str: + calls.bump("query") + return "3 rows" + + return FunctionToolset(tools=[query, *extra], instructions=instructions) + + +def _build_tools(ctx: RunContext[None]) -> FunctionToolset[None]: + """A function that builds a toolset for each run, as a Dag file defines one.""" + return FunctionToolset() + + +class _ToolsBuilder: + def __call__(self, ctx: RunContext[None]) -> FunctionToolset[None]: + return FunctionToolset() + + +class TestToolsetIds: + """Toolsets without an ``id`` replay under the ids ``give_durable_ids`` gives them, as in AgentOperator.""" + + @staticmethod + def tool_step_names(storage) -> list[str]: + return [step["name"] for step in storage.steps() if step["kind"] == "tool"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("built_per_run", "instructions"), + [ + pytest.param(False, None, id="function-toolset"), + # pydantic-ai builds instruction ids from the toolset's id: it rejects a colon in it, and the + # request fingerprint hashes them, so the id has to be the same on every attempt. + pytest.param(False, "Count rows with SQL.", id="function-toolset-with-instructions"), + pytest.param(True, None, id="function-building-one"), + ], + ) + async def test_a_toolset_without_an_id_replays_on_retry( + self, memory_storage, built_per_run, instructions + ): + calls = Calls() + fail_final = [True] + + def build_tools(ctx: RunContext[None]) -> FunctionToolset[None]: + return _query_tools(calls) + + def build() -> Agent[None, str]: + toolset = build_tools if built_per_run else _query_tools(calls, instructions=instructions) + return Agent( + FunctionModel(tool_then_answer(calls, "query", fail_final=fail_final)), + name="analyst", + toolsets=[give_durable_ids(toolset, {})], + capabilities=[AirflowDurability()], + ) + + with pytest.raises(RuntimeError, match="worker died"): + await attempt(memory_storage, build()) + names = self.tool_step_names(memory_storage) + fail_final[0] = False + result = await attempt(memory_storage, build()) + + assert result.output == "answer: 3 rows" + assert names == [ + f"analyst__dynamic_toolset__{__name__}.{build_tools.__qualname__}.call_tool:query" + if built_per_run + else "analyst__function_toolset.call_tool:query" + ] + assert (calls["query"], calls["model"]) == (1, 3) + + @pytest.mark.asyncio + async def test_adding_a_tool_between_attempts_keeps_the_first_tools_step_name(self, memory_storage): + """The retry runs again from the first request, which now offers a new tool, but names its calls alike.""" + calls = Calls() + + def forecast(city: str) -> str: + return "sunny" + + names = [] + for tools in ((), (forecast,)): + agent = Agent( + FunctionModel(tool_then_answer(calls, "query", fail_final=[True])), + name="analyst", + toolsets=[give_durable_ids(_query_tools(calls, *tools), {})], + capabilities=[AirflowDurability()], + ) + with pytest.raises(RuntimeError, match="worker died"): + await attempt(memory_storage, agent) + names.append(self.tool_step_names(memory_storage)) + + assert names == [["analyst__function_toolset.call_tool:query"]] * 2 + assert calls["query"] == 2 + + @pytest.mark.asyncio + async def test_an_id_the_author_chose_keeps_its_step_names(self, memory_storage): + """The step names journals were recorded under before Airflow gave toolsets ids.""" + calls = Calls() + toolset = FunctionToolset(tools=list(_query_tools(calls).tools.values()), id="db") + agent = Agent( + FunctionModel(tool_then_answer(calls, "query", fail_final=[True])), + name="analyst", + toolsets=[give_durable_ids(toolset, {})], + capabilities=[AirflowDurability()], + ) + + with pytest.raises(RuntimeError, match="worker died"): + await attempt(memory_storage, agent) + + assert self.tool_step_names(memory_storage) == ["analyst__function_toolset__db.call_tool:query"] + + def test_gives_ids_to_copies_and_one_copy_to_a_toolset_listed_twice(self): + shared = FunctionToolset() + given: GivenIds = {} + + prefixed = give_durable_ids(shared.prefixed("a"), given) + plain = give_durable_ids(shared, given) + + assert isinstance(prefixed, WrapperToolset) + assert prefixed.wrapped is plain + assert plain.id == "<airflow-function_toolset-0>" + assert shared.id is None + + def test_numbers_the_function_toolsets_it_gives_ids_in_order(self): + given: GivenIds = {} + + ids = [give_durable_ids(FunctionToolset(), given).id for _ in range(2)] + + assert ids == ["<airflow-function_toolset-0>", "<airflow-function_toolset-1>"] + + @pytest.mark.parametrize( + "toolset", + [ + pytest.param(FunctionToolset(id="db").prefixed("a"), id="function-toolset-with-id"), + pytest.param(DynamicToolset(_build_tools, id="tools"), id="dynamic-toolset-with-id"), + pytest.param(MCPToolset("http://mcp.example.com/mcp";), id="mcp-toolset"), + ], + ) + def test_leaves_a_toolset_it_cannot_or_need_not_give_an_id_as_it_is(self, toolset): + assert give_durable_ids(toolset, {}) is toolset + + def test_a_dynamic_toolset_without_an_id_gets_its_functions_name(self): + given = give_durable_ids(DynamicToolset(_build_tools, per_run_step=False), {}) + + assert isinstance(given, DynamicToolset) + assert (given.id, given.per_run_step) == (f"{__name__}._build_tools", False) + + +class TestCallableToolsetId: + def test_leaves_out_the_path_hash_airflow_puts_in_a_dag_files_module(self): + def build_tools(ctx: RunContext[None]) -> FunctionToolset[None]: + return FunctionToolset() + + build_tools.__module__ = f"unusual_prefix_{'0123456789abcdef' * 2}01234567_orders_dag" + + assert callable_toolset_id(build_tools) == f"orders_dag.{build_tools.__qualname__}" + + @pytest.mark.parametrize( + "function", + [ + pytest.param(functools.partial(_build_tools), id="partial"), + pytest.param(_ToolsBuilder(), id="callable-instance"), + ], + ) + def test_none_for_a_callable_without_a_qualified_name(self, function): + assert callable_toolset_id(function) is None + + class Answer(BaseModel): rows: int diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py index 2e9a92f1340..3eb80fd5a07 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py @@ -18,6 +18,7 @@ from __future__ import annotations import asyncio import dataclasses +import functools import inspect import json import logging @@ -46,6 +47,7 @@ from pydantic_ai.capabilities import ( WebSearch, ) from pydantic_ai.exceptions import ApprovalRequired, ModelRetry, ToolFailed, UsageLimitExceeded +from pydantic_ai.mcp import MCPToolset as PydanticAIMCPToolset from pydantic_ai.messages import ( BinaryContent, ModelMessage, @@ -60,6 +62,7 @@ from pydantic_ai.messages import ( ) from pydantic_ai.models.function import AgentInfo, FunctionModel from pydantic_ai.models.wrapper import WrapperModel +from pydantic_ai.toolsets import DynamicToolset from pydantic_ai.toolsets.abstract import AbstractToolset from pydantic_ai.toolsets.combined import CombinedToolset from pydantic_ai.toolsets.function import FunctionToolset @@ -1921,6 +1924,11 @@ def _refund(order_id: int) -> str: return f"refunded order {order_id}" +def _build_forecast_tools(ctx: RunContext) -> FunctionToolset: + """A function that builds a toolset for each run.""" + return FunctionToolset() + + def _tool_then_answer(fail_after_tool: list[bool], live: dict[str, int]): """A model that calls ``my_tool`` once, then answers; it fails once after the tool if asked to.""" @@ -2045,20 +2053,112 @@ class TestAgentOperatorDurable: ), ], ) - def test_durable_refuses_a_toolset_without_an_id(self, kwargs): - with pytest.raises(ValueError, match="durable=True needs a unique id on every FunctionToolset"): - AgentOperator(task_id="t", prompt="p", llm_conn_id="c", durable=True, **kwargs) + def test_durable_accepts_a_function_toolset_without_an_id(self, kwargs): + AgentOperator(task_id="t", prompt="p", llm_conn_id="c", durable=True, **kwargs) + + @pytest.mark.parametrize( + "toolset", + [ + pytest.param(functools.partial(_build_forecast_tools), id="partial"), + pytest.param(DynamicToolset(functools.partial(_build_forecast_tools)), id="dynamic-toolset"), + ], + ) + def test_durable_refuses_a_function_without_a_qualified_name(self, toolset): + with pytest.raises(ValueError, match=r"functools.partial\(.*\) has no __qualname__"): + AgentOperator(task_id="t", prompt="p", llm_conn_id="c", durable=True, toolsets=[toolset]) - def test_durable_refuses_a_function_that_builds_a_toolset(self): - with pytest.raises(ValueError, match="needs a unique id on every DynamicToolset"): + def test_durable_refuses_an_mcp_toolset_without_an_id(self): + with pytest.raises(ValueError, match="needs a unique id on every MCPToolset"): AgentOperator( task_id="t", prompt="p", llm_conn_id="c", durable=True, - toolsets=[lambda ctx: FunctionToolset()], + toolsets=[PydanticAIMCPToolset("http://mcp.example.com/mcp";)], ) + @pytest.mark.parametrize( + "toolsets", + [ + pytest.param([lambda ctx: FunctionToolset(), lambda ctx: FunctionToolset()], id="two-lambdas"), + pytest.param([_build_forecast_tools, _build_forecast_tools], id="one-function-twice"), + pytest.param( + [_build_forecast_tools, DynamicToolset(_build_forecast_tools)], id="function-and-dynamic" + ), + ], + ) + def test_durable_refuses_two_functions_with_the_same_qualified_name(self, toolsets): + with pytest.raises(ValueError, match="more than one toolset has the id"): + AgentOperator(task_id="t", prompt="p", llm_conn_id="c", durable=True, toolsets=toolsets) + + @pytest.mark.parametrize("passed_as", ["toolsets", "agent_params", "toolset_capability"]) + def test_durable_replays_toolsets_that_have_no_id(self, passed_as): + """Id-less FunctionToolsets, one with instructions, and a function that builds a toolset replay on retry.""" + storage = MemoryStorage() + live = {"model": 0, "lookup": 0, "forecast": 0} + fail_final = [True] + + def lookup(order_id: int) -> str: + live["lookup"] += 1 + return "shipped" + + def build_forecast_tools(ctx: RunContext) -> FunctionToolset: + def forecast(city: str) -> str: + live["forecast"] += 1 + return "sunny" + + return FunctionToolset(tools=[forecast]) + + def model_fn(messages, info): + live["model"] += 1 + match sum(isinstance(m, ModelResponse) for m in messages): + case 0: + return ModelResponse(parts=[ToolCallPart("lookup", {"order_id": 7})]) + case 1: + return ModelResponse(parts=[ToolCallPart("forecast", {"city": "Oslo"})]) + if fail_final[0]: + fail_final[0] = False + raise RuntimeError("transient model failure") + return ModelResponse(parts=[TextPart(content="done")]) + + def refund_policy() -> str: + return "30 days" + + for try_number in (1, 2): + toolsets = [ + FunctionToolset(tools=[lookup], instructions="Orders ship from Oslo."), + FunctionToolset(tools=[refund_policy]), + build_forecast_tools, + ] + op = AgentOperator( + task_id="t", + prompt="hi", + llm_conn_id="c", + durable=True, + enable_tool_logging=False, + **{ + "toolsets": {"toolsets": toolsets}, + "agent_params": {"agent_params": {"toolsets": toolsets}}, + # The FunctionToolset with instructions is held by a Toolset capability instead. + "toolset_capability": { + "toolsets": toolsets[1:], + "capabilities": [Toolset(toolsets[0])], + }, + }[passed_as], + ) + op.llm_hook = MagicMock(spec=["create_agent"]) + op.llm_hook.create_agent.side_effect = lambda **kw: Agent(FunctionModel(model_fn), **kw) + context = _make_context(ti=_make_ti(id=f"ti-{try_number}", try_number=try_number)) + with ( + patch.object(AgentOperator, "_build_durable_storage", autospec=True, return_value=storage), + pytest.raises(RuntimeError, match="transient") if try_number == 1 else nullcontext(), + ): + output = op.execute(context=context) + + assert output == "done" + # Attempt 2 replays both model requests and both tool calls; only the failed request runs again. + assert live == {"model": 4, "lookup": 1, "forecast": 1} + def test_non_durable_accepts_a_toolset_without_an_id(self): AgentOperator(task_id="t", prompt="p", llm_conn_id="c", toolsets=[FunctionToolset()]) @@ -2099,9 +2199,15 @@ class TestAgentOperatorDurable: agent_params={"toolsets": [FunctionToolset(id="orders")]}, ) - def test_durable_accepts_one_toolset_passed_twice(self): - orders = FunctionToolset(id="orders") - + @pytest.mark.parametrize( + "orders", + [ + pytest.param(FunctionToolset(id="orders"), id="with-id"), + pytest.param(FunctionToolset(), id="function-toolset-without-id"), + pytest.param(DynamicToolset(_build_forecast_tools), id="dynamic-toolset-without-id"), + ], + ) + def test_durable_accepts_one_toolset_passed_twice(self, orders): AgentOperator( task_id="t", prompt="p", @@ -2111,6 +2217,26 @@ class TestAgentOperatorDurable: capabilities=[Toolset(orders)], ) + def test_durable_gives_a_toolset_passed_twice_one_id(self): + orders = FunctionToolset() + op = AgentOperator( + task_id="t", + prompt="p", + llm_conn_id="c", + durable=True, + toolsets=[orders], + capabilities=[Toolset(orders)], + ) + op.llm_hook = MagicMock(spec=["create_agent"]) + + op._build_agent() + + kwargs = op.llm_hook.create_agent.call_args.kwargs + listed, in_capability = kwargs["toolsets"][0].wrapped, kwargs["capabilities"][0].toolset.wrapped + assert listed is in_capability + assert listed.id is not None + assert orders.id is None + def test_durable_fails_a_tool_that_raises_approval_required_with_unsupported_tool_deferral_error(self): def refund() -> str: raise ApprovalRequired
