pankajastro commented on code in PR #74770: URL: https://github.com/apache/airflow/pull/74770#discussion_r4284669161
########## providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py: ########## @@ -0,0 +1,448 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Replay an OpenAI Agents SDK agent's completed model and tool calls when Airflow retries its task. + +.. note:: Experimental; see :ref:`howto/stability`. +""" + +from __future__ import annotations + +import copy +import json +from typing import TYPE_CHECKING, Any + +from pydantic import TypeAdapter, ValidationError +from pydantic_core import to_jsonable_python + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + +try: + from agents import ( + Agent, + FunctionTool, + Model, + ModelProvider, + ModelResponse, + RunConfig, + Runner, + ToolOutputFileContent, + ToolOutputImage, + ToolOutputText, + Usage, + ) + from agents.exceptions import ( + MaxTurnsExceeded, + ModelBehaviorError, + ModelRefusalError, + OutputGuardrailTripwireTriggered, + ToolInputGuardrailTripwireTriggered, + ) + from agents.items import TResponseOutputItem +except ImportError as e: + raise AirflowOptionalProviderFeatureException(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.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 + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from agents import AgentOutputSchemaBase, Handoff, ModelSettings, ModelTracing, RunResult, Tool + from agents.items import TResponseInputItem, TResponseStreamEvent + from agents.retry import ModelRetryAdvice, ModelRetryAdviceRequest + from agents.tool_context import ToolContext + + from airflow.providers.common.ai.durable.journal import DurableRun, JournalStep + +__all__ = ["run_durably", "run_durably_sync"] + +log = get_task_logger() + +_OUTPUT_ITEMS: TypeAdapter[list[TResponseOutputItem]] = TypeAdapter(list[TResponseOutputItem]) + +# Errors that come from the model's answers rather than from a step. The answers were +# recorded, so a retry must ask the model again instead of replaying them into the same error. +_MODEL_AT_FAULT = ( + ModelBehaviorError, + ModelRefusalError, + MaxTurnsExceeded, + OutputGuardrailTripwireTriggered, + ToolInputGuardrailTripwireTriggered, +) + +# The tool results the SDK hands the model as structured content rather than as text. +_STRUCTURED_OUTPUTS = (ToolOutputText, ToolOutputImage, ToolOutputFileContent) + +# Arguments that keep the conversation on OpenAI's servers, which a replayed model call never reaches. +_SERVER_SIDE_STATE = ("previous_response_id", "conversation_id", "auto_previous_response_id") + + +async def run_durably( + starting_agent: Agent[Any], + input: str | list[TResponseInputItem], + *, + run_config: RunConfig | None = None, + **kwargs: Any, +) -> RunResult: + """ + Run an agent like ``Runner.run``, so that a retry of the task replays what already completed. + + Each model call and each function tool call the run makes is recorded in the task's + durable journal as it completes. When the task fails and Airflow retries it, the agent + runs again from the start, and every call that asks what the previous attempt asked, + in the same order, returns the recorded result instead of calling the model or the tool + again. From the first call that differs, such as a model answer that changed, that call + and every call after it run live. Once the run succeeds what was recorded is deleted, so + clearing the task later starts the agent over. + + .. code-block:: python + + from agents import Agent + + from airflow.providers.common.ai.durable.openai_agents import run_durably_sync + from airflow.providers.common.ai.tools.openai_agents import AirflowTools + from airflow.sdk import task + + + @task(retries=3) + def research(question: str) -> str: + agent = Agent(name="analyst", model=model, tools=AirflowTools(warehouse)) + return run_durably_sync(agent, question).final_output + + Recorded are the calls of every model given by name or in ``run_config``, the model and + function tools of the starting agent and of the agents in its ``handoffs`` list, and an + agent used through ``as_tool``, as one call. These run live on every attempt: tools + other than function tools (hosted tools such as web search, the tools of + ``mcp_servers``, shell and computer tools), a function tool whose ``needs_approval`` is + a function, since the SDK checks its arguments only through its own invoker, the + function tools of an agent reached through ``handoff()`` and its model unless that is + given by name or left to the default, and streamed runs. 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 + before the first model call, so that call no longer asks what it asked then. + + Outside a task there is no journal, and this is ``Runner.run``. + + :param starting_agent: The agent to run. + :param input: The prompt, as for ``Runner.run``. + :param run_config: As for ``Runner.run``. + :param kwargs: Passed to ``Runner.run``, such as ``max_turns`` or ``context``. + :raises ValueError: When ``previous_response_id``, ``conversation_id`` or + ``auto_previous_response_id`` is passed: a replayed call never reaches the + conversation OpenAI keeps for them. + """ + journal = current_journal() + if journal is None: + return await Runner.run(starting_agent, input, run_config=run_config, **kwargs) + if stateful := [name for name in _SERVER_SIDE_STATE if kwargs.get(name)]: + raise ValueError( + f"run_durably() cannot replay a run that keeps its conversation on OpenAI's servers " + f"({', '.join(stateful)}): a replayed model call never reaches it." + ) + durable_run = journal.start_run() + config = copy.copy(run_config) if run_config is not None else RunConfig() + config.model_provider = _DurableModelProvider(config.model_provider, durable_run) + if isinstance(config.model, Model): + config.model = _DurableModel(config.model, durable_run) + agent = _durable_agent(starting_agent, durable_run, {}) + try: + result = await Runner.run(agent, input, run_config=config, **kwargs) + except BaseException as error: + durable_run.fail(error, model_at_fault=isinstance(error, _MODEL_AT_FAULT)) + raise + if journal.clean_up_after_run: + durable_run.cleanup() + return result + + +def run_durably_sync( + starting_agent: Agent[Any], + input: str | list[TResponseInputItem], + *, + run_config: RunConfig | None = None, + **kwargs: Any, +) -> RunResult: + """Run :func:`run_durably` from synchronous code, such as a plain ``@task``.""" + return run_coroutine_sync(run_durably(starting_agent, input, run_config=run_config, **kwargs)) + + +def _durable_agent(agent: Agent[Any], durable_run: DurableRun, seen: dict[int, Agent[Any]]) -> Agent[Any]: + """Return a copy of ``agent`` whose model, function tools and handoff agents go through the journal.""" + if (done := seen.get(id(agent))) is not None: + return done + durable = agent.clone() + # Recorded before the handoffs are copied, so agents that hand off to each other copy once. + seen[id(agent)] = durable + if isinstance(agent.model, Model): + durable.model = _DurableModel(agent.model, durable_run) + durable.tools = [_durable_tool(tool, durable_run) for tool in agent.tools] + durable.handoffs = [ + _durable_agent(target, durable_run, seen) if isinstance(target, Agent) else target + for target in agent.handoffs + ] + return durable + + +def _durable_tool(tool: Tool, durable_run: DurableRun) -> Tool: + # A tool whose approval depends on its arguments is left as it is: the SDK checks that the + # arguments it approved are the ones the tool runs with only through its own invoker. + if not isinstance(tool, FunctionTool) or callable(tool.needs_approval): + return tool + invoke = tool.on_invoke_tool + + async def on_invoke_tool(ctx: ToolContext[Any], arguments: str) -> Any: + if _inside_a_call(durable_run): + return await invoke(ctx, arguments) + try: + parsed = json.loads(arguments) if arguments else {} + except json.JSONDecodeError: + parsed = arguments + fingerprint = fingerprint_tool_call( + tool.name, + # As recorded model answers are masked, so a replayed call's arguments can be too. + mask_secrets(parsed if isinstance(parsed, dict) else {"arguments": parsed}), + ctx.tool_call_id, + ) + return await durable_run.run( + f"tool:{tool.name}", + kind="tool", + fingerprint=fingerprint, + body=lambda: invoke(ctx, arguments), + # Secrets must not reach the journal any more than the model. + to_record=lambda output: mask_secrets(_as_sent(output)), + ) + + # A shallow copy, not dataclasses.replace: the tool's schema is already prepared. + durable = copy.copy(tool) + durable.on_invoke_tool = on_invoke_tool + return durable + + +def _as_sent(output: Any) -> Any: + """ + Return a tool's result in the form the SDK sends it to the model, as a JSON-compatible value. + + The SDK sends a ``ToolOutputText``, ``ToolOutputImage`` or ``ToolOutputFileContent``, a + ``dict`` with a ``type``, or a list of them, as structured content, and anything else as + ``str(output)``. Recording anything else as it is would replay it as a different text + (a model's fields as a dict, a tuple as a list), so the next model call would no longer + match the one the failed try made. + """ + if isinstance(output, str): + return output + items = output if isinstance(output, list | tuple) else [output] + if items and all( + isinstance(item, _STRUCTURED_OUTPUTS) or (isinstance(item, dict) and "type" in item) for item in items + ): + return to_jsonable_python(output) + return str(output) + + +def _inside_a_call(durable_run: DurableRun) -> bool: + """ + Whether a recorded call of ``durable_run`` is running, such as an agent called through ``as_tool``. + + That call is recorded as one step, and a retry replays it whole, so what it does inside + runs live rather than taking positions of the run it belongs to. + """ + step = _EXECUTING_STEP.get() + return step is not None and step.durable_run is durable_run + + +class _DurableModelProvider(ModelProvider): + """Wrap every model the run resolves by name in a :class:`_DurableModel`.""" + + def __init__(self, provider: ModelProvider, durable_run: DurableRun) -> None: + self._provider = provider + self._durable_run = durable_run + self._models: dict[str | None, _DurableModel] = {} + + def get_model(self, model_name: str | None) -> Model: + if model_name not in self._models: + self._models[model_name] = _DurableModel(self._provider.get_model(model_name), self._durable_run) + return self._models[model_name] + + +class _DurableModel(Model): + """Replay a model call the previous attempt recorded, or make it and record the response.""" + + def __init__(self, model: Model, durable_run: DurableRun) -> None: + self._model = model + self._durable_run = durable_run + # The step of a call that raised, which the SDK's own model retry may make again. + self._failed: JournalStep | None = None + + async def get_response( + self, + system_instructions: str | None, + input: str | list[TResponseInputItem], + model_settings: ModelSettings, + tools: list[Tool], + output_schema: AgentOutputSchemaBase | None, + handoffs: list[Handoff], + tracing: ModelTracing, + **kwargs: Any, + ) -> ModelResponse: + if _inside_a_call(self._durable_run): + return await self._model.get_response( + system_instructions, input, model_settings, tools, output_schema, handoffs, tracing, **kwargs + ) + fingerprint = _fingerprint_request( + self._model, system_instructions, input, model_settings, tools, output_schema, handoffs, kwargs + ) + step = self._retried(fingerprint) or self._durable_run.claim( + "model", kind="model", fingerprint=fingerprint + ) + if step.replayed: + try: + return ModelResponse( + output=_OUTPUT_ITEMS.validate_python(step.payload["output"]), + usage=Usage(), + response_id=step.payload.get("response_id"), + ) + except (ValidationError, KeyError, TypeError) as error: Review Comment: Could you add a test for this fallback, with a recorded payload that `_OUTPUT_ITEMS` rejects? It should check that this call and later calls go live. --- Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting ########## providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py: ########## @@ -0,0 +1,448 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Replay an OpenAI Agents SDK agent's completed model and tool calls when Airflow retries its task. + +.. note:: Experimental; see :ref:`howto/stability`. +""" + +from __future__ import annotations + +import copy +import json +from typing import TYPE_CHECKING, Any + +from pydantic import TypeAdapter, ValidationError +from pydantic_core import to_jsonable_python + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + +try: + from agents import ( + Agent, + FunctionTool, + Model, + ModelProvider, + ModelResponse, + RunConfig, + Runner, + ToolOutputFileContent, + ToolOutputImage, + ToolOutputText, + Usage, + ) + from agents.exceptions import ( + MaxTurnsExceeded, + ModelBehaviorError, + ModelRefusalError, + OutputGuardrailTripwireTriggered, + ToolInputGuardrailTripwireTriggered, + ) + from agents.items import TResponseOutputItem +except ImportError as e: + raise AirflowOptionalProviderFeatureException(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.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 + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from agents import AgentOutputSchemaBase, Handoff, ModelSettings, ModelTracing, RunResult, Tool + from agents.items import TResponseInputItem, TResponseStreamEvent + from agents.retry import ModelRetryAdvice, ModelRetryAdviceRequest + from agents.tool_context import ToolContext + + from airflow.providers.common.ai.durable.journal import DurableRun, JournalStep + +__all__ = ["run_durably", "run_durably_sync"] + +log = get_task_logger() + +_OUTPUT_ITEMS: TypeAdapter[list[TResponseOutputItem]] = TypeAdapter(list[TResponseOutputItem]) + +# Errors that come from the model's answers rather than from a step. The answers were +# recorded, so a retry must ask the model again instead of replaying them into the same error. +_MODEL_AT_FAULT = ( + ModelBehaviorError, + ModelRefusalError, + MaxTurnsExceeded, + OutputGuardrailTripwireTriggered, + ToolInputGuardrailTripwireTriggered, +) + +# The tool results the SDK hands the model as structured content rather than as text. +_STRUCTURED_OUTPUTS = (ToolOutputText, ToolOutputImage, ToolOutputFileContent) + +# Arguments that keep the conversation on OpenAI's servers, which a replayed model call never reaches. +_SERVER_SIDE_STATE = ("previous_response_id", "conversation_id", "auto_previous_response_id") + + +async def run_durably( + starting_agent: Agent[Any], + input: str | list[TResponseInputItem], + *, + run_config: RunConfig | None = None, + **kwargs: Any, +) -> RunResult: + """ + Run an agent like ``Runner.run``, so that a retry of the task replays what already completed. + + Each model call and each function tool call the run makes is recorded in the task's + durable journal as it completes. When the task fails and Airflow retries it, the agent + runs again from the start, and every call that asks what the previous attempt asked, + in the same order, returns the recorded result instead of calling the model or the tool + again. From the first call that differs, such as a model answer that changed, that call + and every call after it run live. Once the run succeeds what was recorded is deleted, so + clearing the task later starts the agent over. + + .. code-block:: python + + from agents import Agent + + from airflow.providers.common.ai.durable.openai_agents import run_durably_sync + from airflow.providers.common.ai.tools.openai_agents import AirflowTools + from airflow.sdk import task + + + @task(retries=3) + def research(question: str) -> str: + agent = Agent(name="analyst", model=model, tools=AirflowTools(warehouse)) + return run_durably_sync(agent, question).final_output + + Recorded are the calls of every model given by name or in ``run_config``, the model and + function tools of the starting agent and of the agents in its ``handoffs`` list, and an + agent used through ``as_tool``, as one call. These run live on every attempt: tools + other than function tools (hosted tools such as web search, the tools of + ``mcp_servers``, shell and computer tools), a function tool whose ``needs_approval`` is + a function, since the SDK checks its arguments only through its own invoker, the + function tools of an agent reached through ``handoff()`` and its model unless that is + given by name or left to the default, and streamed runs. 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 + before the first model call, so that call no longer asks what it asked then. + + Outside a task there is no journal, and this is ``Runner.run``. + + :param starting_agent: The agent to run. + :param input: The prompt, as for ``Runner.run``. + :param run_config: As for ``Runner.run``. + :param kwargs: Passed to ``Runner.run``, such as ``max_turns`` or ``context``. + :raises ValueError: When ``previous_response_id``, ``conversation_id`` or + ``auto_previous_response_id`` is passed: a replayed call never reaches the + conversation OpenAI keeps for them. + """ + journal = current_journal() + if journal is None: + return await Runner.run(starting_agent, input, run_config=run_config, **kwargs) + if stateful := [name for name in _SERVER_SIDE_STATE if kwargs.get(name)]: + raise ValueError( + f"run_durably() cannot replay a run that keeps its conversation on OpenAI's servers " + f"({', '.join(stateful)}): a replayed model call never reaches it." + ) + durable_run = journal.start_run() + config = copy.copy(run_config) if run_config is not None else RunConfig() + config.model_provider = _DurableModelProvider(config.model_provider, durable_run) + if isinstance(config.model, Model): + config.model = _DurableModel(config.model, durable_run) + agent = _durable_agent(starting_agent, durable_run, {}) + try: + result = await Runner.run(agent, input, run_config=config, **kwargs) + except BaseException as error: + durable_run.fail(error, model_at_fault=isinstance(error, _MODEL_AT_FAULT)) + raise + if journal.clean_up_after_run: + durable_run.cleanup() + return result + + +def run_durably_sync( + starting_agent: Agent[Any], + input: str | list[TResponseInputItem], + *, + run_config: RunConfig | None = None, + **kwargs: Any, +) -> RunResult: + """Run :func:`run_durably` from synchronous code, such as a plain ``@task``.""" + return run_coroutine_sync(run_durably(starting_agent, input, run_config=run_config, **kwargs)) + + +def _durable_agent(agent: Agent[Any], durable_run: DurableRun, seen: dict[int, Agent[Any]]) -> Agent[Any]: + """Return a copy of ``agent`` whose model, function tools and handoff agents go through the journal.""" + if (done := seen.get(id(agent))) is not None: + return done + durable = agent.clone() + # Recorded before the handoffs are copied, so agents that hand off to each other copy once. + seen[id(agent)] = durable + if isinstance(agent.model, Model): + durable.model = _DurableModel(agent.model, durable_run) + durable.tools = [_durable_tool(tool, durable_run) for tool in agent.tools] + durable.handoffs = [ + _durable_agent(target, durable_run, seen) if isinstance(target, Agent) else target + for target in agent.handoffs + ] + return durable + + +def _durable_tool(tool: Tool, durable_run: DurableRun) -> Tool: + # A tool whose approval depends on its arguments is left as it is: the SDK checks that the + # arguments it approved are the ones the tool runs with only through its own invoker. + if not isinstance(tool, FunctionTool) or callable(tool.needs_approval): + return tool + invoke = tool.on_invoke_tool + + async def on_invoke_tool(ctx: ToolContext[Any], arguments: str) -> Any: + if _inside_a_call(durable_run): + return await invoke(ctx, arguments) + try: + parsed = json.loads(arguments) if arguments else {} + except json.JSONDecodeError: + parsed = arguments + fingerprint = fingerprint_tool_call( + tool.name, + # As recorded model answers are masked, so a replayed call's arguments can be too. + mask_secrets(parsed if isinstance(parsed, dict) else {"arguments": parsed}), + ctx.tool_call_id, + ) + return await durable_run.run( + f"tool:{tool.name}", + kind="tool", + fingerprint=fingerprint, + body=lambda: invoke(ctx, arguments), + # Secrets must not reach the journal any more than the model. + to_record=lambda output: mask_secrets(_as_sent(output)), + ) + + # A shallow copy, not dataclasses.replace: the tool's schema is already prepared. + durable = copy.copy(tool) + durable.on_invoke_tool = on_invoke_tool + return durable + + +def _as_sent(output: Any) -> Any: + """ + Return a tool's result in the form the SDK sends it to the model, as a JSON-compatible value. + + The SDK sends a ``ToolOutputText``, ``ToolOutputImage`` or ``ToolOutputFileContent``, a + ``dict`` with a ``type``, or a list of them, as structured content, and anything else as + ``str(output)``. Recording anything else as it is would replay it as a different text + (a model's fields as a dict, a tuple as a list), so the next model call would no longer + match the one the failed try made. + """ + if isinstance(output, str): + return output + items = output if isinstance(output, list | tuple) else [output] + if items and all( Review Comment: Could you add a list of structured outputs and a dict with a `type` key to the `test_a_tool_result_replays_as_the_model_saw_it` parametrize? Neither shape is covered. --- Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting ########## providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py: ########## @@ -0,0 +1,448 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Replay an OpenAI Agents SDK agent's completed model and tool calls when Airflow retries its task. + +.. note:: Experimental; see :ref:`howto/stability`. +""" + +from __future__ import annotations + +import copy +import json +from typing import TYPE_CHECKING, Any + +from pydantic import TypeAdapter, ValidationError +from pydantic_core import to_jsonable_python + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + +try: + from agents import ( + Agent, + FunctionTool, + Model, + ModelProvider, + ModelResponse, + RunConfig, + Runner, + ToolOutputFileContent, + ToolOutputImage, + ToolOutputText, + Usage, + ) + from agents.exceptions import ( + MaxTurnsExceeded, + ModelBehaviorError, + ModelRefusalError, + OutputGuardrailTripwireTriggered, + ToolInputGuardrailTripwireTriggered, + ) + from agents.items import TResponseOutputItem +except ImportError as e: + raise AirflowOptionalProviderFeatureException(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.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 + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + from agents import AgentOutputSchemaBase, Handoff, ModelSettings, ModelTracing, RunResult, Tool + from agents.items import TResponseInputItem, TResponseStreamEvent + from agents.retry import ModelRetryAdvice, ModelRetryAdviceRequest + from agents.tool_context import ToolContext + + from airflow.providers.common.ai.durable.journal import DurableRun, JournalStep + +__all__ = ["run_durably", "run_durably_sync"] + +log = get_task_logger() + +_OUTPUT_ITEMS: TypeAdapter[list[TResponseOutputItem]] = TypeAdapter(list[TResponseOutputItem]) + +# Errors that come from the model's answers rather than from a step. The answers were +# recorded, so a retry must ask the model again instead of replaying them into the same error. +_MODEL_AT_FAULT = ( + ModelBehaviorError, + ModelRefusalError, + MaxTurnsExceeded, + OutputGuardrailTripwireTriggered, + ToolInputGuardrailTripwireTriggered, +) + +# The tool results the SDK hands the model as structured content rather than as text. +_STRUCTURED_OUTPUTS = (ToolOutputText, ToolOutputImage, ToolOutputFileContent) + +# Arguments that keep the conversation on OpenAI's servers, which a replayed model call never reaches. +_SERVER_SIDE_STATE = ("previous_response_id", "conversation_id", "auto_previous_response_id") + + +async def run_durably( + starting_agent: Agent[Any], + input: str | list[TResponseInputItem], + *, + run_config: RunConfig | None = None, + **kwargs: Any, +) -> RunResult: + """ + Run an agent like ``Runner.run``, so that a retry of the task replays what already completed. + + Each model call and each function tool call the run makes is recorded in the task's + durable journal as it completes. When the task fails and Airflow retries it, the agent + runs again from the start, and every call that asks what the previous attempt asked, + in the same order, returns the recorded result instead of calling the model or the tool + again. From the first call that differs, such as a model answer that changed, that call + and every call after it run live. Once the run succeeds what was recorded is deleted, so + clearing the task later starts the agent over. + + .. code-block:: python + + from agents import Agent + + from airflow.providers.common.ai.durable.openai_agents import run_durably_sync + from airflow.providers.common.ai.tools.openai_agents import AirflowTools + from airflow.sdk import task + + + @task(retries=3) + def research(question: str) -> str: + agent = Agent(name="analyst", model=model, tools=AirflowTools(warehouse)) + return run_durably_sync(agent, question).final_output + + Recorded are the calls of every model given by name or in ``run_config``, the model and + function tools of the starting agent and of the agents in its ``handoffs`` list, and an + agent used through ``as_tool``, as one call. These run live on every attempt: tools + other than function tools (hosted tools such as web search, the tools of + ``mcp_servers``, shell and computer tools), a function tool whose ``needs_approval`` is + a function, since the SDK checks its arguments only through its own invoker, the + function tools of an agent reached through ``handoff()`` and its model unless that is Review Comment: Could you reword this so it says which model setting is recorded for a `handoff()` agent? "That" could mean the agent, the model or the `handoff()`. --- Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting ########## providers/common/ai/docs/frameworks/openai_agents.rst: ########## @@ -145,11 +145,65 @@ To turn the SDK's tracing off altogether, call ``agents.set_tracing_disabled(Tru ``OPENAI_AGENTS_DISABLE_TRACING=1`` in the worker's environment. That stops ``AirflowTraceProcessor`` from receiving anything too. +.. _howto/frameworks:openai-agents-durable: Review Comment: Could you mention the setup here (Airflow 3.3, or `[common.ai] durable_cache_path` before that)? The example docstring has it, but the guide doesn't. --- Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting ########## providers/common/ai/docs/frameworks/openai_agents.rst: ########## @@ -145,11 +145,65 @@ To turn the SDK's tracing off altogether, call ``agents.set_tracing_disabled(Tru ``OPENAI_AGENTS_DISABLE_TRACING=1`` in the worker's environment. That stops ``AirflowTraceProcessor`` from receiving anything too. +.. _howto/frameworks:openai-agents-durable: + +Durable execution +----------------- + +When a task that calls ``Runner.run`` fails and Airflow retries it, the agent starts +again from the prompt: every model call is paid for again, and every tool call happens +again. On a task with ``retries``, run the agent with +:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably_sync`, or +:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably` from async code, +instead. They take the same arguments as ``Runner.run_sync`` and ``Runner.run``: Review Comment: Could you add "except the three conversation arguments, which it refuses"? The guide says this later, so "same arguments" is inconsistent. --- Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting ########## providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py: ########## @@ -0,0 +1,417 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import contextlib +import datetime +import json +from typing import Any + +import pytest + +pytest.importorskip("agents") + +from agents import ( + Agent, + Model, + ModelProvider, + ModelResponse, + ModelSettings, + RunConfig, + ToolOutputText, + Usage, + function_tool, +) +from agents.exceptions import MaxTurnsExceeded +from agents.retry import ModelRetryBackoffSettings, ModelRetrySettings +from openai.types.responses import ResponseFunctionToolCall, ResponseOutputMessage, ResponseOutputText +from pydantic import BaseModel + +from airflow.providers.common.ai.durable.journal import DurableJournal, journal_scope +from airflow.providers.common.ai.durable.openai_agents import _durable_tool, run_durably_sync +from airflow.providers.common.ai.tools import AirflowTool, ToolResult +from airflow.providers.common.ai.tools.openai_agents import AirflowTools + +from unit.common.ai.durable.memory_storage import MemoryStorage + +# Tracing off, so no test exports spans to OpenAI. +NO_TRACING = RunConfig(tracing_disabled=True) + + +class _ScriptedModel(Model): + """ + Calls ``lookup`` until it has ``lookups`` results back, then answers with the last one. + + ``calls`` counts the requests that reached the model. ``fail_on_answer`` makes the + request that would answer raise instead, as a crashed worker would leave it. + """ + + def __init__( + self, + *, + lookups: int = 1, + fail_on_answer: bool = False, + usage: int = 7, + argument: str = "key", + fail_first_request: bool = False, + ) -> None: + self.lookups = lookups + self.argument = argument + self.fail_first_request = fail_first_request + self.fail_on_answer = fail_on_answer + self.usage = usage + self.calls = 0 + # How many tool results each request that reached the model carried. + self.tool_results_seen: list[int] = [] + + async def get_response(self, system_instructions, input, *args: Any, **kwargs: Any) -> ModelResponse: + self.calls += 1 + if self.fail_first_request and self.calls == 1: + raise ConnectionError("connection reset") + outputs = [ + item["output"] + for item in (input if isinstance(input, list) else []) + if isinstance(item, dict) and item.get("type") == "function_call_output" + ] + self.tool_results_seen.append(len(outputs)) + output: list[ResponseFunctionToolCall | ResponseOutputMessage] + if len(outputs) < self.lookups: + output = [ + ResponseFunctionToolCall( + id=f"fc-{len(outputs)}", + call_id=f"call-{len(outputs)}", + name="lookup", + arguments=json.dumps({self.argument: f"k{len(outputs)}"}), + type="function_call", + ) + ] + else: + if self.fail_on_answer: + raise RuntimeError("worker lost") + text = ResponseOutputText(text=f"answer: {outputs[-1]}", type="output_text", annotations=[]) + output = [ + ResponseOutputMessage( + id="msg", role="assistant", status="completed", type="message", content=[text] + ) + ] + usage = Usage( + requests=1, input_tokens=self.usage, output_tokens=self.usage, total_tokens=2 * self.usage + ) + return ModelResponse(output=list(output), usage=usage, response_id=f"resp-{self.calls}") + + def stream_response(self, *args: Any, **kwargs: Any): + raise NotImplementedError + + +class _CountingTool: + """A ``lookup`` tool that counts the calls that reached it.""" + + def __init__(self, result: str = "42") -> None: + self.result = result + self.calls = 0 + + def tool(self) -> AirflowTool: + async def function(arguments: dict[str, Any]) -> ToolResult: + self.calls += 1 + return ToolResult(f"{arguments['key']}={self.result}") + + return AirflowTool( + name="lookup", + description="Look a key up.", + parameters={"type": "object", "properties": {"key": {"type": "string"}}}, + function=function, + ) + + +class _Provider(ModelProvider): + """Resolves every model name to one model.""" + + def __init__(self, model: Model) -> None: + self.model = model + + def get_model(self, model_name: str | None) -> Model: + return self.model + + +class _Report(BaseModel): + rows: int + + +def _attempt(storage: MemoryStorage, model: Model, tool: _CountingTool, 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) + with journal_scope(journal): + return run_durably_sync(agent, prompt, run_config=NO_TRACING, **kwargs), journal + + +class TestRunDurably: + def test_a_retry_replays_the_calls_the_failed_attempt_completed(self): + storage, tool = MemoryStorage(), _CountingTool() + with pytest.raises(RuntimeError, match="worker lost"): + _attempt(storage, _ScriptedModel(fail_on_answer=True), tool) + assert tool.calls == 1 + + retry_model = _ScriptedModel() + result, journal = _attempt(storage, retry_model, tool) + + assert result.final_output == "answer: k0=42" + # The first request and the tool call replay; only the request that answers is new. + assert retry_model.calls == 1 + assert tool.calls == 1 + assert journal.stats.replayed == {"model": 1, "tool": 1} + + def test_every_completed_call_replays_not_only_the_first(self): + """A replayed answer must turn into the same next request as the live one did.""" + storage, tool = MemoryStorage(), _CountingTool() + with pytest.raises(RuntimeError, match="worker lost"): + _attempt(storage, _ScriptedModel(lookups=3, fail_on_answer=True), tool) + + retry_model = _ScriptedModel(lookups=3) + _, journal = _attempt(storage, retry_model, tool) + + assert retry_model.calls == 1 + assert tool.calls == 3 + assert journal.stats.replayed == {"model": 3, "tool": 3} + + def test_the_sdks_generated_prompt_cache_key_does_not_stop_a_replay(self): + """The runner adds a key of its own to every request of a model that accepts one.""" + + class _CachingModel(_ScriptedModel): + def _supports_default_prompt_cache_key(self) -> bool: + return True + + storage, tool = MemoryStorage(), _CountingTool() + with pytest.raises(RuntimeError): + _attempt(storage, _CachingModel(fail_on_answer=True), tool) + + retry_model = _CachingModel() + _attempt(storage, retry_model, tool) + + assert retry_model.calls == 1 + + def test_a_replayed_model_call_reports_no_usage(self): + storage, tool = MemoryStorage(), _CountingTool() + with pytest.raises(RuntimeError): + _attempt(storage, _ScriptedModel(fail_on_answer=True, usage=7), tool) + + result, _ = _attempt(storage, _ScriptedModel(usage=7), tool) + + assert result.context_wrapper.usage.input_tokens == 7 + + def test_a_successful_run_deletes_what_it_recorded(self): + storage, tool = MemoryStorage(), _CountingTool() + + _attempt(storage, _ScriptedModel(), tool) + + assert storage.entries == {} + + def test_a_different_prompt_runs_every_call_again(self): + storage, tool = MemoryStorage(), _CountingTool() + with pytest.raises(RuntimeError): + _attempt(storage, _ScriptedModel(fail_on_answer=True), tool, prompt="first question") + + retry_model = _ScriptedModel() + _attempt(storage, retry_model, tool, prompt="second question") + + assert retry_model.calls == 2 + assert tool.calls == 2 + + def test_a_run_the_model_failed_asks_the_model_again(self): + """Running out of turns comes from the model's answers, which would replay into the same failure.""" + storage, tool = MemoryStorage(), _CountingTool() + with pytest.raises(MaxTurnsExceeded): + _attempt(storage, _ScriptedModel(lookups=3), tool, max_turns=2) + assert tool.calls == 2 + + retry_model = _ScriptedModel(lookups=1) + result, _ = _attempt(storage, retry_model, tool, max_turns=5) + + # The calls before the model's last request replay; that request is asked again, so + # the retry answers after the first lookup instead of replaying the second. + assert result.final_output == "answer: k0=42" + assert retry_model.calls == 1 + assert tool.calls == 2 + + def test_a_model_resolved_by_name_is_recorded(self): + storage, tool = MemoryStorage(), _CountingTool() + + def attempt(model: _ScriptedModel) -> None: + agent = Agent(name="analyst", model="scripted", tools=AirflowTools(tool.tool())) + config = RunConfig(tracing_disabled=True, model_provider=_Provider(model)) + with journal_scope(DurableJournal(storage, clean_up_after_run=True)): + run_durably_sync(agent, "go", run_config=config) + + with pytest.raises(RuntimeError): + attempt(_ScriptedModel(fail_on_answer=True)) + retry_model = _ScriptedModel() + attempt(retry_model) + + assert retry_model.calls == 1 + + def test_an_agent_used_as_a_tool_replays_as_one_call(self): + """Its own model calls run inside the outer tool call and must not take the outer run's positions.""" + storage, tool = MemoryStorage(), _CountingTool() + + def attempt(outer: _ScriptedModel, inner: _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.")], + ) + config = RunConfig(tracing_disabled=True, model_provider=_Provider(inner)) + with journal_scope(DurableJournal(storage, clean_up_after_run=True)): + run_durably_sync(lead, "go", run_config=config) + + # ``as_tool`` takes the question for the specialist as ``input``. + with pytest.raises(RuntimeError): + attempt(_ScriptedModel(lookups=2, fail_on_answer=True, argument="input"), _ScriptedModel()) + assert tool.calls == 2 + outer, inner = _ScriptedModel(lookups=2, argument="input"), _ScriptedModel() + attempt(outer, inner) + + # Both outer requests for the specialist and both of its answers replay. + assert outer.calls == 1 + assert inner.calls == 0 + + @pytest.mark.enable_redact + def test_a_secret_in_a_tool_result_is_not_recorded(self, registered_secret): + """Any function tool, not only the Airflow tools, whose results are masked already.""" + + @function_tool + def lookup(key: str) -> str: + """Look a key up.""" + return f"{key}={registered_secret}" + + agent = Agent(name="analyst", model=_ScriptedModel(fail_on_answer=True), tools=[lookup]) + storage = MemoryStorage() + + with journal_scope(DurableJournal(storage, clean_up_after_run=True)), pytest.raises(RuntimeError): + run_durably_sync(agent, "go", run_config=NO_TRACING) + + (recorded,) = [step for step in storage.steps() if step["name"] == "tool:lookup"] + assert recorded["payload"] == "k0=***" + assert registered_secret not in json.dumps(storage.entries) + + @pytest.mark.parametrize( + "result", + [ + pytest.param(_Report(rows=3), id="pydantic-model"), + pytest.param(("orders", 3), id="tuple"), + pytest.param(datetime.datetime(2026, 10, 14, 9, 30), id="datetime"), + pytest.param({"rows": 3}, id="dict"), + pytest.param(ToolOutputText(text="3 rows"), id="structured"), + ], + ) + def test_a_tool_result_replays_as_the_model_saw_it(self, result): + """The SDK sends most results as str(); a replay must turn into the same next request.""" + + @function_tool + def lookup(key: str) -> Any: + """Look a key up.""" + return result + + storage = MemoryStorage() + + def attempt(model: _ScriptedModel) -> DurableJournal: + agent = Agent(name="analyst", model=model, tools=[lookup]) + journal = DurableJournal(storage, clean_up_after_run=True) + with journal_scope(journal): + run_durably_sync(agent, "go", run_config=NO_TRACING) + return journal + + with pytest.raises(RuntimeError): + attempt(_ScriptedModel(lookups=2, fail_on_answer=True)) + journal = attempt(_ScriptedModel(lookups=2)) + + assert journal.stats.replayed == {"model": 2, "tool": 2} + + @pytest.mark.parametrize( + "change", + [ + pytest.param({"instructions": "Answer in French."}, id="instructions"), + pytest.param({"model_settings": ModelSettings(temperature=0.1)}, id="settings"), + pytest.param({"output_type": _Report}, id="output-type"), + pytest.param({"handoffs": [Agent(name="auditor")]}, id="handoffs"), + pytest.param({"tools": []}, id="tools"), + ], + ) + def test_a_request_that_changed_between_attempts_runs_live(self, change): + storage, tool = MemoryStorage(), _CountingTool() + with pytest.raises(RuntimeError): + _attempt(storage, _ScriptedModel(fail_on_answer=True), tool) + + retry_model = _ScriptedModel() + agent = Agent(name="analyst", model=retry_model, tools=AirflowTools(tool.tool())).clone(**change) + # Some changes make the scripted answers invalid for the agent; only the first request matters. + with journal_scope(DurableJournal(storage, clean_up_after_run=True)), contextlib.suppress(Exception): + run_durably_sync(agent, "go", run_config=NO_TRACING) + + # The first request reached the model, rather than replaying the failed try's answer. + assert retry_model.tool_results_seen[0] == 0 + + def test_a_request_the_sdk_retried_keeps_its_position(self): + """The SDK's model retry makes the same call again; the next try must still line up.""" + settings = ModelSettings( + retry=ModelRetrySettings( + max_retries=1, + backoff=ModelRetryBackoffSettings(initial_delay=0, max_delay=0, jitter=False), + policy=lambda context: True, + ) + ) + storage, tool = MemoryStorage(), _CountingTool() + + def attempt(model: _ScriptedModel) -> DurableJournal: + agent = Agent( + name="analyst", model=model, tools=AirflowTools(tool.tool()), model_settings=settings + ) + journal = DurableJournal(storage, clean_up_after_run=True) + with journal_scope(journal): + run_durably_sync(agent, "go", run_config=NO_TRACING) + return journal + + with pytest.raises(RuntimeError, match="worker lost"): + attempt(_ScriptedModel(fail_first_request=True, fail_on_answer=True)) + journal = attempt(_ScriptedModel()) + + assert journal.stats.replayed == {"model": 1, "tool": 1} + + def test_a_tool_whose_approval_depends_on_its_arguments_is_left_to_the_sdk(self): + async def policy(context, arguments, call_id) -> bool: + return arguments.get("amount", 0) > 100 + + @function_tool(needs_approval=policy) + def pay(amount: int = 1000) -> str: + """Pay an amount.""" + return "paid" + + assert _durable_tool(pay, DurableJournal(MemoryStorage()).start_run()) is pay + + @pytest.mark.parametrize("option", ["previous_response_id", "conversation_id"]) Review Comment: Could you add `auto_previous_response_id` to the parametrize? The code refuses it too. --- Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting ########## providers/common/ai/docs/frameworks/openai_agents.rst: ########## @@ -145,11 +145,65 @@ To turn the SDK's tracing off altogether, call ``agents.set_tracing_disabled(Tru ``OPENAI_AGENTS_DISABLE_TRACING=1`` in the worker's environment. That stops ``AirflowTraceProcessor`` from receiving anything too. +.. _howto/frameworks:openai-agents-durable: + +Durable execution +----------------- + +When a task that calls ``Runner.run`` fails and Airflow retries it, the agent starts +again from the prompt: every model call is paid for again, and every tool call happens +again. On a task with ``retries``, run the agent with +:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably_sync`, or +:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably` from async code, +instead. They take the same arguments as ``Runner.run_sync`` and ``Runner.run``: + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_durable.py + :language: python + :start-after: [START example_openai_agents_durable] + :end-before: [END example_openai_agents_durable] + +Each model call and function tool call is recorded in the task instance's durable +journal as it completes, the journal ``AgentOperator(durable=True)`` uses (see +:doc:`../durable_execution`). On a retry the agent runs again from the start, and each +call that asks what the failed try asked, in the same order, returns what it returned +then. From the first call that differs, that call and every call after it run live. Once +the run succeeds, what was recorded is deleted, so clearing the task later starts over. + +A retry runs some work again: Review Comment: Could this list be split into "What a retry replays" and "What always runs live"? Right now it mixes several kinds of cases. --- Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting ########## providers/common/ai/docs/frameworks/openai_agents.rst: ########## @@ -145,11 +145,65 @@ To turn the SDK's tracing off altogether, call ``agents.set_tracing_disabled(Tru ``OPENAI_AGENTS_DISABLE_TRACING=1`` in the worker's environment. That stops ``AirflowTraceProcessor`` from receiving anything too. +.. _howto/frameworks:openai-agents-durable: + +Durable execution +----------------- + +When a task that calls ``Runner.run`` fails and Airflow retries it, the agent starts +again from the prompt: every model call is paid for again, and every tool call happens +again. On a task with ``retries``, run the agent with +:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably_sync`, or +:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably` from async code, +instead. They take the same arguments as ``Runner.run_sync`` and ``Runner.run``: + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_durable.py + :language: python + :start-after: [START example_openai_agents_durable] + :end-before: [END example_openai_agents_durable] + +Each model call and function tool call is recorded in the task instance's durable +journal as it completes, the journal ``AgentOperator(durable=True)`` uses (see +:doc:`../durable_execution`). On a retry the agent runs again from the start, and each +call that asks what the failed try asked, in the same order, returns what it returned +then. From the first call that differs, that call and every call after it run live. Once +the run succeeds, what was recorded is deleted, so clearing the task later starts over. + +A retry runs some work again: + +- A call that was running when the try failed runs again, so tools that change data + must still be safe to repeat. +- When the run failed because of the model's answers (``MaxTurnsExceeded``, + ``ModelBehaviorError``, ``ModelRefusalError``, or an output or tool input guardrail), + replaying those answers would fail it the same way, so the retry asks the model again + from its last request. A ``ModelBehaviorError`` that a tool raises, such as arguments + that fail validation with ``failure_error_function=None``, counts as the tool's failure + instead: the tool call runs again, but the answer that asked for it replays. +- 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. +- These are not recorded and run live on every try: tools other than function tools + (hosted tools such as web search, the tools of ``mcp_servers``, shell and computer + tools); a function tool whose ``needs_approval`` is a function, since the SDK checks the + arguments it approved only through its own invoker; the function tools of an agent + reached through ``handoff()``, and its model unless that is given by name or left to Review Comment: Could you make this say which model setting is recorded for a `handoff()` agent, to match the docstring? --- Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting ########## providers/common/ai/docs/frameworks/openai_agents.rst: ########## @@ -145,11 +145,65 @@ To turn the SDK's tracing off altogether, call ``agents.set_tracing_disabled(Tru ``OPENAI_AGENTS_DISABLE_TRACING=1`` in the worker's environment. That stops ``AirflowTraceProcessor`` from receiving anything too. +.. _howto/frameworks:openai-agents-durable: + +Durable execution +----------------- + +When a task that calls ``Runner.run`` fails and Airflow retries it, the agent starts +again from the prompt: every model call is paid for again, and every tool call happens +again. On a task with ``retries``, run the agent with +:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably_sync`, or +:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably` from async code, +instead. They take the same arguments as ``Runner.run_sync`` and ``Runner.run``: + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_durable.py + :language: python + :start-after: [START example_openai_agents_durable] + :end-before: [END example_openai_agents_durable] + +Each model call and function tool call is recorded in the task instance's durable +journal as it completes, the journal ``AgentOperator(durable=True)`` uses (see +:doc:`../durable_execution`). On a retry the agent runs again from the start, and each +call that asks what the failed try asked, in the same order, returns what it returned +then. From the first call that differs, that call and every call after it run live. Once +the run succeeds, what was recorded is deleted, so clearing the task later starts over. + +A retry runs some work again: + +- A call that was running when the try failed runs again, so tools that change data + must still be safe to repeat. +- When the run failed because of the model's answers (``MaxTurnsExceeded``, + ``ModelBehaviorError``, ``ModelRefusalError``, or an output or tool input guardrail), + replaying those answers would fail it the same way, so the retry asks the model again + from its last request. A ``ModelBehaviorError`` that a tool raises, such as arguments Review Comment: Could you add a test for a tool-raised `ModelBehaviorError` that reruns the tool and replays the model answer? The guide states it as a guarantee. --- Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting ########## providers/common/ai/docs/frameworks/openai_agents.rst: ########## @@ -145,11 +145,65 @@ To turn the SDK's tracing off altogether, call ``agents.set_tracing_disabled(Tru ``OPENAI_AGENTS_DISABLE_TRACING=1`` in the worker's environment. That stops ``AirflowTraceProcessor`` from receiving anything too. +.. _howto/frameworks:openai-agents-durable: + +Durable execution +----------------- + +When a task that calls ``Runner.run`` fails and Airflow retries it, the agent starts +again from the prompt: every model call is paid for again, and every tool call happens +again. On a task with ``retries``, run the agent with +:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably_sync`, or +:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably` from async code, +instead. They take the same arguments as ``Runner.run_sync`` and ``Runner.run``: + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_durable.py + :language: python + :start-after: [START example_openai_agents_durable] + :end-before: [END example_openai_agents_durable] + +Each model call and function tool call is recorded in the task instance's durable +journal as it completes, the journal ``AgentOperator(durable=True)`` uses (see +:doc:`../durable_execution`). On a retry the agent runs again from the start, and each +call that asks what the failed try asked, in the same order, returns what it returned +then. From the first call that differs, that call and every call after it run live. Once +the run succeeds, what was recorded is deleted, so clearing the task later starts over. + +A retry runs some work again: + +- A call that was running when the try failed runs again, so tools that change data + must still be safe to repeat. +- When the run failed because of the model's answers (``MaxTurnsExceeded``, + ``ModelBehaviorError``, ``ModelRefusalError``, or an output or tool input guardrail), + replaying those answers would fail it the same way, so the retry asks the model again + from its last request. A ``ModelBehaviorError`` that a tool raises, such as arguments + that fail validation with ``failure_error_function=None``, counts as the tool's failure + instead: the tool call runs again, but the answer that asked for it replays. +- 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. +- These are not recorded and run live on every try: tools other than function tools + (hosted tools such as web search, the tools of ``mcp_servers``, shell and computer + tools); a function tool whose ``needs_approval`` is a function, since the SDK checks the + arguments it approved only through its own invoker; the function tools of an agent + reached through ``handoff()``, and its model unless that is given by name or left to + the default; and a streamed run. An agent used through ``as_tool`` is recorded as one + call. +- A ``session`` makes every retry run live: the SDK adds the failed try's prompt to it Review Comment: Could you add a test that passes a `session` to `run_durably`? It would check that the retry runs live. --- Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
