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