This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch openai-agents-replayed-errors-retry-budget in repository https://gitbox.apache.org/repos/asf/airflow.git
commit d8b8738eee028ddda7eb859a90a0df97045c8371 Author: Kaxil Naik <[email protected]> AuthorDate: Thu Nov 5 23:30:14 2026 +0000 Count replayed tool errors against the retry limit in OpenAI Agents durable runs Under run_durably, an AirflowTools tool's correctable error replays from the journal without the tool running, and each tool's retry budget starts every run at zero. A task retry therefore gave the model max_retries new corrections on top of the ones it had already used. Pydantic AI does not have the gap: a replayed ModelRetry is raised again, so its own count includes it. Each AirflowTools call now records the count it left in the tool's retry budget, with when it finished, and a replay restores the latest-finished count, so the limit spans the task's tries. A run that uses the limit up fails on the model's account, as a Pydantic AI run does, so a retry asks the model again from its last request instead of replaying its answers into the same failure. --- .../common/ai/docs/frameworks/openai_agents.rst | 18 +- .../providers/common/ai/durable/openai_agents.py | 26 +- .../providers/common/ai/tools/_from_toolset.py | 138 ++++++++-- .../providers/common/ai/tools/openai_agents.py | 61 ++++- .../unit/common/ai/durable/test_openai_agents.py | 294 ++++++++++++++++++++- .../unit/common/ai/tools/test__from_toolset.py | 46 +++- 6 files changed, 544 insertions(+), 39 deletions(-) diff --git a/providers/common/ai/docs/frameworks/openai_agents.rst b/providers/common/ai/docs/frameworks/openai_agents.rst index 8b55cdd03c0..be6c0e572b7 100644 --- a/providers/common/ai/docs/frameworks/openai_agents.rst +++ b/providers/common/ai/docs/frameworks/openai_agents.rst @@ -360,7 +360,8 @@ that has an ``arr`` column but no ``revenue`` column: A toolset's tool may fail this way once in a row by default; pass ``max_retries`` to the toolset to allow more. The count starts over after a call that succeeds and on every run -of the agent. A tool still failing after that raises +of the agent, except that a retry of a durable run counts the failures it replays as well +(see :ref:`howto/frameworks:openai-agents-durable`). A tool still failing after that raises :class:`~airflow.providers.common.ai.tools.ToolCallError` out of ``Runner.run`` and ``Runner.run_sync``, so the task fails and Airflow's own retry takes over. Catch it as ``ToolCallError``; the traceback names a private subclass of it. For the same query sent @@ -581,12 +582,25 @@ starting agent and of the agents in its ``handoffs`` list, and an agent used thr ``as_tool``, which replays as one call. Recorded tool results pass through Airflow's secret masker before they are saved, and a replayed model call reports no token usage. +A failure the model could correct, which an ``AirflowTools`` tool answers with +``{"error": ...}``, is recorded and replays like any other result, together with how many of +the tool's corrections were used when it finished. The replay restores that count, so the +retry limit spans the task's tries: a model that used one of two corrections before the +crash has one left on the retry, not two. A run that used the limit up fails on the model's +answers, so each retry asks the model once more, with the errors it already got in its +history. The task log notes when a replay changes the count: + +.. code-block:: text + + Tool query: an earlier attempt left 1 of its 2 correction(s) used + From the first call that differs, that call and every call after it run live: - A retry whose prompt, instructions, model settings or tools differ from the failed try's runs everything live from the first call that sees the difference. - When the run failed because of the model's answers (``MaxTurnsExceeded``, - ``ModelBehaviorError``, ``ModelRefusalError``, or an output or tool input guardrail), + ``ModelBehaviorError``, ``ModelRefusalError``, an output or tool input guardrail, or an + ``AirflowTools`` tool that used up its retry limit), replaying those answers would fail it the same way, so the retry asks the model again from its last request. diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py b/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py index 9f2388c3dd5..155d3e6217e 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py +++ b/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py @@ -64,6 +64,7 @@ except ImportError as e: from airflow.providers.common.ai.durable.fingerprint import _digest, fingerprint_tool_call from airflow.providers.common.ai.durable.journal import _EXECUTING_STEP, current_journal +from airflow.providers.common.ai.tools.openai_agents import _REPLAY_RETRY, _RETRY_COUNT, _RetryLimitReached from airflow.providers.common.ai.utils.coroutines import run_coroutine_sync from airflow.providers.common.ai.utils.masking import mask_secrets from airflow.providers.common.ai.utils.task_logger import get_task_logger @@ -92,6 +93,7 @@ _MODEL_AT_FAULT = ( MaxTurnsExceeded, OutputGuardrailTripwireTriggered, ToolInputGuardrailTripwireTriggered, + _RetryLimitReached, ) # The tool results the SDK hands the model as structured content rather than as text. @@ -150,7 +152,10 @@ async def run_durably( function tools of an agent reached through ``handoff()``, and its model too unless the agent's ``model`` is a model name or not set, and streamed runs. A tool call that fails, including one whose exception ``failure_error_function`` turns into text for the model, - is not recorded either, so a retry runs it again. A replayed model call reports no + is not recorded either, so a retry runs it again. A failure the model could correct, + which an ``AirflowTools`` tool answers with ``{"error": ...}``, is recorded, and still + counts against the tool's retry limit when it replays; a run that used that limit up asks + the model again from its last request on a retry. A replayed model call reports no usage, since it costs nothing. A ``session`` makes every retry run live: the SDK adds the failed try's prompt to it @@ -239,6 +244,9 @@ def _get_durable_tool(tool: Tool, durable_run: DurableRun) -> Tool: set_function_tool_failure_error_function(durable_tool, remember) invoke = bind(durable_tool) + # An AirflowTools tool, whose correctable failures count against its retry limit. + replay_retry = getattr(tool, _REPLAY_RETRY, None) + async def on_invoke_tool(ctx: ToolContext[Any], arguments: str) -> Any: if _is_inside_a_call(durable_run): return await invoke(ctx, arguments) @@ -246,12 +254,21 @@ def _get_durable_tool(tool: Tool, durable_run: DurableRun) -> Tool: fingerprint = _build_tool_fingerprint(tool.name, arguments, ctx.tool_call_id) step = durable_run.claim_reserved(ctx.tool_call_id, name, kind="tool", fingerprint=fingerprint) if step.replayed: + # Such a tool's result is text, so a dict is one recorded with the count it left in the + # tool's retry budget. Restored, so the retry limit spans attempts. + if replay_retry is not None and isinstance(step.payload, dict): + await replay_retry(ctx, step.payload["retry"]) + return step.payload["output"] return step.payload try: with step.executing(): output = await invoke(ctx, arguments) except Exception as error: - step.fail(error) + # The model's answers used this tool's retry limit up, not this call, so the run fails + # on the model's account and a retry asks it again from its last request. The limit of + # a tool inside an agent this one calls through as_tool is that call's failure. + if not (replay_retry is not None and isinstance(error, _RetryLimitReached)): + step.fail(error) raise # A failure the SDK turned into text for the model is not recorded, so a retry runs the # call again, as it does a call that raised. @@ -259,7 +276,10 @@ def _get_durable_tool(tool: Tool, durable_run: DurableRun) -> Tool: step.fail(handled_error) elif not getattr(ctx, _DEFAULT_FAILURE_HANDLED, False): # Secrets must not reach the journal any more than the model. - step.record(mask_secrets(_as_sent(output))) + payload = mask_secrets(_as_sent(output)) + if replay_retry is not None and (retry := getattr(ctx, _RETRY_COUNT, None)) is not None: + payload = {"output": payload, "retry": retry} + step.record(payload) return output durable_tool.on_invoke_tool = on_invoke_tool diff --git a/providers/common/ai/src/airflow/providers/common/ai/tools/_from_toolset.py b/providers/common/ai/src/airflow/providers/common/ai/tools/_from_toolset.py index 21bd028f338..1c1eaf1f42c 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/tools/_from_toolset.py +++ b/providers/common/ai/src/airflow/providers/common/ai/tools/_from_toolset.py @@ -21,13 +21,15 @@ from __future__ import annotations import asyncio import concurrent.futures import contextlib +import logging +import math import threading import time from collections.abc import Hashable, Iterator from contextlib import contextmanager from contextvars import ContextVar -from dataclasses import replace -from typing import TYPE_CHECKING, Any +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any, TypedDict from pydantic import ValidationError from pydantic_ai import RunContext @@ -39,27 +41,67 @@ from airflow.providers.common.ai.tools import AirflowTool, ToolCallError, ToolRe from airflow.providers.common.ai.utils.coroutines import run_coroutine_sync if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool +log = logging.getLogger(__name__) + + +class RetryCount(TypedDict): + """The corrections a tool had used when a call finished, and when it finished.""" + + failures: int + # Wall-clock time, not monotonic: a retry compares it with times recorded by earlier tries, + # which may have run on other workers. + finished: float + + +@dataclass(slots=True) +class ToolCallScope: + """Which agent run and model turn a tool call belongs to, and what it did to the tool's retry budget.""" -# Set by a framework adapter around each call: which agent run and which model turn the call -# belongs to, where the framework says. The retry budget counts by them. -_call_scope: ContextVar[tuple[Hashable | None, Hashable | None]] = ContextVar( - "common_ai_tool_call_scope", default=(None, None) -) + run: Hashable | None = None + turn: Hashable | None = None + # The tool function the adapter calls. Only its own count is recorded and replayed, not that + # of a toolset tool a custom function calls in turn. + function: Callable[[dict[str, Any]], Awaitable[ToolResult]] | None = None + # When set, the call takes this count, recorded by an earlier attempt, instead of running. + replayed: RetryCount | None = None + # Set by a call that failed or succeeded; None when it left the count alone, such as a ToolFailed. + counted: RetryCount | None = None + # Set when the call failed once more than the budget allows, which is the model's doing. + limit_reached: bool = False + + +# Set by a framework adapter around each call. The retry budget counts by its run and turn. +_call_scope: ContextVar[ToolCallScope | None] = ContextVar("common_ai_tool_call_scope", default=None) @contextmanager -def tool_call_scope(*, run: Hashable | None, turn: Hashable | None = None) -> Iterator[None]: +def tool_call_scope( + *, + run: Hashable | None, + turn: Hashable | None = None, + function: Callable[[dict[str, Any]], Awaitable[ToolResult]] | None = None, + replayed: RetryCount | None = None, +) -> Iterator[ToolCallScope]: """ Tell the tools which agent run and which model turn the calls inside the block belong to. Each tool's retry budget starts over when the run changes, and counts at most one failure per turn, as pydantic-ai does. Without a turn, calls that run at the same time count as one. + + Given the ``function`` of the tool the block calls, the block yields the count that call left + in the tool's retry budget, so durable execution can record it with the call's result. On a + retry that replays that result, ``replayed`` makes the call restore the recorded count + instead of running, so the tool gets ``max_retries`` corrections across attempts, not per + attempt. """ - token = _call_scope.set((run, turn)) + scope = ToolCallScope(run=run, turn=turn, function=function, replayed=replayed) + token = _call_scope.set(scope) try: - yield + yield scope finally: _call_scope.reset(token) @@ -165,6 +207,10 @@ class _RetryBudget: self._run: Hashable | None = None self._failed_turn: Hashable | None = None self._counted_until = float("-inf") + # When the call whose count this is finished, live or in an earlier attempt, and whether a + # call has counted live in this run. + self._finished = float("-inf") + self._live = False @property def failures(self) -> int: @@ -177,6 +223,8 @@ class _RetryBudget: self._run = run self._failures = 0 self._failed_turn = None + self._finished = float("-inf") + self._live = False def succeeded(self, turn: Hashable | None) -> None: # A success does not clear a failure of the same turn. @@ -196,6 +244,31 @@ class _RetryBudget: self._failures = 0 return True + def take(self) -> RetryCount: + """Return the count a call that just failed or succeeded left, for durable execution to record.""" + # Later than every count restored before it even when this worker's clock is behind the + # one an earlier attempt ran on, so a later retry orders the attempts' counts as they ran. + # A live call that counts before a replayed one of its turn restores is ordered by the + # clocks alone, which holds while workers' clocks differ by less than the time between tries. + self._finished = max(time.time(), math.nextafter(self._finished, math.inf)) + self._live = True + return {"failures": self._failures, "finished": self._finished} + + def restore(self, count: RetryCount) -> None: + """ + Take a count an earlier attempt recorded, unless a call that finished later left the current one. + + A replay reaches calls in the order they started, but the attempt counted them in the + order they finished, which differs for calls that ran alongside each other. A call that + runs live finished after any an earlier attempt recorded. + """ + if count["failures"]: + # A call running alongside this failure fails in the same turn, as it did live. + self._counted_until = time.monotonic() + if not self._live and count["finished"] > self._finished: + self._finished = count["finished"] + self._failures = count["failures"] + def _as_airflow_tool( toolset: AbstractToolset[Any], @@ -206,16 +279,36 @@ def _as_airflow_tool( ) -> AirflowTool: budget = _RetryBudget(tool.max_retries) - def correctable(started: float, turn: Hashable | None, message: str) -> ToolResult: - if budget.failed(started, turn): + def correctable(started: float, scope: ToolCallScope, message: str) -> ToolResult: + if budget.failed(started, scope.turn): + scope.limit_reached = True raise ToolCallError(f"{name} kept failing after {budget.max_retries} correction(s): {message}") + scope.counted = budget.take() return ToolResult(content=message, is_error=True) - async def call(arguments: dict[str, Any], started: float, turn: Hashable | None) -> ToolResult: + def replay(count: RetryCount) -> ToolResult: + # The result the earlier attempt recorded reaches the model from the journal, not from here. + restored = budget.failures + budget.restore(count) + if budget.failures != restored: + log.info( + "Tool %s: an earlier attempt left %d of its %d correction(s) used", + name, + budget.failures, + budget.max_retries, + ) + return ToolResult(content=None) + + async def call(arguments: dict[str, Any], started: float, scope: ToolCallScope) -> ToolResult: + if scope.replayed is not None: + # A sequential call holds the line without awaiting, so without this the calls the + # framework started alongside it would not have started when it marks its turn. + await asyncio.sleep(0) + return replay(scope.replayed) try: validated = tool.args_validator.validate_python(arguments) except ValidationError as e: - return correctable(started, turn, str(e)) + return correctable(started, scope, str(e)) # A ValidationError raised by the tool itself is not the model's to fix: the call # may already have had a side effect, so it propagates rather than inviting a retry. # Per call, as pydantic-ai's ToolManager does, so ``ctx.retry`` and ``ctx.last_attempt`` @@ -224,26 +317,31 @@ def _as_airflow_tool( try: result = await toolset.call_tool(name, validated, call_ctx, tool) except ModelRetry as e: - return correctable(started, turn, e.message) + return correctable(started, scope, e.message) except ToolFailed as e: return ToolResult(content=e.message, is_error=True) - budget.succeeded(turn) + budget.succeeded(scope.turn) + scope.counted = budget.take() return ToolResult(content=result) async def function(arguments: dict[str, Any]) -> ToolResult: started = time.monotonic() - run, turn = _call_scope.get() - budget.start(run) + scope = _call_scope.get() or ToolCallScope() + if scope.function is not function: + # Called by another tool, or by an adapter that records nothing: count by the scope's run + # and turn, but keep this call's count out of the scope the adapter reads. + scope = ToolCallScope(run=scope.run, turn=scope.turn) + budget.start(scope.run) if not tool.tool_def.sequential: # Validation does not await, so without this a call the framework started alongside # this one would not have started yet when this one fails, and the budget would # count both failures of one turn. await asyncio.sleep(0) - return await call(arguments, started, turn) + return await call(arguments, started, scope) ticket = in_order.take() await in_order.wait(ticket) try: - return await call(arguments, started, turn) + return await call(arguments, started, scope) finally: in_order.done(ticket) diff --git a/providers/common/ai/src/airflow/providers/common/ai/tools/openai_agents.py b/providers/common/ai/src/airflow/providers/common/ai/tools/openai_agents.py index 8d855171f82..73133c9df67 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/tools/openai_agents.py +++ b/providers/common/ai/src/airflow/providers/common/ai/tools/openai_agents.py @@ -26,6 +26,7 @@ from __future__ import annotations import copy import json +from contextlib import contextmanager from typing import TYPE_CHECKING, Any from urllib.parse import urlsplit @@ -60,6 +61,7 @@ from airflow.providers.common.ai.utils.tool_definition import serialize_for_llm from airflow.providers.common.ai.utils.tool_metrics import calling_framework if TYPE_CHECKING: + from collections.abc import Iterator from contextvars import Token from agents.tool_context import ToolContext @@ -68,6 +70,7 @@ if TYPE_CHECKING: from opentelemetry.trace import Span as OtelSpan, TracerProvider from airflow.providers.common.ai.tools import ToolProvider + from airflow.providers.common.ai.tools._from_toolset import RetryCount, ToolCallScope __all__ = ["AirflowTools", "AirflowTraceProcessor"] @@ -76,6 +79,13 @@ _MCP_HINT = ( "such as MCPServerStreamableHttp, instead." ) +# Set on a tool call's context: the count the call left in its tool's retry budget, for durable +# execution to record with the result. +_RETRY_COUNT = "_airflow_retry_count" +# Set on each function tool these tools become: given the count a call the durable journal +# replays left in the tool's retry budget, restores it. +_REPLAY_RETRY = "_airflow_replay_retry" + class AirflowTools(list[Tool]): """ @@ -106,8 +116,9 @@ class AirflowTools(list[Tool]): with ``strict_json_schema=False``: the SDK's strict mode would make every argument required. A failure the model can correct, such as a query naming a missing column, reaches it as ``{"error": ...}``, counted against the tool's retry limit afresh on - every run of the agent. Arguments that are not a JSON object also come back as - ``{"error": ...}``, bounded by the run's ``max_turns`` rather than that limit. A tool + every run of the agent; a retry of a durable run counts the ones it replays as well. + Arguments that are not a JSON object also come back as ``{"error": ...}``, bounded by + the run's ``max_turns`` rather than that limit. A tool still failing once the limit is used up, or any other failure, raises :class:`~airflow.providers.common.ai.tools.ToolCallError` out of ``Runner.run``, so the task fails and Airflow retries it. Tools a toolset marks sequential, such as the @@ -154,6 +165,32 @@ class _SameRun: return id(self._usage) +class _RetryLimitReached(_FailedToolCall): + """ + A tool kept failing in ways the model could correct until its retry limit ran out. + + The model's answers used the limit up, not the call, so durable execution asks the model + again on a retry rather than replaying them into the same failure. + """ + + +@contextmanager +def _calling_tool( + ctx: ToolContext[Any], tool: AirflowTool, replayed: RetryCount | None = None +) -> Iterator[ToolCallScope]: + # The SDK does not identify the model turn, so calls that run at the same time count once. + with ( + calling_framework("openai_agents"), + tool_call_scope(run=_SameRun(ctx.usage), function=tool.function, replayed=replayed) as scope, + ): + try: + yield scope + except ToolCallError as e: + # Already masked, and logged to the task log with its traceback by AirflowTool.call. + error = _RetryLimitReached if scope.limit_reached else _FailedToolCall + raise error(str(e)) from None + + def _to_function_tool(tool: AirflowTool) -> FunctionTool: async def call_airflow_tool(ctx: ToolContext[Any], arguments: str) -> str: # The SDK passes the model's arguments as it sent them, without parsing them first. @@ -163,18 +200,20 @@ def _to_function_tool(tool: AirflowTool) -> FunctionTool: return serialize_for_llm({"error": f"The arguments are not valid JSON: {e}"}) if not isinstance(parsed, dict): return serialize_for_llm({"error": "The arguments must be a JSON object."}) - # The SDK does not identify the model turn, so calls that run at the same time count once. - try: - with calling_framework("openai_agents"), tool_call_scope(run=_SameRun(ctx.usage)): - result = await tool.call(parsed) - except ToolCallError as e: - # Already masked, and logged to the task log with its traceback by AirflowTool.call. - raise _FailedToolCall(str(e)) from None + with _calling_tool(ctx, tool) as scope: + result = await tool.call(parsed) + setattr(ctx, _RETRY_COUNT, scope.counted) if result.is_error: return serialize_for_llm({"error": result.content}) return serialize_for_llm(result.content) - return FunctionTool( + async def replay_retry(ctx: ToolContext[Any], count: RetryCount) -> None: + # Not through AirflowTool.call: nothing reaches the model, and its task log lines would + # show a call that never ran. Only a toolset tool records a count, and it does not run. + with _calling_tool(ctx, tool, replayed=count): + await tool.function({}) + + function_tool = FunctionTool( name=tool.name, description=tool.description, # The source schema can be a toolset's module-level constant, so hand over a copy. @@ -182,6 +221,8 @@ def _to_function_tool(tool: AirflowTool) -> FunctionTool: on_invoke_tool=call_airflow_tool, strict_json_schema=False, ) + setattr(function_tool, _REPLAY_RETRY, replay_retry) + return function_tool class AirflowTraceProcessor(TracingProcessor): diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py b/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py index da4adcaf7fb..c712b373687 100644 --- a/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py +++ b/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py @@ -19,8 +19,10 @@ from __future__ import annotations import asyncio import collections import contextlib +import dataclasses import datetime import json +from collections.abc import Container from typing import Any from unittest import mock @@ -51,6 +53,8 @@ from agents.retry import ModelRetryBackoffSettings, ModelRetrySettings from agents.tool_context import ToolContext from openai.types.responses import ResponseFunctionToolCall, ResponseOutputMessage, ResponseOutputText from pydantic import BaseModel +from pydantic_ai.exceptions import ModelRetry +from pydantic_ai.toolsets import FunctionToolset from airflow.providers.common.ai.durable.base import build_step_key from airflow.providers.common.ai.durable.journal import DurableJournal, journal_scope @@ -60,7 +64,8 @@ from airflow.providers.common.ai.durable.openai_agents import ( _get_durable_tool, run_durably_sync, ) -from airflow.providers.common.ai.tools import AirflowTool, ToolResult +from airflow.providers.common.ai.tools import AirflowTool, ToolCallError, ToolResult +from airflow.providers.common.ai.tools._from_toolset import airflow_tools_from_toolset from airflow.providers.common.ai.tools.openai_agents import AirflowTools from unit.common.ai.durable.memory_storage import MemoryStorage @@ -154,6 +159,98 @@ class _CountingTool: ) +class _ToolsetLookup: + """ + A ``lookup`` tool from a pydantic-ai toolset, which counts the calls that reached it. + + It raises ``ModelRetry``, a failure the model can correct, for the keys in ``retry``, and + ``RuntimeError``, which fails the try, for those in ``broken``, after sleeping + ``delays[key]`` seconds; a key without a delay answers without awaiting. + """ + + def __init__( + self, + *, + max_retries: int, + retry: Container[str], + broken: Container[str] = (), + delays: dict[str, float] | None = None, + sequential: bool = False, + ) -> None: + self.max_retries = max_retries + self.retry = retry + self.broken = broken + self.delays = delays or {} + self.sequential = sequential + self.calls = 0 + + def tool(self) -> AirflowTool: + async def lookup(key: str) -> str: + self.calls += 1 + if key in self.delays: + await asyncio.sleep(self.delays[key]) + if key in self.broken: + raise RuntimeError(f"the warehouse dropped the connection looking up {key}") + if key in self.retry: + raise ModelRetry(f"There is no {key}.") + return f"{key}=42" + + # A fresh tool per attempt, as a retried task builds its agent again. + toolset = FunctionToolset([lookup], max_retries=self.max_retries, sequential=self.sequential) + (tool,) = airflow_tools_from_toolset(toolset) + return tool + + +class _TurnsModel(Model): + """ + Asks for one ``lookup`` per key of each turn, all of a turn's in one response, then answers. + + ``fail_on_answer`` makes the request that would answer raise, as a crashed worker would leave it. + """ + + def __init__(self, turns: list[list[str]], *, fail_on_answer: bool = False) -> None: + self.turns = turns + self.fail_on_answer = fail_on_answer + self.calls = 0 + + async def get_response(self, system_instructions, input, *args: Any, **kwargs: Any) -> ModelResponse: + self.calls += 1 + answered = sum( + 1 + for item in (input if isinstance(input, list) else []) + if isinstance(item, dict) and item.get("type") == "function_call_output" + ) + done = 0 + output: list[ResponseFunctionToolCall | ResponseOutputMessage] = [] + for keys in self.turns: + if answered == done: + output = [ + ResponseFunctionToolCall( + id=f"fc-{key}", + call_id=f"call-{key}", + name="lookup", + arguments=json.dumps({"key": key}), + type="function_call", + ) + for key in keys + ] + break + done += len(keys) + else: + if self.fail_on_answer: + raise RuntimeError("worker lost") + text = ResponseOutputText(text="answer", type="output_text", annotations=[]) + output = [ + ResponseOutputMessage( + id="msg", role="assistant", status="completed", type="message", content=[text] + ) + ] + return ModelResponse(output=list(output), usage=Usage(), response_id=f"resp-{self.calls}") + + def stream_response(self, *args: Any, **kwargs: Any): + raise NotImplementedError + + class _ParallelModel(Model): """ Asks for one ``lookup`` per key in a single response, then answers with their results. @@ -263,7 +360,13 @@ class _Report(BaseModel): rows: int -def _attempt(storage: MemoryStorage, model: Model, tool: _CountingTool, prompt: str = "go", **kwargs: Any): +def _attempt( + storage: MemoryStorage, + model: Model, + tool: _CountingTool | _ToolsetLookup, + prompt: str = "go", + **kwargs: Any, +): """Run one task attempt: a fresh journal over the storage the attempts share.""" agent = Agent(name="analyst", model=model, tools=AirflowTools(tool.tool())) journal = DurableJournal(storage, clean_up_after_run=True) @@ -610,6 +713,193 @@ class TestRunDurably: # The answer that asked for the lookup replays; only the final answer is asked for. assert retry_model.calls == 1 + def test_a_replayed_correctable_failure_counts_against_the_tools_retry_limit(self): + """Across attempts the tool gets max_retries corrections in all, not max_retries per attempt.""" + storage = MemoryStorage() + tool = _ToolsetLookup(max_retries=2, retry={"k0", "k1", "k2"}) + with pytest.raises(RuntimeError, match="worker lost"): + _attempt(storage, _ScriptedModel(fail_on_answer=True), tool) + tool.calls = 0 + + # The replayed failure of k0 is the first correction, k1 the second, and k2 one too many. + with pytest.raises(ToolCallError, match="lookup kept failing after 2 correction"): + _attempt(storage, _ScriptedModel(lookups=3), tool) + assert tool.calls == 2 + + def test_a_run_that_used_up_a_tools_retry_limit_asks_the_model_again_on_retry(self): + """The model's answers used the limit up, so a retry asks it again rather than replaying them.""" + storage = MemoryStorage() + tool = _ToolsetLookup(max_retries=1, retry={"k0", "k1"}) + with pytest.raises(ToolCallError, match="lookup kept failing after 1 correction"): + _attempt(storage, _ScriptedModel(lookups=2), tool) + tool.calls = 0 + retry_model = _ScriptedModel(lookups=1) + + result, _ = _attempt(storage, retry_model, tool) + + # The answer that asked for k1 is asked for again, and this time the model answers. + assert retry_model.calls == 1 + assert tool.calls == 0 + assert result.final_output.startswith("answer: ") + + def test_a_replayed_success_starts_the_tools_retry_count_over(self): + storage = MemoryStorage() + tool = _ToolsetLookup(max_retries=1, retry={"k0", "k2"}) + with pytest.raises(RuntimeError, match="worker lost"): + _attempt(storage, _ScriptedModel(lookups=2, fail_on_answer=True), tool) + tool.calls = 0 + + # k1 succeeded after k0 failed, so k2 is the one correction the limit allows. + result, _ = _attempt(storage, _ScriptedModel(lookups=3), tool) + + assert tool.calls == 1 + assert result.final_output.startswith("answer: ") + + def test_parallel_calls_replay_the_count_they_left_in_the_order_they_finished(self): + """ + Live, b fails and then a, still running, succeeds, so d is the first correction again. + + A replay reaches a and b in the order they started, which would count d as the second. + """ + storage = MemoryStorage() + tool = _ToolsetLookup(max_retries=1, retry={"b", "d"}, delays={"a": 0.05}) + with pytest.raises(RuntimeError, match="worker lost"): + _attempt(storage, _TurnsModel([["a", "b"]], fail_on_answer=True), tool) + tool.calls = 0 + + result, _ = _attempt(storage, _TurnsModel([["a", "b"], ["d"]]), tool) + + assert result.final_output == "answer" + assert tool.calls == 1 + + def test_a_third_attempt_counts_the_corrections_of_both_earlier_attempts(self): + storage = MemoryStorage() + tool = _ToolsetLookup(max_retries=2, retry={"k0", "k1", "k2"}) + with pytest.raises(RuntimeError, match="worker lost"): + _attempt(storage, _ScriptedModel(lookups=1, fail_on_answer=True), tool) + with pytest.raises(RuntimeError, match="worker lost"): + _attempt(storage, _ScriptedModel(lookups=2, fail_on_answer=True), tool) + tool.calls = 0 + + # k0 from the first attempt and k1 from the second are both corrections already used. + with pytest.raises(ToolCallError, match="lookup kept failing after 2 correction"): + _attempt(storage, _ScriptedModel(lookups=3), tool) + assert tool.calls == 1 + + @pytest.mark.parametrize("sequential", [False, True], ids=["parallel", "sequential"]) + def test_a_live_failure_alongside_a_replayed_one_counts_once(self, sequential): + """b failed the first try outright; on the retry it fails in a way the model can correct, as a did.""" + storage = MemoryStorage() + tool = _ToolsetLookup(max_retries=1, retry={"a", "b"}, broken={"b"}, sequential=sequential) + with pytest.raises(ToolCallError, match="warehouse dropped the connection"): + _attempt(storage, _TurnsModel([["a", "b"]]), tool) + tool.broken, tool.calls = (), 0 + + result, _ = _attempt(storage, _TurnsModel([["a", "b"]]), tool) + + assert result.final_output == "answer" + assert tool.calls == 1 + + def test_a_replayed_success_does_not_hide_a_live_failure_alongside_it(self): + storage = MemoryStorage() + tool = _ToolsetLookup(max_retries=1, retry={"b", "d"}, broken={"b"}) + with pytest.raises(ToolCallError, match="warehouse dropped the connection"): + _attempt(storage, _TurnsModel([["a", "b"]]), tool) + tool.broken, tool.calls = (), 0 + + # b fails live before a's replay restores its success, which must not wipe b's correction: + # b is the first correction and d one too many. + with pytest.raises(ToolCallError, match="lookup kept failing after 1 correction"): + _attempt(storage, _TurnsModel([["a", "b"], ["d"]]), tool) + assert tool.calls == 2 + + def test_each_agent_run_in_a_task_restores_its_own_count(self): + """A second run_durably call in the task replays its count after a first one ran live.""" + storage, tool = MemoryStorage(), _ToolsetLookup(max_retries=1, retry={"k1", "k2"}) + + def attempt(second: _ScriptedModel) -> None: + tools = AirflowTools(tool.tool()) + with journal_scope(DurableJournal(storage, clean_up_after_run=True)): + run_durably_sync( + Agent(name="a", model=_ScriptedModel(), tools=tools), "1", run_config=NO_TRACING + ) + run_durably_sync(Agent(name="b", model=second, tools=tools), "2", run_config=NO_TRACING) + + with pytest.raises(RuntimeError, match="worker lost"): + attempt(_ScriptedModel(lookups=2, fail_on_answer=True)) + tool.calls = 0 + + # The first run succeeded, so it runs live again; the second's k1 is its one correction. + with pytest.raises(ToolCallError, match="lookup kept failing after 1 correction"): + attempt(_ScriptedModel(lookups=3)) + assert tool.calls == 2 + + def test_a_custom_tool_that_calls_a_toolset_tool_replays_without_running(self): + storage = MemoryStorage() + inner = _ToolsetLookup(max_retries=1, retry=set()) + outer_calls = [] + + def custom_tool() -> AirflowTool: + lookup = inner.tool() + + async def function(arguments: dict[str, Any]) -> ToolResult: + outer_calls.append(arguments["key"]) + return await lookup.call(arguments) + + return dataclasses.replace(lookup, function=function) + + def attempt(model: Model) -> Any: + agent = Agent(name="analyst", model=model, tools=AirflowTools(custom_tool())) + with journal_scope(DurableJournal(storage, clean_up_after_run=True)): + return run_durably_sync(agent, "go", run_config=NO_TRACING) + + with pytest.raises(RuntimeError, match="worker lost"): + attempt(_ScriptedModel(fail_on_answer=True)) + + result = attempt(_ScriptedModel()) + + assert outer_calls == ["k0"] + assert result.final_output == "answer: k0=42" + + def test_a_tool_rebuilt_without_the_replay_hook_replays_its_result_as_text(self): + storage, tool = MemoryStorage(), _ToolsetLookup(max_retries=1, retry=set()) + + def attempt(model: Model) -> Any: + tools = [dataclasses.replace(function_tool) for function_tool in AirflowTools(tool.tool())] + with journal_scope(DurableJournal(storage, clean_up_after_run=True)): + return run_durably_sync( + Agent(name="analyst", model=model, tools=tools), "go", run_config=NO_TRACING + ) + + with pytest.raises(RuntimeError, match="worker lost"): + attempt(_ScriptedModel(fail_on_answer=True)) + + assert attempt(_ScriptedModel()).final_output == "answer: k0=42" + + def test_a_used_up_limit_inside_an_agent_used_as_a_tool_fails_that_call_not_the_model(self): + """The outer model's answer was fine: the retry replays it and runs the inner agent again.""" + storage, tool = MemoryStorage(), _ToolsetLookup(max_retries=1, retry={"k0", "k1"}) + + def attempt(outer: _ScriptedModel) -> None: + specialist = Agent(name="specialist", model="inner", tools=AirflowTools(tool.tool())) + lead = Agent( + name="lead", + model=outer, + tools=[specialist.as_tool("lookup", "Ask the specialist.", failure_error_function=None)], + ) + config = RunConfig(tracing_disabled=True, model_provider=_Provider(_ScriptedModel(lookups=2))) + with journal_scope(DurableJournal(storage, clean_up_after_run=True)): + run_durably_sync(lead, "go", run_config=config) + + with pytest.raises(ToolCallError, match="lookup kept failing after 1 correction"): + attempt(_ScriptedModel(argument="input")) + outer = _ScriptedModel(argument="input") + + with pytest.raises(ToolCallError, match="lookup kept failing after 1 correction"): + attempt(outer) + + assert outer.calls == 0 + def test_a_run_with_a_session_runs_live_on_retry(self, tmp_path): """The session already holds the failed try's prompt, so the first request differs.""" storage, tool = MemoryStorage(), _CountingTool() diff --git a/providers/common/ai/tests/unit/common/ai/tools/test__from_toolset.py b/providers/common/ai/tests/unit/common/ai/tools/test__from_toolset.py index 76f5030eaec..359e422b7f7 100644 --- a/providers/common/ai/tests/unit/common/ai/tools/test__from_toolset.py +++ b/providers/common/ai/tests/unit/common/ai/tools/test__from_toolset.py @@ -32,8 +32,12 @@ from pydantic_ai.models.test import TestModel from pydantic_ai.toolsets.function import FunctionToolset from pydantic_ai.usage import RunUsage -from airflow.providers.common.ai.tools import ToolCallError, ToolResult -from airflow.providers.common.ai.tools._from_toolset import airflow_tools_from_toolset, tool_call_scope +from airflow.providers.common.ai.tools import AirflowTool, ToolCallError, ToolResult +from airflow.providers.common.ai.tools._from_toolset import ( + RetryCount, + airflow_tools_from_toolset, + tool_call_scope, +) from airflow.providers.common.ai.toolsets.hook import HookToolset from airflow.providers.common.ai.toolsets.sql import SQLToolset @@ -275,6 +279,44 @@ class TestRetryBudget: with pytest.raises(ToolCallError, match="kept failing"): asyncio.run(call_in("turn-2")) + def test_a_count_an_earlier_attempt_recorded_does_not_replace_a_live_one(self): + """The earlier attempt's worker had a clock an hour ahead; the live failure still counts.""" + step = _scripted(ModelRetry("bad"), ModelRetry("bad"), max_retries=1) + recorded: RetryCount = {"failures": 0, "finished": time.time() + 3600} + + async def call_in(turn: str, replayed: RetryCount | None = None) -> ToolResult: + with tool_call_scope(run="run", turn=turn, function=step.function, replayed=replayed): + return await step.function({}) + + assert asyncio.run(call_in("turn-1")).is_error + asyncio.run(call_in("turn-2", replayed=recorded)) + with pytest.raises(ToolCallError, match="kept failing"): + asyncio.run(call_in("turn-3")) + + def test_a_later_attempts_count_outranks_an_earlier_one_from_a_clock_ahead(self): + """ + The first attempt's worker had a clock an hour ahead and recorded one correction used. + + The second replayed it and then succeeded live; the third must restore that success. + """ + first: RetryCount = {"failures": 1, "finished": time.time() + 3600} + + async def call_in(step: AirflowTool, turn: str, replayed: RetryCount | None = None): + with tool_call_scope(run="run", turn=turn, function=step.function, replayed=replayed) as scope: + result = await step.function({}) + return result, scope.counted + + second_attempt = _scripted("ok", max_retries=1) + asyncio.run(call_in(second_attempt, "turn-1", first)) + _, second = asyncio.run(call_in(second_attempt, "turn-2")) + third_attempt = _scripted(ModelRetry("bad"), max_retries=1) + asyncio.run(call_in(third_attempt, "turn-1", first)) + asyncio.run(call_in(third_attempt, "turn-2", second)) + + result, _ = asyncio.run(call_in(third_attempt, "turn-3")) + + assert result.is_error + def test_each_tool_list_starts_with_a_fresh_budget(self): toolset = FunctionToolset([_failing_step], max_retries=0)
