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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
