You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "kaxil (via GitHub)" <gi...@apache.org> on 2026/10/14 22:05:32 UTC
[PR] Replay an OpenAI Agents SDK agent's completed calls when Airflow retries its task [airflow]
kaxil opened a new pull request, #74770:
URL: https://github.com/apache/airflow/pull/74770
When a task running an OpenAI Agents SDK agent 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. `run_durably()` / `run_durably_sync()` take the same arguments as `Runner.run` / `Runner.run_sync` and record each model call and function tool call in the task's durable journal, the one `AgentOperator(durable=True)` uses. A retry replays every call that asks what the failed try asked, in order, and runs live from the first that differs. Once the run succeeds the recording is deleted.
```python
@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
```
From Airflow 3.3 the journal is kept in the task state store; before that, under `[common.ai] durable_cache_path`, as for `AgentOperator(durable=True)`. Stacked on the tracing PR.
## Design rationale
**Model calls are intercepted by wrapping the model, not by run hooks.** `RunHooks` and `AgentHooks` return `None`, so they can observe a call but not replace it, and `RunState.to_json()` captures a run only at an interruption, not mid-run. The SDK resolves a model from `RunConfig.model`, then the `model_provider` for names, then a `Model` on the agent, so `run_durably` wraps the provider and the `run_config` model and clones the agent graph (following `handoffs`) to wrap `Model` instances and function tools. This is the same interception point the SDK's Temporal integration uses.
**A replayed call has to turn into the same next request as the live one did.** Each request is fingerprinted from the model, instructions, conversation, settings, tools (with descriptions and strictness), output schema, handoffs and prompt, and four things keep that stable across a replay:
- Recorded answers keep only the fields the response set (`exclude_unset`), and `None` fields are left out of the fingerprint. Otherwise replayed items carry explicit `null`s the live ones do not, and every request after the first replay mismatches.
- Tool results are recorded as the SDK sends them to the model: structured `ToolOutput*` content as it is, anything else as `str(output)`. Recording a Pydantic model, tuple or datetime as JSON would replay as different text.
- The `prompt_cache_key` the SDK generates once per run is left out of the settings.
- A call the SDK's own model retry makes again reuses the failed call's position instead of taking a new one, so the next try still lines up.
**An agent used through `as_tool` is one recorded call.** Its own model calls run inside the outer tool call, so they run live rather than taking positions of the outer run; the outer call replays whole.
**A run the model failed asks the model again.** For `MaxTurnsExceeded`, `ModelBehaviorError`, `ModelRefusalError` and output or tool-input guardrail tripwires the recorded answers would replay into the same failure, so the journal reruns from the model's last request.
**Server-side conversations are refused.** `previous_response_id`, `conversation_id` and `auto_previous_response_id` keep the conversation on OpenAI's servers, which a replayed call never reaches, so `run_durably` raises `ValueError` for them.
## Known issues
- Not recorded, so they run live on every try: tools other than function tools (hosted tools, `mcp_servers` tools, shell and computer tools), a function tool whose `needs_approval` is a function (the SDK checks the approved arguments only through its own invoker, so wrapping it would weaken that check), an agent reached through `handoff()` unless its model is given by name, and streamed runs. A `session` makes every retry run live. The guide lists all of these.
- A `ModelBehaviorError` that a tool raises, such as arguments failing validation with `failure_error_function=None`, counts as the tool's failure: the tool reruns, but the answer that asked for it replays.
- Small follow-ups: passing `result.last_agent` back into `run_durably` wraps it twice; the wrapped model provider is not closed when `run_config` is omitted, which matters only for the 0.23 websocket transport; a dict `run_config` is not accepted.
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Replay an OpenAI Agents SDK agent's completed calls when Airflow retries its task [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74770:
URL: https://github.com/apache/airflow/pull/74770#discussion_r4296239688
##########
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:
Review Comment:
Done in 7715f05b85d.
##########
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]:
Review Comment:
Renamed in 7715f05b85d.
##########
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()
Review Comment:
Done in 7715f05b85d.
##########
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:
Review Comment:
Renamed, and good catch on the error path. Fixed in 7715f05b85d: a call whose exception `failure_error_function` turned into text is no longer recorded, so a retry runs it again. For a custom `failure_error_function` I took your approach: a wrapper on the copy stashes the exception per `tool_call_id`, then `claim`/`step.executing()` and `step.fail`. For the SDK's default formatter I read the flag the SDK sets on the tool context when it handled the error, instead of wrapping it, because on 0.23 a non-default formatter makes the SDK put the exception text into its own trace spans, which it leaves out for the default one. Tests cover both.
##########
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)
Review Comment:
Done in 7715f05b85d.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Replay an OpenAI Agents SDK agent's completed calls when Airflow retries its task [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil closed pull request #74770: Replay an OpenAI Agents SDK agent's completed calls when Airflow retries its task
URL: https://github.com/apache/airflow/pull/74770
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Replay an OpenAI Agents SDK agent's completed calls when Airflow retries its task [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74770:
URL: https://github.com/apache/airflow/pull/74770#discussion_r4296278664
##########
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:
Added in 7715f05b85d: a recorded answer the type adapter rejects makes that call and every call after it run live.
##########
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:
Added both cases to the parametrize in 7715f05b85d.
##########
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:
Added in 7715f05b85d.
##########
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:
Reworded in 7715f05b85d: its model is recorded only when the agent's `model` is a model name or not set.
##########
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:
Added in 7715f05b85d: the journal is kept in the task state store on Airflow 3.3 and later, and earlier versions need `[common.ai] durable_cache_path`.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Replay an OpenAI Agents SDK agent's completed calls when Airflow retries its task [airflow]
Posted by "pankajastro (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Replay an OpenAI Agents SDK agent's completed calls when Airflow retries its task [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil merged PR #74770:
URL: https://github.com/apache/airflow/pull/74770
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Replay an OpenAI Agents SDK agent's completed calls when Airflow retries its task [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74770:
URL: https://github.com/apache/airflow/pull/74770#discussion_r4296282011
##########
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:
Done in 7715f05b85d, with the reason next to it.
##########
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:
Split into "What a retry replays" and "What always runs live" in 7715f05b85d.
##########
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:
Added in 7715f05b85d: the tool runs again and the model answer before it replays.
##########
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:
Reworded to match the docstring in 7715f05b85d.
##########
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:
Added in 7715f05b85d, with a SQLite session shared across attempts: nothing replays on the retry.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Replay an OpenAI Agents SDK agent's completed calls when Airflow retries its task [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #74770:
URL: https://github.com/apache/airflow/pull/74770#discussion_r4284488865
##########
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:
Review Comment:
```suggestion
if (journal := current_journal()) is None:
```
##########
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]:
Review Comment:
```suggestion
def _get_durable_agent(agent: Agent[Any], durable_run: DurableRun, seen: dict[int, Agent[Any]]) -> Agent[Any]:
```
##########
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()
Review Comment:
```suggestion
durable_agent = agent.clone()
```
##########
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:
Review Comment:
```suggestion
def _is_inside_a_call(durable_run: DurableRun) -> bool:
```
##########
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:
+ self._durable_run.reject(step, reason=f"the recorded model response cannot be read: {error}")
+
+ async def call() -> ModelResponse:
+ return await self._model.get_response(
+ system_instructions, input, model_settings, tools, output_schema, handoffs, tracing, **kwargs
+ )
+
+ try:
+ return await step.run(call, to_record=_record_response)
+ except Exception:
+ self._failed = step
+ raise
+
+ def _retried(self, fingerprint: str | None) -> JournalStep | None:
+ """
+ Return the step of the call that just raised, when this call is the SDK retrying it.
+
+ The SDK's model retry asks again with the same request and nothing in between; it is
+ the same call, so it takes the same position, and the next try lines up with this one.
+ """
+ failed, self._failed = self._failed, None
+ if failed is None or failed.fingerprint != fingerprint:
+ return None
+ return failed if self._durable_run.position == failed.position + 1 else None
+
+ def stream_response(
Review Comment:
```suggestion
def get_stream_response(
```
##########
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:
Review Comment:
A tool that raises is recorded as a successful result, so a retry replays the error text and never reruns the tool.
The SDK's default `failure_error_function` runs inside the tool's own `on_invoke_tool` and turns the exception into text. `durable_run.run(...)` sees a normal return and records it.
##########
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:
Review Comment:
We can wrap the formatter on the shallow copy of the tool so it stashes the exception per `ctx.tool_call_id` and delegates to the SDK's formatter. Then use `claim` and `step.executing()` instead of `durable_run.run`, and call `step.fail(error)` instead of recording when the stash holds an error.
##########
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:
Review Comment:
```suggestion
def _get_durable_tool(tool: Tool, durable_run: DurableRun) -> Tool:
```
##########
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)
Review Comment:
```suggestion
durable_tool = copy.copy(tool)
```
##########
providers/common/ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_durable.py:
##########
@@ -0,0 +1,90 @@
+# 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.
+"""
+Run an OpenAI Agents SDK agent whose completed model and tool calls a task retry replays.
+
+The agent is the one in ``example_openai_agents_agent``: a model whose API key comes from
+an Airflow connection, and a ``SQLToolset``'s tools through
+:class:`~airflow.providers.common.ai.tools.openai_agents.AirflowTools`.
+:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably_sync` runs it, so
+when the task fails and Airflow retries it, the calls the failed try completed return
+what they returned then instead of running again.
+
+Before running:
+
+1. Install the SDK: ``pip install "openai-agents>=0.20.0"``.
+2. Create a connection (``LLM_CONN_ID``, default ``openai_default``) whose password is
+ an OpenAI API key. Set its host to the base URL of any other OpenAI-compatible
+ endpoint, such as a gateway or a self-hosted server.
+3. Create a database connection (``DB_CONN_ID``, default ``sql_default``) whose
+ hook is a ``DbApiHook`` (e.g. SQLite, Postgres, MySQL).
+4. On Airflow before 3.3, set ``[common.ai] durable_cache_path`` to where the journal
+ is kept; from 3.3 on it is kept in the task state store.
+"""
+
+from __future__ import annotations
+
+import os
+from datetime import timedelta
+
+from airflow.providers.common.compat.sdk import BaseHook, dag, task
+
+LLM_CONN_ID = os.environ.get("LLM_CONN_ID", "openai_default")
+LLM_MODEL = os.environ.get("LLM_MODEL", "gpt-5.6-luna")
+DB_CONN_ID = os.environ.get("DB_CONN_ID", "sql_default")
+
+DEFAULT_QUESTION = "Which tables exist, and how many rows does each contain?"
+
+
+# [START example_openai_agents_durable]
+@dag(tags=["example"])
+def example_openai_agents_durable():
+ """Answer a question about a database with an OpenAI Agents SDK agent that survives task retries."""
+
+ @task(retries=2, retry_delay=timedelta(minutes=1))
+ def run_durable_openai_agent(question: str = DEFAULT_QUESTION) -> str:
+ from agents import Agent, OpenAIChatCompletionsModel, RunConfig
+ from openai import AsyncOpenAI
+
+ from airflow.providers.common.ai.durable.openai_agents import run_durably_sync
+ from airflow.providers.common.ai.tools.openai_agents import AirflowTools
+ from airflow.providers.common.ai.toolsets.sql import SQLToolset
Review Comment:
should we move them to the top insetad?
##########
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
Review Comment:
```suggestion
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
```
same question as https://github.com/apache/airflow/pull/74770/changes#r4284539628
##########
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] = {}
Review Comment:
```suggestion
self.provider = provider
self.durable_run = durable_run
self.models: dict[str | None, _DurableModel] = {}
```
The class itself comes with a `_` already. do we need to make these `_` prefix as well?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Replay an OpenAI Agents SDK agent's completed calls when Airflow retries its task [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74770:
URL: https://github.com/apache/airflow/pull/74770#discussion_r4296272396
##########
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:
Review Comment:
Renamed in 7715f05b85d.
##########
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] = {}
Review Comment:
Dropped the `_` prefixes in 7715f05b85d. On `_DurableModel` I named the wrapped model `wrapped` rather than `model`: `_fingerprint_request` reads a model's `model` attribute for its name, so if a `_DurableModel` is ever wrapped again, a `model` attribute holding another model would fingerprint differently on every attempt.
##########
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
Review Comment:
Same here: `wrapped`, `durable_run` and `failed_step`.
##########
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:
+ self._durable_run.reject(step, reason=f"the recorded model response cannot be read: {error}")
+
+ async def call() -> ModelResponse:
+ return await self._model.get_response(
+ system_instructions, input, model_settings, tools, output_schema, handoffs, tracing, **kwargs
+ )
+
+ try:
+ return await step.run(call, to_record=_record_response)
+ except Exception:
+ self._failed = step
+ raise
+
+ def _retried(self, fingerprint: str | None) -> JournalStep | None:
+ """
+ Return the step of the call that just raised, when this call is the SDK retrying it.
+
+ The SDK's model retry asks again with the same request and nothing in between; it is
+ the same call, so it takes the same position, and the next try lines up with this one.
+ """
+ failed, self._failed = self._failed, None
+ if failed is None or failed.fingerprint != fingerprint:
+ return None
+ return failed if self._durable_run.position == failed.position + 1 else None
+
+ def stream_response(
Review Comment:
This one has to stay `stream_response`: it overrides `agents.Model.stream_response`, which the SDK calls by that name.
##########
providers/common/ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_durable.py:
##########
@@ -0,0 +1,90 @@
+# 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.
+"""
+Run an OpenAI Agents SDK agent whose completed model and tool calls a task retry replays.
+
+The agent is the one in ``example_openai_agents_agent``: a model whose API key comes from
+an Airflow connection, and a ``SQLToolset``'s tools through
+:class:`~airflow.providers.common.ai.tools.openai_agents.AirflowTools`.
+:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably_sync` runs it, so
+when the task fails and Airflow retries it, the calls the failed try completed return
+what they returned then instead of running again.
+
+Before running:
+
+1. Install the SDK: ``pip install "openai-agents>=0.20.0"``.
+2. Create a connection (``LLM_CONN_ID``, default ``openai_default``) whose password is
+ an OpenAI API key. Set its host to the base URL of any other OpenAI-compatible
+ endpoint, such as a gateway or a self-hosted server.
+3. Create a database connection (``DB_CONN_ID``, default ``sql_default``) whose
+ hook is a ``DbApiHook`` (e.g. SQLite, Postgres, MySQL).
+4. On Airflow before 3.3, set ``[common.ai] durable_cache_path`` to where the journal
+ is kept; from 3.3 on it is kept in the task state store.
+"""
+
+from __future__ import annotations
+
+import os
+from datetime import timedelta
+
+from airflow.providers.common.compat.sdk import BaseHook, dag, task
+
+LLM_CONN_ID = os.environ.get("LLM_CONN_ID", "openai_default")
+LLM_MODEL = os.environ.get("LLM_MODEL", "gpt-5.6-luna")
+DB_CONN_ID = os.environ.get("DB_CONN_ID", "sql_default")
+
+DEFAULT_QUESTION = "Which tables exist, and how many rows does each contain?"
+
+
+# [START example_openai_agents_durable]
+@dag(tags=["example"])
+def example_openai_agents_durable():
+ """Answer a question about a database with an OpenAI Agents SDK agent that survives task retries."""
+
+ @task(retries=2, retry_delay=timedelta(minutes=1))
+ def run_durable_openai_agent(question: str = DEFAULT_QUESTION) -> str:
+ from agents import Agent, OpenAIChatCompletionsModel, RunConfig
+ from openai import AsyncOpenAI
+
+ from airflow.providers.common.ai.durable.openai_agents import run_durably_sync
+ from airflow.providers.common.ai.tools.openai_agents import AirflowTools
+ from airflow.providers.common.ai.toolsets.sql import SQLToolset
Review Comment:
They stay inside the task on purpose: the Dag processor parses this file, and `agents` and `openai` only need to be installed where the task runs. The other example Dags in this provider do the same.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org