This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch durable-tool-approval in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 3a43b649b75e2749be69a6e7c9dbec2d3c101540 Author: Kaxil Naik <[email protected]> AuthorDate: Tue Nov 10 18:27:15 2026 +0000 Let durable agent tasks run chosen tools on every attempt instead of replaying them A tool whose answer must be current (today's date, a job's state) or whose effect does not outlive the attempt (a temporary file) can now opt out of replay with the tool metadata {"airflow": False}, set on the tool or with pydantic-ai's SetToolMetadata capability for toolsets such as SQLToolset. On a retry the tool runs live while the model requests and other tool calls still replay. The opt-out is read by AirflowDurability's operation backend through RoleBasedOperationConfig(resolve_tool=...) for pydantic-ai's function and dynamic toolsets, and by the journaling wrapper for Airflow's own toolsets. An opted-out call keeps its position in the journal and its failures still count toward asking the model again. The key is left out of the model request fingerprint, so toggling it between attempts does not invalidate the journal. pydantic-ai keeps refusing the opt-out on its own MCPToolset. The pydantic-ai-slim floor moves to 2.46.0: earlier releases drop metadata that a wrapper such as SetToolMetadata adds before the call reaches the toolset, which would silently replay a tool the user opted out. --- providers/common/ai/README.rst | 12 +- providers/common/ai/docs/decision_models.rst | 7 +- providers/common/ai/docs/durable_execution.rst | 73 +++++++- providers/common/ai/docs/features.rst | 3 +- providers/common/ai/docs/index.rst | 12 +- providers/common/ai/pyproject.toml | 21 +-- .../airflow/providers/common/ai/durable/base.py | 3 + .../providers/common/ai/durable/capability.py | 59 ++++-- .../providers/common/ai/durable/fingerprint.py | 25 ++- .../airflow/providers/common/ai/durable/journal.py | 7 +- .../ai/example_dags/example_agent_durable.py | 38 +++- .../airflow/providers/common/ai/operators/agent.py | 4 +- .../unit/common/ai/durable/test_capability.py | 198 ++++++++++++++++++++- .../unit/common/ai/durable/test_fingerprint.py | 18 +- uv.lock | 14 +- 15 files changed, 423 insertions(+), 71 deletions(-) diff --git a/providers/common/ai/README.rst b/providers/common/ai/README.rst index cd4d21bf72f..02a0ed164f9 100644 --- a/providers/common/ai/README.rst +++ b/providers/common/ai/README.rst @@ -57,7 +57,7 @@ PIP package Version required ``apache-airflow-providers-common-compat`` ``>=1.22.0`` ``apache-airflow-providers-standard`` ``>=1.20.0`` ``httpx2`` ``>=2.7`` -``pydantic-ai-slim`` ``>=2.36.0`` +``pydantic-ai-slim`` ``>=2.46.0`` ``structlog`` ``>=24.2.0`` ========================================== ================== @@ -88,12 +88,12 @@ Optional dependencies ==================== ======================================================================================================================================= Extra Dependencies ==================== ======================================================================================================================================= -``anthropic`` ``pydantic-ai-slim[anthropic]>=2.36.0``, ``anthropic>=1.1.0,<2`` -``bedrock`` ``pydantic-ai-slim[bedrock]>=2.36.0`` -``google`` ``pydantic-ai-slim[google]>=2.36.0`` -``openai`` ``pydantic-ai-slim[openai]>=2.36.0``, ``openai>=2.47.0`` +``anthropic`` ``pydantic-ai-slim[anthropic]>=2.46.0``, ``anthropic>=1.1.0,<2`` +``bedrock`` ``pydantic-ai-slim[bedrock]>=2.46.0`` +``google`` ``pydantic-ai-slim[google]>=2.46.0`` +``openai`` ``pydantic-ai-slim[openai]>=2.46.0``, ``openai>=2.47.0`` ``typesafe`` ``typesafe-sdk>=0.6.0`` -``mcp`` ``pydantic-ai-slim[mcp]>=2.36.0`` +``mcp`` ``pydantic-ai-slim[mcp]>=2.46.0`` ``modal`` ``apache-airflow-providers-modal``, ``modal>=1.5.2`` ``opensandbox`` ``opensandbox>=1.1.0`` ``boat`` ``boat-sdk>=1.9.3`` diff --git a/providers/common/ai/docs/decision_models.rst b/providers/common/ai/docs/decision_models.rst index a84c895ab90..88ad02f6623 100644 --- a/providers/common/ai/docs/decision_models.rst +++ b/providers/common/ai/docs/decision_models.rst @@ -42,8 +42,7 @@ model, so a model id is the whole integration. - Needs - Runs * - ``typesafe:`` - - ``pydantic-ai-slim`` 2.45.0 or later (2.46.0 for the examples on this page, which - describe options with ``UseEnumMemberDocstrings``) and the provider's ``typesafe`` extra + - The provider's ``typesafe`` extra - `TypeSafe <https://typesafe.ai>`__'s hosted Jev. * - ``system-one:`` - ``pydantic-ai-slim`` 2.53.0 or later, no extra @@ -55,8 +54,8 @@ model, so a model id is the whole integration. it. `pydantic-ai's System One page <https://pydantic.dev/docs/ai/models/system-one/>`__ lists more. -Both prefixes need a newer ``pydantic-ai-slim`` than the floor the provider's other extras -set, so check the release you have installed. +``system-one:`` needs a newer ``pydantic-ai-slim`` than the provider's minimum, so check the +release you have installed. Setup ----- diff --git a/providers/common/ai/docs/durable_execution.rst b/providers/common/ai/docs/durable_execution.rst index fa402d0ef4e..679be441518 100644 --- a/providers/common/ai/docs/durable_execution.rst +++ b/providers/common/ai/docs/durable_execution.rst @@ -210,11 +210,13 @@ The journal records: - model requests, including streamed ones; - calls to tools from any toolset, including the provider's own toolsets, MCP servers, - toolsets built per run, and toolsets supplied by capabilities; + toolsets built per run, and toolsets supplied by capabilities, except the tools you + :ref:`set to run on every attempt <durable-tool-opt-out>`; - tool discovery for MCP and per-run toolsets; - ``@durable_operation`` methods of other capabilities, such as the spend accounting of pydantic-ai-harness ``SpendLimits`` and the summarization of ``SummarizingCompaction``, - so a retry neither charges the spend twice nor pays for a summary again. + so a retry neither charges the spend twice nor pays for a summary again, unless the + summary follows a tool set to run on every attempt. A tool that raises ``ModelRetry``, ``ToolFailed``, ``ApprovalRequired`` or ``CallDeferred`` replays as the same signal. ``AgentOperator`` does not support tool @@ -398,10 +400,6 @@ reject the duplicate. Tool results are masked by Airflow's secret masker before they are recorded, so a connection password in a result reaches the journal as ``***``. -A toolset can opt out of replay. ``ManagedAgentToolset`` does by default -(``replayable=False``), because a managed agent acts on a system Airflow cannot observe: -its calls run again on every attempt, while the steps around them still replay. - A result has to survive a JSON round trip to be recorded. Bytes a tool returns, on their own or as values anywhere inside dicts and lists, are recorded as base64, and a retry replays the same bytes. Only ``bytes`` is converted. ``bytearray``, ``memoryview``, a @@ -414,6 +412,69 @@ such as an object with no JSON form, fails the task with ``DurableJournalError`` retry, since it would fail the same way every time. A write the storage refuses or fails does not fail the task: the task log names the step, and a retry runs it again. +.. _durable-tool-opt-out: + +Run a tool on every attempt +--------------------------- + +Some tools must run again on a retry rather than replay: one whose answer has to be +current, such as today's date or the state of a job, and one whose effect does not outlive +the attempt, such as a file in a temporary directory. Give such a tool the metadata +``{"airflow": False}``. For a tool in a toolset you do not build yourself, such as +``SQLToolset``, select it by name with pydantic-ai's ``SetToolMetadata`` capability: + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py + :language: python + :start-after: [START howto_agent_durable_tool_opt_out] + :end-before: [END howto_agent_durable_tool_opt_out] + +On a retry, ``today`` and every ``query`` call run again, and the model requests and the +other tool calls still replay. A live call that returns what it returned on the last +attempt leaves the model requests and tool calls after it replaying. One that returns +something new changes the conversation the model sees, so the model requests after it, +and the tool calls they make, run again. A capability that writes into the conversation +right after the live call (a compaction summary, say) runs again too: its +``@durable_operation`` has no request to compare, so it replays only after a step that +replayed. The summary it writes can then make the steps after it run again. + +You can turn the opt-out on or off between attempts. The call keeps its place in the +journal and the setting is left out when a model request is compared, so later steps still +replay. Failures still count, too. Two in a row with the same arguments, and the next retry +asks the model again. + +Where the setting goes depends on the toolset: + +.. list-table:: + :header-rows: 1 + :widths: 35 65 + + * - Toolset + - How to opt a tool out + * - ``FunctionToolset``, tools passed to the agent + - ``metadata={"airflow": False}`` on the tool, or ``SetToolMetadata``. + * - ``SQLToolset``, ``HookToolset``, the provider's ``MCPToolset``, and any other toolset + - ``SetToolMetadata``. On a toolset built on ``AirflowToolset``, setting + ``replayable = False`` opts all its tools out; ``ManagedAgentToolset`` does by + default, because a managed agent acts on a system Airflow cannot observe. + * - ``DynamicToolset`` + - ``SetToolMetadata``. Its tool list replays from the journal, so metadata set inside + the toolset can keep the value an earlier attempt saw. + * - ``pydantic_ai.mcp.MCPToolset`` + - Not supported; pydantic-ai raises the error below when an opted-out tool is called. + +The error from pydantic-ai's ``MCPToolset``: + +.. code-block:: text + + pydantic_ai.exceptions.UserError: Airflow durable config for MCP tool 'get_forecast' has been explicitly set to `False` (durable execution disabled), but MCP tools perform I/O and cannot run outside a durable step. Remove the metadata so the call stays durable. + +Any value other than ``True``, ``False`` or ``None`` is refused when the tool is called. The +attempt fails, as do its retries: + +.. code-block:: text + + pydantic_ai.exceptions.UserError: Tool 'today' has invalid 'airflow' metadata: expected True or False, got str. + Limitations ----------- diff --git a/providers/common/ai/docs/features.rst b/providers/common/ai/docs/features.rst index 9ecee96b711..8b5969c348b 100644 --- a/providers/common/ai/docs/features.rst +++ b/providers/common/ai/docs/features.rst @@ -87,8 +87,7 @@ Choose a model - Uses a model that cannot write text, such as TypeSafe's hosted Jev or Ollama's decision models, to answer a typed question with one value from a set you define, plus a confidence. - - Jev: ``typesafe``, and install ``pydantic-ai-slim`` 2.45+; others: install - ``pydantic-ai-slim`` 2.53+ + - Jev: ``typesafe``; others: install ``pydantic-ai-slim`` 2.53+ * - :doc:`LlamaIndex models <hooks/llamaindex>` - Builds LlamaIndex language and embedding models from a connection. - ``llamaindex``, and install the vendor's package for vendors other than OpenAI diff --git a/providers/common/ai/docs/index.rst b/providers/common/ai/docs/index.rst index cc627cf2842..c0926af6387 100644 --- a/providers/common/ai/docs/index.rst +++ b/providers/common/ai/docs/index.rst @@ -189,7 +189,7 @@ PIP package Version required ``apache-airflow-providers-common-compat`` ``>=1.22.0`` ``apache-airflow-providers-standard`` ``>=1.20.0`` ``httpx2`` ``>=2.7`` -``pydantic-ai-slim`` ``>=2.36.0`` +``pydantic-ai-slim`` ``>=2.46.0`` ``structlog`` ``>=24.2.0`` ========================================== ================== @@ -228,12 +228,12 @@ Install them when installing from PyPI. For example: ==================== ======================================================================================================================================= Extra Dependencies ==================== ======================================================================================================================================= -``anthropic`` ``pydantic-ai-slim[anthropic]>=2.36.0``, ``anthropic>=1.1.0,<2`` -``bedrock`` ``pydantic-ai-slim[bedrock]>=2.36.0`` -``google`` ``pydantic-ai-slim[google]>=2.36.0`` -``openai`` ``pydantic-ai-slim[openai]>=2.36.0``, ``openai>=2.47.0`` +``anthropic`` ``pydantic-ai-slim[anthropic]>=2.46.0``, ``anthropic>=1.1.0,<2`` +``bedrock`` ``pydantic-ai-slim[bedrock]>=2.46.0`` +``google`` ``pydantic-ai-slim[google]>=2.46.0`` +``openai`` ``pydantic-ai-slim[openai]>=2.46.0``, ``openai>=2.47.0`` ``typesafe`` ``typesafe-sdk>=0.6.0`` -``mcp`` ``pydantic-ai-slim[mcp]>=2.36.0`` +``mcp`` ``pydantic-ai-slim[mcp]>=2.46.0`` ``modal`` ``apache-airflow-providers-modal``, ``modal>=1.5.2`` ``opensandbox`` ``opensandbox>=1.1.0`` ``boat`` ``boat-sdk>=1.9.3`` diff --git a/providers/common/ai/pyproject.toml b/providers/common/ai/pyproject.toml index 9dcff8706a4..63c0ecde342 100644 --- a/providers/common/ai/pyproject.toml +++ b/providers/common/ai/pyproject.toml @@ -73,8 +73,10 @@ dependencies = [ # the SDK and stopped passing temperature/top_p/top_k as messages.create() kwargs, both # of which raise TypeError on 2.31 and earlier. The cost API this provider relies on # (RunUsage.cost, UsageLimits.cost_limit) landed earlier, in 2.23.0. Durable execution - # (AirflowDurability) builds on the durable backend API pydantic-ai-slim published in 2.36.0. - "pydantic-ai-slim>=2.36.0", + # (AirflowDurability) builds on the durable backend API pydantic-ai-slim published in 2.36.0; + # 2.46.0 is the first release whose CombinedToolset passes a tool call the tool definition it + # was handed, so tool metadata a wrapper sets (SetToolMetadata's replay opt-out) reaches it. + "pydantic-ai-slim>=2.46.0", # Airflow 3 brings structlog in through the Task SDK; Airflow 2 does not. 24.2.0 is the first # release whose render_to_log_kwargs hands ``stacklevel`` to stdlib logging under that name. "structlog>=24.2.0", @@ -92,19 +94,18 @@ dependencies = [ # openai>=3, so this floor only matters to an install that overrides that requirement. # anthropic is capped below 2 because pydantic-ai's Anthropic model (up to 2.59.0) calls # BetaMessageParam(...), which anthropic 2 turned into a Union that cannot be instantiated. -"anthropic" = ["pydantic-ai-slim[anthropic]>=2.36.0", "anthropic>=1.1.0,<2"] -"bedrock" = ["pydantic-ai-slim[bedrock]>=2.36.0"] -"google" = ["pydantic-ai-slim[google]>=2.36.0"] -"openai" = ["pydantic-ai-slim[openai]>=2.36.0", "openai>=2.47.0"] +"anthropic" = ["pydantic-ai-slim[anthropic]>=2.46.0", "anthropic>=1.1.0,<2"] +"bedrock" = ["pydantic-ai-slim[bedrock]>=2.46.0"] +"google" = ["pydantic-ai-slim[google]>=2.46.0"] +"openai" = ["pydantic-ai-slim[openai]>=2.46.0", "openai>=2.47.0"] # TypeSafe Jev: answers typed questions (pick one of N, a bounded score, a yes/no # probability) and cannot emit text, so it fits a branch or a classification but not a # summary. It arrives through pydantic-ai's own ``typesafe:`` provider, so the hook needs # no changes -- setting ``model_id`` or the connection's Model field is the integration. # The adapter itself ships inside pydantic-ai-slim from 2.45.0 (pydantic/pydantic-ai#8450); -# this extra only adds the pre-1.0 typesafe-sdk it talks to, which is why it is opt-in. It -# deliberately does not raise the pydantic-ai-slim floor to 2.45; the guide says 2.45+. +# this extra only adds the pre-1.0 typesafe-sdk it talks to, which is why it is opt-in. "typesafe" = ["typesafe-sdk>=0.6.0"] -"mcp" = ["pydantic-ai-slim[mcp]>=2.36.0"] +"mcp" = ["pydantic-ai-slim[mcp]>=2.46.0"] # Modal-hosted sandboxes for SandboxToolset, the backend to use where the worker host # cannot run microVMs -- Kubernetes, most notably. Needs no host-level installation, # only Modal credentials, which the Modal provider's `modal` connection supplies. @@ -194,7 +195,7 @@ dev = [ "apache-airflow-providers-standard", # Additional devel dependencies (do not remove this line and add extra development dependencies) "sqlglot>=30.0.0", - "pydantic-ai-slim[mcp]>=2.36.0", + "pydantic-ai-slim[mcp]>=2.46.0", "pydantic-ai-skills>=1.2.0,<2", "apache-airflow-providers-common-sql[datafusion]", "langchain>=1.0.0", diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/base.py b/providers/common/ai/src/airflow/providers/common/ai/durable/base.py index e9f1ff8165a..a85584ae10b 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/durable/base.py +++ b/providers/common/ai/src/airflow/providers/common/ai/durable/base.py @@ -45,6 +45,9 @@ RUN_ID_KEY = f"{DURABLE_KEY_PREFIX}run_id" # Lists every agent run any attempt started, so cleanup reaches the ones this attempt did not. RUNS_KEY = f"{DURABLE_KEY_PREFIX}runs" +# The tool metadata key that runs a tool live on every attempt: ``metadata={"airflow": False}``. +TOOL_CONFIG_KEY = "airflow" + @runtime_checkable class DurableStorageProtocol(Protocol): 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 75a5142e995..a6a8a3e3b68 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 @@ -25,7 +25,7 @@ from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass, field from functools import partial -from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias, cast +from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, cast from pydantic_ai.durable_exec import ( JSON_CODEC, @@ -46,13 +46,15 @@ from pydantic_ai.exceptions import ( ModelRetry, ToolFailed, UnexpectedModelBehavior, + UserError, ) from pydantic_ai.messages import ModelResponse, ToolReturn from pydantic_ai.tools import AgentDepsT -from pydantic_ai.toolsets import DynamicToolset, FunctionToolset +from pydantic_ai.toolsets import DynamicToolset, FunctionToolset, ToolsetTool from pydantic_ai.toolsets.wrapper import WrapperToolset from pydantic_core import PydanticSerializationError +from airflow.providers.common.ai.durable.base import TOOL_CONFIG_KEY from airflow.providers.common.ai.durable.fingerprint import fingerprint_model_request, fingerprint_tool_call from airflow.providers.common.ai.durable.journal import Unrecordable, current_journal from airflow.providers.common.ai.durable.replay_usage import ReplayUsageLedger, is_successful_tool_payload @@ -72,8 +74,8 @@ if TYPE_CHECKING: from pydantic_ai.models import Model, ModelRequestParameters from pydantic_ai.run import AgentRunResult from pydantic_ai.settings import ModelSettings - from pydantic_ai.tools import RunContext - from pydantic_ai.toolsets import AbstractToolset, ToolsetTool + from pydantic_ai.tools import RunContext, ToolDefinition + from pydantic_ai.toolsets import AbstractToolset from airflow.providers.common.ai.durable.journal import DurableRun, JournalStep, StepKind @@ -99,8 +101,33 @@ else: # 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" -_NO_CONFIG: RoleBasedOperationConfig[None] = RoleBasedOperationConfig( - model=None, event=None, capability=None, tool=None +# How an operation's step is run: looked up in the journal and recorded, or always run live. +_StepMode: TypeAlias = Literal["journal", "live"] + + +def _journals_tool(tool_def: ToolDefinition) -> bool: + """Whether calls of a tool are journaled; ``metadata={"airflow": False}`` runs them live on every attempt.""" + setting = (tool_def.metadata or {}).get(TOOL_CONFIG_KEY) + if setting is None: + return True + if isinstance(setting, bool): + return setting + raise UserError( + f"Tool {tool_def.name!r} has invalid {TOOL_CONFIG_KEY!r} metadata: expected True or False, " + f"got {type(setting).__name__}." + ) + + +def _tool_step_mode(operation_id: DurableOperationId, tool: object | None, tool_name: str) -> _StepMode: + # Never pydantic-ai's ``False``, which would take the call out of the backend: a live call still + # takes its position, so the steps after it keep theirs, and a failure is still the step's own. + if isinstance(tool, ToolsetTool) and not _journals_tool(tool.tool_def): + return "live" + return "journal" + + +_STEP_MODES: RoleBasedOperationConfig[_StepMode] = RoleBasedOperationConfig( + model="journal", event="journal", capability="journal", tool="journal", resolve_tool=_tool_step_mode ) @@ -253,6 +280,12 @@ class AirflowDurability(BaseDurabilityCapability[AgentDepsT]): soon as the run succeeds, so that clearing the task later starts it fresh: a task that runs several agents replays only the run that failed, and runs the ones before it again. + A tool that must run on every attempt, such as one that reads the current time or + creates a temporary table, opts out with ``metadata={"airflow": False}``; for tools of + a toolset you do not define, such as ``SQLToolset``, add the capability + ``SetToolMetadata(tools=[...], airflow=False)``. pydantic-ai's own ``MCPToolset`` refuses the + opt-out; the provider's ``MCPToolset`` accepts it. + 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. @@ -350,12 +383,12 @@ class AirflowDurability(BaseDurabilityCapability[AgentDepsT]): return model -class _AirflowOperationBackend(JournalCallableOperationBackend[None]): +class _AirflowOperationBackend(JournalCallableOperationBackend[_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=_NO_CONFIG + agent_name=durability.name, default_model_id=durability.default_model_id, config=_STEP_MODES ) self._durability = durability # How pydantic-ai decodes each operation's result, by operation. @@ -398,7 +431,7 @@ class _AirflowOperationBackend(JournalCallableOperationBackend[None]): name: str, body: Callable[[], Awaitable[object]], cache_key: tuple[object, ...], - config: None, + config: _StepMode, ) -> object: active = _ACTIVE_RUN.get() if active is None: @@ -414,6 +447,7 @@ class _AirflowOperationBackend(JournalCallableOperationBackend[None]): name, _tool_fingerprint(cache_key), valid_payload=partial(self._loads, operation_id), + replayable=config == "journal", ) if step.replayed: return step.payload @@ -528,8 +562,9 @@ class _JournaledToolset(WrapperToolset[Any]): fingerprint_tool_call(name, tool_args, ctx.tool_call_id), valid_payload=_is_journaled_tool_payload, # A toolset whose calls act on a system Airflow cannot observe, such as a managed - # agent, runs them again on every attempt. - replayable=leaf.replayable if isinstance(leaf, AirflowToolset) else True, + # agent, runs them again on every attempt, and so does a tool opted out by its metadata. + replayable=(not isinstance(leaf, AirflowToolset) or leaf.replayable) + and _journals_tool(tool.tool_def), ) if step.replayed: if isinstance(leaf, AirflowToolset): @@ -582,6 +617,8 @@ async def _run_and_record_tool(step: JournalStep, call: Callable[[], Awaitable[A def _record_tool_payload(step: JournalStep, payload: dict[str, Any]) -> None: + if not step.replayable: + return encoded = _JOURNAL_CODEC.dump(Any, payload) if isinstance(encoded, Unrecordable): step.cannot_record(encoded.error) diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py b/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py index 9904be74abf..b6a7297dd2f 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py +++ b/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py @@ -44,8 +44,9 @@ as base64. So a value that is not a JSON type but renders the same way on every attempt -- a ``datetime`` or ``Decimal`` tool argument, a dataclass in ``tool_choice`` -- still produces a usable fingerprint. The ``capability_id`` of a tool from a capability without an ``id`` is left out -of the request parameters, since it is random per run, and so is -``deferred_capability_ids`` (see ``_normalize_params``). +of the request parameters, since it is random per run, and so are +``deferred_capability_ids`` and the ``airflow`` key of a tool's metadata, which +only decides whether the tool replays (see ``_normalize_params``). That rendering is hashed as it is, with two exceptions. Both are found by dumping the payload a second time in pydantic's python mode, which has the same @@ -107,6 +108,7 @@ from pydantic_ai.messages import ModelMessagesTypeAdapter from pydantic_ai.models import ModelRequestParameters from pydantic_core import PydanticSerializationError, to_jsonable_python +from airflow.providers.common.ai.durable.base import TOOL_CONFIG_KEY from airflow.providers.common.ai.utils.masking import mask_secrets from airflow.providers.common.ai.utils.prompt_cache import PROMPT_CACHE_SETTING_NAMES from airflow.providers.common.ai.utils.task_logger import get_task_logger @@ -217,11 +219,13 @@ def _strip_volatile(messages_dump: list[dict[str, Any]]) -> list[dict[str, Any]] def _normalize_params(params_dump: dict[str, Any]) -> dict[str, Any]: """ - Drop capability ids and sort set-valued fields from dumped request parameters. + Drop capability ids and the replay opt-out, and sort set-valued fields, in dumped request parameters. A capability without an explicit ``id`` gets a random one per run (``<toolset:d0d75e>``), stamped on its tools' ``capability_id``, so hashing it would make every retry miss the - cache. A set dumps in iteration order, which differs between processes (``PYTHONHASHSEED``), + cache. A tool's ``airflow`` metadata decides only whether the tool replays, and the model + never sees it, so turning it on or off between attempts leaves the requests' fingerprints + as they were. A set dumps in iteration order, which differs between processes (``PYTHONHASHSEED``), and a retry runs in a new process, so ``revealed_tool_names`` is sorted. ``deferred_capability_ids`` is a set of explicit ids and is dropped rather than sorted: a deferred capability's id reaches the model through the capability catalog in the @@ -230,7 +234,16 @@ def _normalize_params(params_dump: dict[str, Any]) -> dict[str, Any]: cleaned = {k: v for k, v in params_dump.items() if k != "deferred_capability_ids"} cleaned["revealed_tool_names"] = sorted(cleaned["revealed_tool_names"]) for key in ("function_tools", "output_tools"): - cleaned[key] = [{k: v for k, v in tool.items() if k != "capability_id"} for tool in cleaned[key]] + cleaned[key] = [_normalize_tool(tool) for tool in cleaned[key]] + return cleaned + + +def _normalize_tool(tool: dict[str, Any]) -> dict[str, Any]: + cleaned = {k: v for k, v in tool.items() if k != "capability_id"} + metadata = cleaned.get("metadata") + if isinstance(metadata, dict): + # Empty metadata hashes like none, so ``{}`` and ``{"airflow": False}`` match too. + cleaned["metadata"] = {k: v for k, v in metadata.items() if k != TOOL_CONFIG_KEY} or None return cleaned @@ -538,7 +551,7 @@ def fingerprint_model_request( The ``ModelRequestParameters`` object is hashed (tool definitions, output mode and schema, native tools, ...) so any change to what is sent to the model invalidates the cached response; only each tool's ``capability_id`` and - ``deferred_capability_ids`` are left out. + ``airflow`` metadata key, and ``deferred_capability_ids``, are left out. Returns ``None`` when the request cannot be fingerprinted. The journal then replays the step without verification, and only right after a step that diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py b/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py index 8fc42d77f14..f77e25cd154 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py +++ b/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py @@ -167,6 +167,8 @@ class JournalStep: except Exception as error: self.fail(error) raise + if not self.replayable: + return payload if isinstance(payload, Unrecordable): self.cannot_record(payload.error) else: @@ -403,8 +405,9 @@ class DurableRun: fingerprinted. A step without one replays on a matching name, and only right after a step that replayed. :param replayable: ``False`` for a step that must always run live, such as a tool - whose effect Airflow cannot observe. It still takes a position, so the steps - after it keep theirs, but nothing is looked up or recorded for it. + whose effect Airflow cannot observe or one its metadata opts out. It still takes + a position, so the steps after it keep theirs, but nothing is looked up or + recorded for it. """ return self._claim(name, kind=kind, fingerprint=fingerprint, replayable=replayable, count_replay=True) 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 87ddf9e6367..2ce8d0a84ec 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 @@ -18,7 +18,11 @@ from __future__ import annotations -from datetime import timedelta +from datetime import UTC, datetime, timedelta + +from pydantic_ai import Tool +from pydantic_ai.capabilities import SetToolMetadata +from pydantic_ai.toolsets import FunctionToolset from airflow.providers.common.ai.durable import AirflowDurability from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook @@ -120,3 +124,35 @@ if SQLToolset is not None: # [END howto_capability_airflow_durability] example_airflow_durability_capability() + + +# --------------------------------------------------------------------------- +# 4. Tools that run on every attempt instead of replaying +# --------------------------------------------------------------------------- + + +# [START howto_agent_durable_tool_opt_out] +def today() -> str: + """Return today's date in UTC.""" + return datetime.now(UTC).date().isoformat() + + +if SQLToolset is not None: + + @dag(default_args={"retries": 3, "retry_delay": timedelta(seconds=30)}, tags=["example"]) + def example_agent_durable_tool_opt_out(): + AgentOperator( + task_id="durable_freshness_check", + prompt="Which tables in load_log have not been loaded today?", + llm_conn_id="pydanticai_default", + durable=True, + toolsets=[ + FunctionToolset(tools=[Tool(today, metadata={"airflow": False})], id="calendar"), + SQLToolset(db_conn_id="postgres_default", allowed_tables=["load_log"]), + ], + capabilities=[SetToolMetadata(tools=["query"], airflow=False)], + ) + + # [END howto_agent_durable_tool_opt_out] + + example_agent_durable_tool_opt_out() 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 71b2b626564..1c85bf7a83f 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 @@ -380,7 +380,9 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa 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. + capabilities with ``@durable_operation`` methods. 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 not: a replayed tool result describes a workspace state the replay did not reproduce, and the first call that misses the cache runs against 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 b4ca0481e88..9f8db919b92 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 @@ -33,8 +33,8 @@ from unittest import mock import pytest from pydantic import BaseModel from pydantic_ai import Agent, CancellationToken, RunContext -from pydantic_ai.capabilities import AbstractCapability, durable_operation -from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior +from pydantic_ai.capabilities import AbstractCapability, SetToolMetadata, durable_operation +from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior, UserError from pydantic_ai.messages import ( ModelMessage, ModelResponse, @@ -44,7 +44,7 @@ from pydantic_ai.messages import ( ToolReturnPart, ) from pydantic_ai.models.function import AgentInfo, FunctionModel -from pydantic_ai.toolsets import FunctionToolset +from pydantic_ai.toolsets import DynamicToolset, FunctionToolset from pydantic_ai.usage import RequestUsage, RunUsage, UsageLimits from airflow.providers.common.ai.durable import AirflowDurability @@ -477,6 +477,15 @@ def _function_toolset(calls: Calls, result: Any) -> FunctionToolset[None]: return toolset +def _dynamic_toolset(calls: Calls, result: Any) -> DynamicToolset[None]: + """A dynamic toolset that builds :func:`_function_toolset` per run step.""" + + async def build(ctx: RunContext[None]) -> FunctionToolset[None]: + return _function_toolset(calls, result) + + return DynamicToolset(build, id="dynamic_db") + + @dataclasses.dataclass class _Event: """A tool result of the user's own whose ``kind`` matches pydantic-ai's recorded tool return.""" @@ -887,11 +896,182 @@ class TestAirflowToolsets: assert calls["model"] == 3 +def _ops_toolset(calls: Calls, *, clock_metadata: dict[str, Any] | None = None, clock_changes: bool = False): + toolset = FunctionToolset(id="ops") + + @toolset.tool_plain(metadata=clock_metadata) + def clock() -> str: + """Read the time.""" + calls.bump("clock") + return f"tick {calls['clock']}" if clock_changes else "tick" + + @toolset.tool_plain + def lookup(key: str) -> str: + """Look a key up.""" + calls.bump("lookup") + return "value" + + return toolset + + +def clock_then_lookup(calls: Calls, *, fail_final: list[bool]): + """A model that calls ``clock``, then ``lookup``, then answers.""" + + def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + calls.bump("model") + match responses_so_far(messages): + case 0: + return ModelResponse(parts=[ToolCallPart("clock", {})]) + case 1: + return ModelResponse(parts=[ToolCallPart("lookup", {"key": "a"})]) + if fail_final[0]: + raise RuntimeError("worker died") + return ModelResponse(parts=[TextPart("done")]) + + return model_fn + + +class TestToolOptOut: + """A tool whose metadata says ``{"airflow": False}`` runs live on every attempt.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("clock_metadata", "extra_capabilities"), + [ + pytest.param({"airflow": False}, [], id="tool-metadata"), + pytest.param(None, [SetToolMetadata(tools=["clock"], airflow=False)], id="SetToolMetadata"), + ], + ) + @pytest.mark.parametrize( + ("clock_changes", "expected_model_calls", "expected_lookups"), + [ + pytest.param(False, 4, 1, id="same-result"), + # The request after the clock sees a different result, so it and the steps after it run again. + pytest.param(True, 5, 2, id="new-result"), + ], + ) + async def test_an_opted_out_function_tool_runs_again_and_the_other_steps_replay( + self, + memory_storage, + clock_metadata, + extra_capabilities, + clock_changes, + expected_model_calls, + expected_lookups, + ): + calls = Calls() + fail_final = [True] + + def build() -> Agent[None, str]: + return Agent( + FunctionModel(clock_then_lookup(calls, fail_final=fail_final)), + name="ops", + toolsets=[_ops_toolset(calls, clock_metadata=clock_metadata, clock_changes=clock_changes)], + capabilities=[AirflowDurability(), *extra_capabilities], + ) + + with pytest.raises(RuntimeError, match="worker died"): + await attempt(memory_storage, build()) + fail_final[0] = False + result = await attempt(memory_storage, build()) + + assert result.output == "done" + assert calls["clock"] == 2 + assert (calls["model"], calls["lookup"]) == (expected_model_calls, expected_lookups) + + @pytest.mark.asyncio + @pytest.mark.parametrize("opt_out_on_attempt", [1, 2], ids=["opted-out-first", "opted-out-on-retry"]) + @pytest.mark.parametrize("dynamic", [False, True], ids=["function", "dynamic"]) + async def test_turning_the_opt_out_on_or_off_between_attempts_keeps_the_rest_replaying( + self, memory_storage, opt_out_on_attempt, dynamic + ): + calls = Calls() + fail_final = [True] + + def build(attempt_number: int) -> Agent[None, str]: + opt_out = [SetToolMetadata(tools=["clock"], airflow=False)] + + async def ops(ctx: RunContext[None]) -> FunctionToolset[None]: + return _ops_toolset(calls) + + return Agent( + FunctionModel(clock_then_lookup(calls, fail_final=fail_final)), + name="ops", + toolsets=[DynamicToolset(ops, id="dynamic_ops") if dynamic else _ops_toolset(calls)], + capabilities=[ + AirflowDurability(), + *(opt_out if attempt_number == opt_out_on_attempt else []), + ], + ) + + with pytest.raises(RuntimeError, match="worker died"): + await attempt(memory_storage, build(1)) + fail_final[0] = False + await attempt(memory_storage, build(2)) + + # The clock is live on the attempt that opts it out, and has nothing to replay after one + # that did; the model requests and the lookup replay either way. + assert (calls["clock"], calls["lookup"], calls["model"]) == (2, 1, 4) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("airflow", "expected_queries"), [(False, 2), (True, 1)], ids=["opted-out", "journaled"] + ) + @pytest.mark.parametrize( + "build_toolset", + [_function_toolset, _WarehouseToolset, _dynamic_toolset], + ids=["function", "airflow", "dynamic"], + ) + async def test_set_tool_metadata_decides_whether_a_tool_replays( + self, memory_storage, build_toolset, airflow, expected_queries + ): + calls = Calls() + fail_final = [True] + + def build() -> Agent[None, str]: + return Agent( + FunctionModel(tool_then_answer(calls, "query", fail_final=fail_final)), + name="analyst", + toolsets=[build_toolset(calls, "3 rows")], + capabilities=[AirflowDurability(), SetToolMetadata(tools=["query"], airflow=airflow)], + ) + + with pytest.raises(RuntimeError, match="worker died"): + await attempt(memory_storage, build()) + fail_final[0] = False + result = await attempt(memory_storage, build()) + + assert result.output == "answer: 3 rows" + assert calls["query"] == expected_queries + # The model request that chose the query replayed either way. + assert calls["model"] == 3 + + @pytest.mark.asyncio + @pytest.mark.parametrize("value", ["false", 0]) + @pytest.mark.parametrize("build_toolset", [_function_toolset, _WarehouseToolset]) + async def test_an_airflow_metadata_value_that_is_not_a_bool_is_refused( + self, memory_storage, value, build_toolset + ): + calls = Calls() + agent = Agent( + FunctionModel(tool_then_answer(calls, "query", fail_final=[False])), + name="analyst", + toolsets=[build_toolset(calls, "3 rows")], + capabilities=[AirflowDurability(), SetToolMetadata(tools=["query"], airflow=value)], + ) + + with pytest.raises(UserError, match="Tool 'query' has invalid 'airflow' metadata"): + await attempt(memory_storage, agent) + assert calls["query"] == 0 + + class Answer(BaseModel): rows: int -def build_analyst(calls: Calls, queries: list[str]) -> Agent[None, str]: +def build_analyst( + calls: Calls, queries: list[str], metadata: dict[str, Any] | None = None +) -> Agent[None, str]: """An analyst whose first answer queries a column that does not exist; later answers query one that does.""" def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: @@ -903,7 +1083,7 @@ def build_analyst(calls: Calls, queries: list[str]) -> Agent[None, str]: toolset = FunctionToolset(id="db") - @toolset.tool_plain + @toolset.tool_plain(metadata=metadata) def query(sql: str) -> str: queries.append(sql) if "nope" in sql: @@ -1014,8 +1194,10 @@ class TestModelAtFault: assert (calls["model"], calls["query"]) == (2, 2) @pytest.mark.asyncio + # An opted-out call still takes its step, so its failures count the same way. + @pytest.mark.parametrize("metadata", [None, {"airflow": False}], ids=["journaled", "opted-out"]) async def test_a_tool_that_keeps_failing_on_the_models_arguments_asks_the_model_again( - self, memory_storage + self, memory_storage, metadata ): """The second attempt retries the same call in case the failure was transient; the third asks the model.""" calls = Calls() @@ -1023,8 +1205,8 @@ class TestModelAtFault: for _ in range(2): with pytest.raises(ValueError, match="column nope"): - await attempt(memory_storage, build_analyst(calls, queries)) - result = await attempt(memory_storage, build_analyst(calls, queries)) + await attempt(memory_storage, build_analyst(calls, queries, metadata)) + result = await attempt(memory_storage, build_analyst(calls, queries, metadata)) assert result.output == "answer: 3 rows" assert queries == ["select nope", "select nope", "select 1"] diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_fingerprint.py b/providers/common/ai/tests/unit/common/ai/durable/test_fingerprint.py index 1a281193941..c3143fb3b4b 100644 --- a/providers/common/ai/tests/unit/common/ai/durable/test_fingerprint.py +++ b/providers/common/ai/tests/unit/common/ai/durable/test_fingerprint.py @@ -200,10 +200,11 @@ class TestModelRequestFingerprint: {"parameters_json_schema": {"type": "object", "properties": {"q": {"type": "string"}}}}, id="schema", ), + pytest.param({"metadata": {"scopes": ["read"]}}, id="metadata"), ], ) def test_changes_with_tool_content_next_to_the_capability_id(self, tools_param, change): - """Only ``capability_id`` is dropped from a tool definition; what the model sees still counts.""" + """Only ``capability_id`` and the ``airflow`` metadata key are dropped from a tool definition.""" base = {"name": "t", "parameters_json_schema": {"type": "object"}, "capability_id": "lookup"} fp1 = fingerprint_model_request( "m", make_messages(), None, ModelRequestParameters(**{tools_param: [ToolDefinition(**base)]}) @@ -217,6 +218,21 @@ class TestModelRequestFingerprint: assert fp1 != fp2 + @pytest.mark.parametrize("tools_param", ["function_tools", "output_tools"]) + @pytest.mark.parametrize( + "other_metadata", [None, {}, {"scopes": ["read"]}], ids=["no-metadata", "empty", "with-other-keys"] + ) + def test_stable_across_the_replay_opt_out(self, tools_param, other_metadata): + """Whether a tool replays is not part of the request, so turning it on or off between attempts keeps the match.""" + + def fingerprint(metadata): + tool = ToolDefinition(name="t", parameters_json_schema={"type": "object"}, metadata=metadata) + return fingerprint_model_request( + "m", make_messages(), None, ModelRequestParameters(**{tools_param: [tool]}) + ) + + assert fingerprint({**(other_metadata or {}), "airflow": False}) == fingerprint(other_metadata) + def test_changes_with_revealed_tool_names(self): fp1 = fingerprint_model_request( "m", make_messages(), None, ModelRequestParameters(revealed_tool_names={"search"}) diff --git a/uv.lock b/uv.lock index 111c1a0fcc2..c6d4f6b2c83 100644 --- a/uv.lock +++ b/uv.lock @@ -4977,12 +4977,12 @@ requires-dist = [ { name = "pydantic-ai-harness", extras = ["codemode"], marker = "extra == 'code-mode'", specifier = ">=0.35.0" }, { name = "pydantic-ai-shields", marker = "extra == 'shields'", specifier = ">=0.3.4" }, { name = "pydantic-ai-skills", marker = "extra == 'skills'", specifier = ">=1.2.0,<2" }, - { name = "pydantic-ai-slim", specifier = ">=2.36.0" }, - { name = "pydantic-ai-slim", extras = ["anthropic"], marker = "extra == 'anthropic'", specifier = ">=2.36.0" }, - { name = "pydantic-ai-slim", extras = ["bedrock"], marker = "extra == 'bedrock'", specifier = ">=2.36.0" }, - { name = "pydantic-ai-slim", extras = ["google"], marker = "extra == 'google'", specifier = ">=2.36.0" }, - { name = "pydantic-ai-slim", extras = ["mcp"], marker = "extra == 'mcp'", specifier = ">=2.36.0" }, - { name = "pydantic-ai-slim", extras = ["openai"], marker = "extra == 'openai'", specifier = ">=2.36.0" }, + { name = "pydantic-ai-slim", specifier = ">=2.46.0" }, + { name = "pydantic-ai-slim", extras = ["anthropic"], marker = "extra == 'anthropic'", specifier = ">=2.46.0" }, + { name = "pydantic-ai-slim", extras = ["bedrock"], marker = "extra == 'bedrock'", specifier = ">=2.46.0" }, + { name = "pydantic-ai-slim", extras = ["google"], marker = "extra == 'google'", specifier = ">=2.46.0" }, + { name = "pydantic-ai-slim", extras = ["mcp"], marker = "extra == 'mcp'", specifier = ">=2.46.0" }, + { name = "pydantic-ai-slim", extras = ["openai"], marker = "extra == 'openai'", specifier = ">=2.46.0" }, { name = "pypdf", marker = "extra == 'pdf'", specifier = ">=4.0.0" }, { name = "python-docx", marker = "extra == 'docx'", specifier = ">=1.0.0" }, { name = "sqlglot", marker = "extra == 'sql'", specifier = ">=30.0.0" }, @@ -5014,7 +5014,7 @@ dev = [ { name = "protobuf", specifier = ">=6.31.1" }, { name = "pydantic-ai-harness", extras = ["codemode"], specifier = ">=0.35.0" }, { name = "pydantic-ai-skills", specifier = ">=1.2.0,<2" }, - { name = "pydantic-ai-slim", extras = ["mcp"], specifier = ">=2.36.0" }, + { name = "pydantic-ai-slim", extras = ["mcp"], specifier = ">=2.46.0" }, { name = "sqlglot", specifier = ">=30.0.0" }, ] docs = [{ name = "apache-airflow-devel-common", extras = ["docs"], editable = "devel-common" }]
