This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch openai-agents-durable in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 76baf91097337c2b0d6d907c6233686cb04f4197 Author: Kaxil Naik <[email protected]> AuthorDate: Wed Oct 14 21:48:44 2026 +0100 Replay an OpenAI Agents SDK agent's completed calls when Airflow retries its task run_durably() and run_durably_sync() run an agent like Runner.run, recording each model call and function tool call in the task's durable journal, the one AgentOperator(durable=True) uses. A retry replays every call that asks what the failed try asked, in order, and runs live from the first that differs. A run the model failed (max turns, malformed output, output guardrail) asks the model again from its last request; an agent called through as_tool replays as one call; server-side conversations are refused because a replay never reaches them. --- providers/common/ai/docs/durable_execution.rst | 3 +- providers/common/ai/docs/features.rst | 3 +- .../common/ai/docs/frameworks/openai_agents.rst | 58 ++- providers/common/ai/docs/stability.rst | 3 +- .../providers/common/ai/durable/openai_agents.py | 448 +++++++++++++++++++++ .../example_dags/example_openai_agents_durable.py | 90 +++++ .../unit/common/ai/durable/test_openai_agents.py | 417 +++++++++++++++++++ 7 files changed, 1017 insertions(+), 5 deletions(-) diff --git a/providers/common/ai/docs/durable_execution.rst b/providers/common/ai/docs/durable_execution.rst index 5c670045998..bf5476649ad 100644 --- a/providers/common/ai/docs/durable_execution.rst +++ b/providers/common/ai/docs/durable_execution.rst @@ -30,7 +30,8 @@ them again. Only the steps the failed attempt never finished run live. Durable execution only helps a task that has retries. To decide *whether* a failure is worth retrying, see :doc:`retry_policies`; a retried ``LLMBatchOperator`` re-attaches to its running batch instead (:ref:`llm-batch-reattach`). To resume a Strands Agents agent on a -retry, see :ref:`howto/frameworks:strands-durable`. +retry, see :ref:`howto/frameworks:strands-durable`; for an OpenAI Agents SDK agent, see +:ref:`howto/frameworks:openai-agents-durable`. Turn it on ---------- diff --git a/providers/common/ai/docs/features.rst b/providers/common/ai/docs/features.rst index 50b3dd6e5eb..f2420d51381 100644 --- a/providers/common/ai/docs/features.rst +++ b/providers/common/ai/docs/features.rst @@ -268,7 +268,8 @@ toolset except - Gives a loop on the Anthropic SDK's ``tool_runner`` the toolsets. - ``anthropic`` * - :doc:`OpenAI Agents SDK <frameworks/openai_agents>` - - Gives an OpenAI Agents SDK agent the toolsets. + - Gives an OpenAI Agents SDK agent the toolsets. Run through ``run_durably``, a retry + replays its completed calls. - None; install ``openai-agents`` * - :doc:`LangChain <hooks/langchain>` - Builds LangChain chat and embedding models from a connection, and diff --git a/providers/common/ai/docs/frameworks/openai_agents.rst b/providers/common/ai/docs/frameworks/openai_agents.rst index 093546cf3b7..a0256450b68 100644 --- a/providers/common/ai/docs/frameworks/openai_agents.rst +++ b/providers/common/ai/docs/frameworks/openai_agents.rst @@ -145,11 +145,65 @@ To turn the SDK's tracing off altogether, call ``agents.set_tracing_disabled(Tru ``OPENAI_AGENTS_DISABLE_TRACING=1`` in the worker's environment. That stops ``AirflowTraceProcessor`` from receiving anything too. +.. _howto/frameworks:openai-agents-durable: + +Durable execution +----------------- + +When a task that calls ``Runner.run`` fails and Airflow retries it, the agent starts +again from the prompt: every model call is paid for again, and every tool call happens +again. On a task with ``retries``, run the agent with +:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably_sync`, or +:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably` from async code, +instead. They take the same arguments as ``Runner.run_sync`` and ``Runner.run``: + +.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_durable.py + :language: python + :start-after: [START example_openai_agents_durable] + :end-before: [END example_openai_agents_durable] + +Each model call and function tool call is recorded in the task instance's durable +journal as it completes, the journal ``AgentOperator(durable=True)`` uses (see +:doc:`../durable_execution`). On a retry the agent runs again from the start, and each +call that asks what the failed try asked, in the same order, returns what it returned +then. From the first call that differs, that call and every call after it run live. Once +the run succeeds, what was recorded is deleted, so clearing the task later starts over. + +A retry runs some work again: + +- A call that was running when the try failed runs again, so tools that change data + must still be safe to repeat. +- When the run failed because of the model's answers (``MaxTurnsExceeded``, + ``ModelBehaviorError``, ``ModelRefusalError``, or an output or tool input guardrail), + replaying those answers would fail it the same way, so the retry asks the model again + from its last request. A ``ModelBehaviorError`` that a tool raises, such as arguments + that fail validation with ``failure_error_function=None``, counts as the tool's failure + instead: the tool call runs again, but the answer that asked for it replays. +- A retry whose prompt, instructions, model settings or tools differ from the failed + try's runs everything live from the first call that sees the difference. +- These are not recorded and run live on every try: tools other than function tools + (hosted tools such as web search, the tools of ``mcp_servers``, shell and computer + tools); a function tool whose ``needs_approval`` is a function, since the SDK checks the + arguments it approved only through its own invoker; the function tools of an agent + reached through ``handoff()``, and its model unless that is given by name or left to + the default; and a streamed run. An agent used through ``as_tool`` is recorded as one + call. +- A ``session`` makes every retry run live: the SDK adds the failed try's prompt to it + before the first model call, so that call no longer asks what it asked then. +- A ``SandboxToolset`` without ``attach_to`` is destroyed when the try ends, so the retry + starts with a new, empty sandbox that the replayed conversation does not describe. + +Recorded tool results pass through Airflow's secret masker before they are saved, and a +replayed model call reports no token usage. ``run_durably`` refuses +``previous_response_id``, ``conversation_id`` and ``auto_previous_response_id``: with them +OpenAI keeps the conversation, and a replayed call never reaches it. Outside a task there +is nothing to replay into, and it runs the agent as ``Runner.run`` does. + Differences from ``AgentOperator`` ---------------------------------- -- A task retry runs the agent again from the start, so tools that change data must be - safe to repeat. +- A task retry runs the agent again from the start unless it runs through + ``run_durably``; see :ref:`howto/frameworks:openai-agents-durable`. - There is no human review step like ``AgentOperator(enable_hitl_review=True)``, and the SDK's own ``needs_approval`` interruptions are not connected to Airflow's approval UI. diff --git a/providers/common/ai/docs/stability.rst b/providers/common/ai/docs/stability.rst index beefd7aa1eb..9298e171835 100644 --- a/providers/common/ai/docs/stability.rst +++ b/providers/common/ai/docs/stability.rst @@ -260,7 +260,8 @@ Everything this provider ships that is not in the table above is experimental. * - The Anthropic tool runner adapter (:doc:`frameworks/anthropic`) - Built on ``tool_runner``, a beta helper of the Anthropic SDK that can change between SDK releases. - * - The OpenAI Agents SDK adapter and trace processor (:doc:`frameworks/openai_agents`) + * - The OpenAI Agents SDK adapter, trace processor and ``run_durably`` + (:doc:`frameworks/openai_agents`) - Written against OpenAI Agents SDK 0.20 and 0.23. The SDK is pre-1.0 and changes its tool and runner APIs between minor releases. * - The ``common_ai.tool_calls`` metric and diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py b/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py new file mode 100644 index 00000000000..bc0af085627 --- /dev/null +++ b/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( + 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, + ) -> AsyncIterator[TResponseStreamEvent]: + # A streamed response is not recorded: it reaches the caller event by event. + return self._model.stream_response( + system_instructions, input, model_settings, tools, output_schema, handoffs, tracing, **kwargs + ) + + def _supports_default_prompt_cache_key(self) -> bool: + # The runner asks the model whether to add a prompt cache key; the answer is the wrapped model's. + supports = getattr(self._model, "_supports_default_prompt_cache_key", None) + return bool(supports()) if callable(supports) else False + + def get_retry_advice(self, request: ModelRetryAdviceRequest) -> ModelRetryAdvice | None: + return self._model.get_retry_advice(request) + + async def _cleanup_on_run_end(self, owner: object) -> None: + await self._model._cleanup_on_run_end(owner) + + async def close(self) -> None: + await self._model.close() + + +def _record_response(response: ModelResponse) -> dict[str, Any]: + # Only the fields the response set, so a replayed item turns into the same next request as + # the live one: the SDK carries every field of an item into the conversation it sends. + output = [item.model_dump(mode="json", exclude_unset=True) for item in response.output] + # Masked, as the journal's other steps are; a replay hands the model its own masked answer. + return {"output": mask_secrets(output), "response_id": response.response_id} + + +def _fingerprint_request( + model: Model, + system_instructions: str | None, + input: str | list[TResponseInputItem], + model_settings: ModelSettings, + tools: list[Tool], + output_schema: AgentOutputSchemaBase | None, + handoffs: list[Handoff], + options: dict[str, Any], +) -> str | None: + """ + Fingerprint what a model call asks: the model, instructions, conversation, settings and tools. + + ``None`` when the request cannot be serialized; the call then replays on position alone, + and only right after a call that replayed. + """ + try: + settings = to_jsonable_python(model_settings) + # Generated once per run, so it differs on every attempt without changing the request. + if isinstance(extra := settings.get("extra_args"), dict): + extra.pop("prompt_cache_key", None) + payload = { + "model": f"{type(model).__module__}.{type(model).__qualname__}:{getattr(model, 'model', None)}", + "instructions": system_instructions, + # Masked as the recorded steps are, so a replayed conversation matches a live one. + "input": _without_none(mask_secrets(to_jsonable_python(input))), + "settings": settings, + "tools": [_describe_tool(tool) for tool in tools], + "output": output_schema.json_schema() + if output_schema and not output_schema.is_plain_text() + else None, + "handoffs": [[handoff.tool_name, handoff.input_json_schema] for handoff in handoffs], + "prompt": to_jsonable_python(options.get("prompt")), + } + return _digest(payload) + except (TypeError, ValueError): + log.warning( + "Durable: could not fingerprint the model request; its recorded response replays " + "without verification" + ) + return None + + +def _without_none(value: Any) -> Any: + """ + Drop ``None`` fields, so an item that leaves a field unset matches one that sets it to ``None``. + + Which of the two the SDK sends depends on how the item was built, not on what was asked. + """ + if isinstance(value, dict): + return {key: _without_none(item) for key, item in value.items() if item is not None} + if isinstance(value, list): + return [_without_none(item) for item in value] + return value + + +def _describe_tool(tool: Tool) -> Any: + """Return what the model is told about a tool, for the request fingerprint.""" + if isinstance(tool, FunctionTool): + return [tool.name, tool.description, tool.params_json_schema, tool.strict_json_schema] + # A hosted tool's configuration, such as a file search's vector stores, without its callbacks. + return [type(tool).__name__, to_jsonable_python(tool, fallback=lambda value: type(value).__name__)] diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_durable.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_durable.py new file mode 100644 index 00000000000..ca27e8da8ab --- /dev/null +++ b/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 + + llm = BaseHook.get_connection(LLM_CONN_ID) + agent = Agent( + name="analyst", + model=OpenAIChatCompletionsModel( + model=LLM_MODEL, + openai_client=AsyncOpenAI(api_key=llm.password, base_url=llm.host or None), + ), + instructions=( + "You are a SQL analyst. Use list_tables and get_schema to explore " + "the database, then run read-only queries to answer the question." + ), + tools=AirflowTools(SQLToolset(db_conn_id=DB_CONN_ID)), + ) + # The SDK sends traces to OpenAI by default; see the guide's tracing section. + result = run_durably_sync(agent, question, run_config=RunConfig(tracing_disabled=True)) + return result.final_output + + run_durable_openai_agent() + + +# [END example_openai_agents_durable] + + +example_openai_agents_durable() diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py b/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py new file mode 100644 index 00000000000..f1bffe35775 --- /dev/null +++ b/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py @@ -0,0 +1,417 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import contextlib +import datetime +import json +from typing import Any + +import pytest + +pytest.importorskip("agents") + +from agents import ( + Agent, + Model, + ModelProvider, + ModelResponse, + ModelSettings, + RunConfig, + ToolOutputText, + Usage, + function_tool, +) +from agents.exceptions import MaxTurnsExceeded +from agents.retry import ModelRetryBackoffSettings, ModelRetrySettings +from openai.types.responses import ResponseFunctionToolCall, ResponseOutputMessage, ResponseOutputText +from pydantic import BaseModel + +from airflow.providers.common.ai.durable.journal import DurableJournal, journal_scope +from airflow.providers.common.ai.durable.openai_agents import _durable_tool, run_durably_sync +from airflow.providers.common.ai.tools import AirflowTool, ToolResult +from airflow.providers.common.ai.tools.openai_agents import AirflowTools + +from unit.common.ai.durable.memory_storage import MemoryStorage + +# Tracing off, so no test exports spans to OpenAI. +NO_TRACING = RunConfig(tracing_disabled=True) + + +class _ScriptedModel(Model): + """ + Calls ``lookup`` until it has ``lookups`` results back, then answers with the last one. + + ``calls`` counts the requests that reached the model. ``fail_on_answer`` makes the + request that would answer raise instead, as a crashed worker would leave it. + """ + + def __init__( + self, + *, + lookups: int = 1, + fail_on_answer: bool = False, + usage: int = 7, + argument: str = "key", + fail_first_request: bool = False, + ) -> None: + self.lookups = lookups + self.argument = argument + self.fail_first_request = fail_first_request + self.fail_on_answer = fail_on_answer + self.usage = usage + self.calls = 0 + # How many tool results each request that reached the model carried. + self.tool_results_seen: list[int] = [] + + async def get_response(self, system_instructions, input, *args: Any, **kwargs: Any) -> ModelResponse: + self.calls += 1 + if self.fail_first_request and self.calls == 1: + raise ConnectionError("connection reset") + outputs = [ + item["output"] + for item in (input if isinstance(input, list) else []) + if isinstance(item, dict) and item.get("type") == "function_call_output" + ] + self.tool_results_seen.append(len(outputs)) + output: list[ResponseFunctionToolCall | ResponseOutputMessage] + if len(outputs) < self.lookups: + output = [ + ResponseFunctionToolCall( + id=f"fc-{len(outputs)}", + call_id=f"call-{len(outputs)}", + name="lookup", + arguments=json.dumps({self.argument: f"k{len(outputs)}"}), + type="function_call", + ) + ] + else: + if self.fail_on_answer: + raise RuntimeError("worker lost") + text = ResponseOutputText(text=f"answer: {outputs[-1]}", type="output_text", annotations=[]) + output = [ + ResponseOutputMessage( + id="msg", role="assistant", status="completed", type="message", content=[text] + ) + ] + usage = Usage( + requests=1, input_tokens=self.usage, output_tokens=self.usage, total_tokens=2 * self.usage + ) + return ModelResponse(output=list(output), usage=usage, response_id=f"resp-{self.calls}") + + def stream_response(self, *args: Any, **kwargs: Any): + raise NotImplementedError + + +class _CountingTool: + """A ``lookup`` tool that counts the calls that reached it.""" + + def __init__(self, result: str = "42") -> None: + self.result = result + self.calls = 0 + + def tool(self) -> AirflowTool: + async def function(arguments: dict[str, Any]) -> ToolResult: + self.calls += 1 + return ToolResult(f"{arguments['key']}={self.result}") + + return AirflowTool( + name="lookup", + description="Look a key up.", + parameters={"type": "object", "properties": {"key": {"type": "string"}}}, + function=function, + ) + + +class _Provider(ModelProvider): + """Resolves every model name to one model.""" + + def __init__(self, model: Model) -> None: + self.model = model + + def get_model(self, model_name: str | None) -> Model: + return self.model + + +class _Report(BaseModel): + rows: int + + +def _attempt(storage: MemoryStorage, model: Model, tool: _CountingTool, prompt: str = "go", **kwargs: Any): + """Run one task attempt: a fresh journal over the storage the attempts share.""" + agent = Agent(name="analyst", model=model, tools=AirflowTools(tool.tool())) + journal = DurableJournal(storage, clean_up_after_run=True) + with journal_scope(journal): + return run_durably_sync(agent, prompt, run_config=NO_TRACING, **kwargs), journal + + +class TestRunDurably: + def test_a_retry_replays_the_calls_the_failed_attempt_completed(self): + storage, tool = MemoryStorage(), _CountingTool() + with pytest.raises(RuntimeError, match="worker lost"): + _attempt(storage, _ScriptedModel(fail_on_answer=True), tool) + assert tool.calls == 1 + + retry_model = _ScriptedModel() + result, journal = _attempt(storage, retry_model, tool) + + assert result.final_output == "answer: k0=42" + # The first request and the tool call replay; only the request that answers is new. + assert retry_model.calls == 1 + assert tool.calls == 1 + assert journal.stats.replayed == {"model": 1, "tool": 1} + + def test_every_completed_call_replays_not_only_the_first(self): + """A replayed answer must turn into the same next request as the live one did.""" + storage, tool = MemoryStorage(), _CountingTool() + with pytest.raises(RuntimeError, match="worker lost"): + _attempt(storage, _ScriptedModel(lookups=3, fail_on_answer=True), tool) + + retry_model = _ScriptedModel(lookups=3) + _, journal = _attempt(storage, retry_model, tool) + + assert retry_model.calls == 1 + assert tool.calls == 3 + assert journal.stats.replayed == {"model": 3, "tool": 3} + + def test_the_sdks_generated_prompt_cache_key_does_not_stop_a_replay(self): + """The runner adds a key of its own to every request of a model that accepts one.""" + + class _CachingModel(_ScriptedModel): + def _supports_default_prompt_cache_key(self) -> bool: + return True + + storage, tool = MemoryStorage(), _CountingTool() + with pytest.raises(RuntimeError): + _attempt(storage, _CachingModel(fail_on_answer=True), tool) + + retry_model = _CachingModel() + _attempt(storage, retry_model, tool) + + assert retry_model.calls == 1 + + def test_a_replayed_model_call_reports_no_usage(self): + storage, tool = MemoryStorage(), _CountingTool() + with pytest.raises(RuntimeError): + _attempt(storage, _ScriptedModel(fail_on_answer=True, usage=7), tool) + + result, _ = _attempt(storage, _ScriptedModel(usage=7), tool) + + assert result.context_wrapper.usage.input_tokens == 7 + + def test_a_successful_run_deletes_what_it_recorded(self): + storage, tool = MemoryStorage(), _CountingTool() + + _attempt(storage, _ScriptedModel(), tool) + + assert storage.entries == {} + + def test_a_different_prompt_runs_every_call_again(self): + storage, tool = MemoryStorage(), _CountingTool() + with pytest.raises(RuntimeError): + _attempt(storage, _ScriptedModel(fail_on_answer=True), tool, prompt="first question") + + retry_model = _ScriptedModel() + _attempt(storage, retry_model, tool, prompt="second question") + + assert retry_model.calls == 2 + assert tool.calls == 2 + + def test_a_run_the_model_failed_asks_the_model_again(self): + """Running out of turns comes from the model's answers, which would replay into the same failure.""" + storage, tool = MemoryStorage(), _CountingTool() + with pytest.raises(MaxTurnsExceeded): + _attempt(storage, _ScriptedModel(lookups=3), tool, max_turns=2) + assert tool.calls == 2 + + retry_model = _ScriptedModel(lookups=1) + result, _ = _attempt(storage, retry_model, tool, max_turns=5) + + # The calls before the model's last request replay; that request is asked again, so + # the retry answers after the first lookup instead of replaying the second. + assert result.final_output == "answer: k0=42" + assert retry_model.calls == 1 + assert tool.calls == 2 + + def test_a_model_resolved_by_name_is_recorded(self): + storage, tool = MemoryStorage(), _CountingTool() + + def attempt(model: _ScriptedModel) -> None: + agent = Agent(name="analyst", model="scripted", tools=AirflowTools(tool.tool())) + config = RunConfig(tracing_disabled=True, model_provider=_Provider(model)) + with journal_scope(DurableJournal(storage, clean_up_after_run=True)): + run_durably_sync(agent, "go", run_config=config) + + with pytest.raises(RuntimeError): + attempt(_ScriptedModel(fail_on_answer=True)) + retry_model = _ScriptedModel() + attempt(retry_model) + + assert retry_model.calls == 1 + + def test_an_agent_used_as_a_tool_replays_as_one_call(self): + """Its own model calls run inside the outer tool call and must not take the outer run's positions.""" + storage, tool = MemoryStorage(), _CountingTool() + + def attempt(outer: _ScriptedModel, inner: _ScriptedModel) -> None: + specialist = Agent(name="specialist", model="inner", tools=AirflowTools(tool.tool())) + lead = Agent( + name="lead", + model=outer, + tools=[specialist.as_tool("lookup", "Ask the specialist.")], + ) + config = RunConfig(tracing_disabled=True, model_provider=_Provider(inner)) + with journal_scope(DurableJournal(storage, clean_up_after_run=True)): + run_durably_sync(lead, "go", run_config=config) + + # ``as_tool`` takes the question for the specialist as ``input``. + with pytest.raises(RuntimeError): + attempt(_ScriptedModel(lookups=2, fail_on_answer=True, argument="input"), _ScriptedModel()) + assert tool.calls == 2 + outer, inner = _ScriptedModel(lookups=2, argument="input"), _ScriptedModel() + attempt(outer, inner) + + # Both outer requests for the specialist and both of its answers replay. + assert outer.calls == 1 + assert inner.calls == 0 + + @pytest.mark.enable_redact + def test_a_secret_in_a_tool_result_is_not_recorded(self, registered_secret): + """Any function tool, not only the Airflow tools, whose results are masked already.""" + + @function_tool + def lookup(key: str) -> str: + """Look a key up.""" + return f"{key}={registered_secret}" + + agent = Agent(name="analyst", model=_ScriptedModel(fail_on_answer=True), tools=[lookup]) + storage = MemoryStorage() + + with journal_scope(DurableJournal(storage, clean_up_after_run=True)), pytest.raises(RuntimeError): + run_durably_sync(agent, "go", run_config=NO_TRACING) + + (recorded,) = [step for step in storage.steps() if step["name"] == "tool:lookup"] + assert recorded["payload"] == "k0=***" + assert registered_secret not in json.dumps(storage.entries) + + @pytest.mark.parametrize( + "result", + [ + pytest.param(_Report(rows=3), id="pydantic-model"), + pytest.param(("orders", 3), id="tuple"), + pytest.param(datetime.datetime(2026, 10, 14, 9, 30), id="datetime"), + pytest.param({"rows": 3}, id="dict"), + pytest.param(ToolOutputText(text="3 rows"), id="structured"), + ], + ) + def test_a_tool_result_replays_as_the_model_saw_it(self, result): + """The SDK sends most results as str(); a replay must turn into the same next request.""" + + @function_tool + def lookup(key: str) -> Any: + """Look a key up.""" + return result + + storage = MemoryStorage() + + def attempt(model: _ScriptedModel) -> DurableJournal: + agent = Agent(name="analyst", model=model, tools=[lookup]) + journal = DurableJournal(storage, clean_up_after_run=True) + with journal_scope(journal): + run_durably_sync(agent, "go", run_config=NO_TRACING) + return journal + + with pytest.raises(RuntimeError): + attempt(_ScriptedModel(lookups=2, fail_on_answer=True)) + journal = attempt(_ScriptedModel(lookups=2)) + + assert journal.stats.replayed == {"model": 2, "tool": 2} + + @pytest.mark.parametrize( + "change", + [ + pytest.param({"instructions": "Answer in French."}, id="instructions"), + pytest.param({"model_settings": ModelSettings(temperature=0.1)}, id="settings"), + pytest.param({"output_type": _Report}, id="output-type"), + pytest.param({"handoffs": [Agent(name="auditor")]}, id="handoffs"), + pytest.param({"tools": []}, id="tools"), + ], + ) + def test_a_request_that_changed_between_attempts_runs_live(self, change): + storage, tool = MemoryStorage(), _CountingTool() + with pytest.raises(RuntimeError): + _attempt(storage, _ScriptedModel(fail_on_answer=True), tool) + + retry_model = _ScriptedModel() + agent = Agent(name="analyst", model=retry_model, tools=AirflowTools(tool.tool())).clone(**change) + # Some changes make the scripted answers invalid for the agent; only the first request matters. + with journal_scope(DurableJournal(storage, clean_up_after_run=True)), contextlib.suppress(Exception): + run_durably_sync(agent, "go", run_config=NO_TRACING) + + # The first request reached the model, rather than replaying the failed try's answer. + assert retry_model.tool_results_seen[0] == 0 + + def test_a_request_the_sdk_retried_keeps_its_position(self): + """The SDK's model retry makes the same call again; the next try must still line up.""" + settings = ModelSettings( + retry=ModelRetrySettings( + max_retries=1, + backoff=ModelRetryBackoffSettings(initial_delay=0, max_delay=0, jitter=False), + policy=lambda context: True, + ) + ) + storage, tool = MemoryStorage(), _CountingTool() + + def attempt(model: _ScriptedModel) -> DurableJournal: + agent = Agent( + name="analyst", model=model, tools=AirflowTools(tool.tool()), model_settings=settings + ) + journal = DurableJournal(storage, clean_up_after_run=True) + with journal_scope(journal): + run_durably_sync(agent, "go", run_config=NO_TRACING) + return journal + + with pytest.raises(RuntimeError, match="worker lost"): + attempt(_ScriptedModel(fail_first_request=True, fail_on_answer=True)) + journal = attempt(_ScriptedModel()) + + assert journal.stats.replayed == {"model": 1, "tool": 1} + + def test_a_tool_whose_approval_depends_on_its_arguments_is_left_to_the_sdk(self): + async def policy(context, arguments, call_id) -> bool: + return arguments.get("amount", 0) > 100 + + @function_tool(needs_approval=policy) + def pay(amount: int = 1000) -> str: + """Pay an amount.""" + return "paid" + + assert _durable_tool(pay, DurableJournal(MemoryStorage()).start_run()) is pay + + @pytest.mark.parametrize("option", ["previous_response_id", "conversation_id"]) + def test_refuses_a_conversation_kept_on_openais_servers(self, option): + with pytest.raises(ValueError, match=option): + _attempt(MemoryStorage(), _ScriptedModel(), _CountingTool(), **{option: "resp_1"}) + + def test_outside_a_task_it_runs_the_agent_as_runner_run_does(self): + model, tool = _ScriptedModel(), _CountingTool() + agent = Agent(name="analyst", model=model, tools=AirflowTools(tool.tool())) + + result = run_durably_sync(agent, "go", run_config=NO_TRACING) + + assert result.final_output == "answer: k0=42" + assert model.calls == 2
