This is an automated email from the ASF dual-hosted git repository.
Lee-W pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f5172d84929 Make OpenAI Responses token ceilings templated and surface
truncation (#72150)
f5172d84929 is described below
commit f5172d84929e805d24362b955b89da9b2ebbc887
Author: Wei Lee <[email protected]>
AuthorDate: Tue Sep 15 23:36:07 2026 +0800
Make OpenAI Responses token ceilings templated and surface truncation
(#72150)
---
providers/openai/docs/operators/openai.rst | 30 +-
.../airflow/providers/openai/operators/openai.py | 146 ++++++++-
.../openai/tests/system/openai/example_openai.py | 12 +
.../tests/unit/openai/operators/test_openai.py | 360 ++++++++++++++++++++-
4 files changed, 540 insertions(+), 8 deletions(-)
diff --git a/providers/openai/docs/operators/openai.rst
b/providers/openai/docs/operators/openai.rst
index b6bd4ab5b3f..06689ac041d 100644
--- a/providers/openai/docs/operators/openai.rst
+++ b/providers/openai/docs/operators/openai.rst
@@ -53,6 +53,27 @@ The OpenAIResponseOperator requires the ``input_text``
prompt. Use the ``conn_id
specify the OpenAI connection to use, and ``response_kwargs`` to pass through
options such as
``tools``, ``conversation`` or ``previous_response_id``.
+Use ``max_output_tokens`` and ``max_tool_calls`` to cap generation per run --
both are templated,
+so a ceiling can vary by environment or Dag run without hardcoding it.
``max_output_tokens`` caps
+the number of tokens generated; ``max_tool_calls`` caps the number of built-in
tool calls the model
+may make. Both limits are enforced by the OpenAI API itself; OpenAI exposes no
monetary cost limit
+on the Responses API, so this operator has no cost cap. For a monetary limit,
use
+:doc:`apache-airflow-providers-common-ai:index` instead. Hitting
``max_output_tokens`` does not
+fail the request: the response comes back with ``status="incomplete"``, so
``return_value`` will
+not raise -- but it is not guaranteed to be truncated text either. A reasoning
model can spend
+the entire ceiling on reasoning tokens and return an empty ``output_text``, in
which case
+``return_value`` is an empty string. Hitting ``max_tool_calls`` is different:
the OpenAI API
+silently drops any tool calls beyond the ceiling without changing ``status``
or setting
+``incomplete_details`` -- there is no log warning and no signal in
``return_value``, so a run
+truncated by ``max_tool_calls`` looks identical to a clean run.
+
+A rendered ``max_output_tokens`` or ``max_tool_calls`` that is blank or
whitespace-only -- for
+example ``max_output_tokens="{{ params.tokens | default('', true) }}"`` when
``params.tokens`` is
+unset -- is treated as "no ceiling for this run" rather than raising. This
only applies when the same run does
+not also set the corresponding key in ``response_kwargs``: the
mutually-exclusive-with-``response_kwargs``
+check happens when the operator is constructed and fires regardless of what
the template later
+renders to.
+
.. exampleinclude:: /../../openai/tests/system/openai/example_openai.py
:language: python
:start-after: [START howto_operator_openai_response]
@@ -103,8 +124,13 @@ know about yet. Options worth knowing about:
by ``execute``. Use ``OpenAIHook`` directly to access it.
- ``metadata``: a mapping of key-value pairs attached to the response for your
own bookkeeping.
- ``max_output_tokens``: an upper bound on the number of tokens the model can
generate, including
- reasoning tokens as well as visible output tokens.
-- ``max_tool_calls``: an upper bound on the number of built-in tool calls the
model can make.
+ reasoning tokens as well as visible output tokens. Prefer the operator's own
``max_output_tokens``
+ parameter (see above) instead of setting this key here: the operator
parameter is templated, this
+ ``response_kwargs`` key is not, and setting the same ceiling in both places
raises when the
+ operator is constructed.
+- ``max_tool_calls``: an upper bound on the number of built-in tool calls the
model can make. Same
+ trade-off as ``max_output_tokens`` above: prefer the operator's own,
templated ``max_tool_calls``
+ parameter instead of this non-templated ``response_kwargs`` key.
.. note::
diff --git a/providers/openai/src/airflow/providers/openai/operators/openai.py
b/providers/openai/src/airflow/providers/openai/operators/openai.py
index 1f12593e23b..0dad7792f18 100644
--- a/providers/openai/src/airflow/providers/openai/operators/openai.py
+++ b/providers/openai/src/airflow/providers/openai/operators/openai.py
@@ -19,7 +19,7 @@ from __future__ import annotations
from collections.abc import Sequence
from functools import cached_property
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, ClassVar, Literal
from airflow.providers.common.compat.sdk import BaseOperator, conf
from airflow.providers.openai.exceptions import OpenAIBatchJobException
@@ -87,6 +87,18 @@ class OpenAIResponseOperator(BaseOperator):
``previous_response_id`` chaining, ``background=True`` responses, or
access to the full
structured response, use
:class:`~airflow.providers.openai.hooks.openai.OpenAIHook` directly.
+ ``max_output_tokens`` caps the number of tokens generated for the
response; ``max_tool_calls``
+ caps the number of built-in tool calls the model may make. Both limits are
enforced by the
+ OpenAI API itself -- OpenAI exposes no monetary cost limit on the
Responses API, so this
+ operator has no cost cap. For a monetary limit, use
+ :doc:`apache-airflow-providers-common-ai:index` instead. When
``max_output_tokens`` is hit, the
+ request does not fail: the response comes back with
``status="incomplete"`` -- but
+ ``output_text`` is not guaranteed to contain any content, since a
reasoning model can spend
+ the entire ceiling on reasoning tokens without producing visible output.
Hitting
+ ``max_tool_calls`` is different: the OpenAI SDK documents it as silently
dropping further
+ tool calls, with no ``status`` change and no ``incomplete_details`` -- a
run truncated this
+ way looks identical to a clean one in both the logs and ``return_value``.
+
:param conn_id: The OpenAI connection ID to use.
:param input_text: The input prompt for the model. This can be a string or
a structured list of
input items.
@@ -97,7 +109,25 @@ class OpenAIResponseOperator(BaseOperator):
completes, so this operator logs a warning and the returned output
text may be empty, while
``stream=True`` returns an object without ``status`` or
``output_text``, so the task raises
``AttributeError``. See :ref:`howto/operator:OpenAIResponseOperator`
for these and other
- options this operator can pass through, such as ``truncation`` and
``max_output_tokens``.
+ options this operator can pass through, such as ``truncation`` and
``metadata``.
+ :param max_output_tokens: Optional upper bound on the number of tokens
generated for the
+ response. Templated, so it renders to a string; accepts an ``int`` or
a string containing one.
+ Must be a positive integer -- an invalid value raises instead of
silently disabling the
+ ceiling. A literal ``bool``, ``float``, or ``int`` value is validated
when the operator is
+ constructed; any other non-string literal (for example ``Decimal`` or
``Fraction``) is
+ coerced -- and rejected if invalid -- only when the task executes. A
string value --
+ whether a template or a plain literal string -- is also validated when
the task executes,
+ after templating has resolved it. A blank or whitespace-only rendered
value (for example
+ ``{{ params.tokens | default('', true) }}`` rendering to ``''``) is
treated as unset,
+ disabling the ceiling; the literal strings ``"None"``, ``"none"`` and
``"null"`` are **not**
+ treated as blank and still raise. A value that was supplied but
resolves to ``None`` (for
+ example an unresolved ``XComArg``, or a Jinja-native-rendered null)
also raises -- it is not
+ treated as unset. Mutually exclusive with ``max_output_tokens`` in
``response_kwargs`` --
+ this is checked when the operator is constructed, regardless of what
the templated value
+ later renders to.
+ :param max_tool_calls: Optional upper bound on the number of built-in tool
calls the model may
+ make while generating the response. Same templating, type, validation,
blank-as-unset, and
+ mutual-exclusion rules as ``max_output_tokens``.
.. seealso::
For more information on how to use this operator, take a look at the
guide:
@@ -106,7 +136,9 @@ class OpenAIResponseOperator(BaseOperator):
https://platform.openai.com/docs/api-reference/responses/create
"""
- template_fields: Sequence[str] = ("input_text",)
+ template_fields: Sequence[str] = ("input_text", "max_output_tokens",
"max_tool_calls")
+
+ _TOKEN_CEILING_PARAM_NAMES: ClassVar[tuple[str, ...]] =
("max_output_tokens", "max_tool_calls")
def __init__(
self,
@@ -114,6 +146,9 @@ class OpenAIResponseOperator(BaseOperator):
input_text: str | list[Any],
model: str = "gpt-4o-mini",
response_kwargs: dict | None = None,
+ *,
+ max_output_tokens: int | str | None = None,
+ max_tool_calls: int | str | None = None,
**kwargs: Any,
):
super().__init__(**kwargs)
@@ -121,15 +156,116 @@ class OpenAIResponseOperator(BaseOperator):
self.input_text = input_text
self.model = model
self.response_kwargs = response_kwargs or {}
+ self.max_output_tokens = max_output_tokens
+ self.max_tool_calls = max_tool_calls
+ self._supplied_ceilings: frozenset[str] = frozenset(
+ name for name in self._TOKEN_CEILING_PARAM_NAMES if getattr(self,
name) is not None
+ )
+ self._validate_no_response_kwargs_conflict()
+ self._validate_literal_ceiling_values()
+
+ def _validate_no_response_kwargs_conflict(self) -> None:
+ """Reject a ceiling set both as an operator argument and in
``response_kwargs``."""
+ for param_name in self._TOKEN_CEILING_PARAM_NAMES:
+ value = getattr(self, param_name)
+ if value is not None and param_name in self.response_kwargs:
+ raise ValueError(
+ f"Task {self.task_id!r}: {param_name!r} was set both as an
operator argument "
+ "and in 'response_kwargs'; set it in only one place."
+ )
+
+ def _validate_literal_ceiling_values(self) -> None:
+ """
+ Eagerly validate a ceiling value that is already a final literal, not
a template.
+
+ Only ``bool``, ``float``, and ``int`` are recognized as literals here
-- these are the raw
+ values passed at construction, before any templating runs, so an
invalid one is rejected
+ when the operator is constructed instead of surfacing only when the
task runs. Anything
+ else (``str`` templates awaiting ``render_template_fields()``, or
template values such as
+ ``XComArg`` that resolve later -- including a ``bool``, ``float``, or
``int`` produced by
+ Jinja's native rendering with ``render_template_as_native_obj=True``)
must wait for
+ ``_build_response_kwargs()`` at ``execute()`` time.
+ """
+ for param_name in self._TOKEN_CEILING_PARAM_NAMES:
+ value = getattr(self, param_name)
+ if value is not None and isinstance(value, (bool, float, int)):
+ self._coerce_token_ceiling(param_name, value)
@cached_property
def hook(self) -> OpenAIHook:
"""Return an instance of the OpenAIHook."""
return OpenAIHook(conn_id=self.conn_id)
+ @staticmethod
+ def _coerce_token_ceiling(param_name: str, value: int | float | str) ->
int:
+ """Coerce a templated token-ceiling argument to a positive int, or
raise ``ValueError``."""
+ # bool is an int subclass (isinstance(True, int) is True) and must be
rejected before the
+ # allowlist check below. Only int and str are accepted as real values
to coerce; anything
+ # else -- float, Decimal, Fraction, or any other numeric type -- is
rejected here instead
+ # of being handed to int(), since int() silently truncates those (e.g.
int(10.5) == 10,
+ # int(Decimal("10.5")) == 10) rather than raising. Such values can
reach here as real Python
+ # objects, not just strings, when a Dag uses
render_template_as_native_obj=True.
+ if isinstance(value, bool):
+ raise ValueError(f"{param_name!r} must be an integer, got
{value!r}.")
+ if not isinstance(value, (int, str)):
+ raise ValueError(f"{param_name!r} must be an integer, got
{value!r}.")
+ try:
+ coerced = int(value)
+ except (TypeError, ValueError):
+ raise ValueError(f"{param_name!r} must be an integer, got
{value!r}.")
+ if coerced <= 0:
+ raise ValueError(f"{param_name!r} must be a positive integer, got
{coerced}.")
+ return coerced
+
+ def _build_response_kwargs(self) -> dict[str, Any]:
+ """Merge the token-ceiling arguments into ``response_kwargs``,
skipping unset ceilings."""
+ response_kwargs = dict(self.response_kwargs)
+ for param_name in self._TOKEN_CEILING_PARAM_NAMES:
+ value = getattr(self, param_name)
+ if param_name not in self._supplied_ceilings:
+ continue
+ # Blank means unset; the response_kwargs conflict was already
rejected in __init__.
+ if isinstance(value, str) and value.strip() == "":
+ continue
+ if value is None:
+ raise ValueError(
+ f"{param_name!r} was supplied but resolved to None (e.g.
an unresolved "
+ "XComArg, or a Jinja-native-rendered null); pass a
positive integer, or "
+ "leave the argument unset entirely to disable the ceiling."
+ )
+ response_kwargs[param_name] =
self._coerce_token_ceiling(param_name, value)
+ return response_kwargs
+
def execute(self, context: Context) -> str:
- response = self.hook.create_response(input=self.input_text,
model=self.model, **self.response_kwargs)
- if response.status != "completed":
+ response = self.hook.create_response(
+ input=self.input_text, model=self.model,
**self._build_response_kwargs()
+ )
+ if response.status == "incomplete":
+ reason = response.incomplete_details.reason if
response.incomplete_details else None
+ if reason and response.output_text:
+ # Any reason -- including max_output_tokens -- can fire before
any output text
+ # is produced (e.g. a reasoning model spends the whole ceiling
on reasoning
+ # tokens), so whether truncated content actually exists is
decided by looking at
+ # output_text itself, not by the reason string.
+ self.log.warning(
+ "Response %s is incomplete (incomplete_details.reason=%s);
the returned output "
+ "text is truncated, not empty.",
+ response.id,
+ reason,
+ )
+ elif reason:
+ self.log.warning(
+ "Response %s is incomplete (incomplete_details.reason=%s);
the returned output "
+ "text may be empty.",
+ response.id,
+ reason,
+ )
+ else:
+ self.log.warning(
+ "Response %s is incomplete; the returned output text may
be truncated or empty.",
+ response.id,
+ )
+ elif response.status != "completed":
self.log.warning(
"Response %s ended with status %s; the returned output text
may be empty.",
response.id,
diff --git a/providers/openai/tests/system/openai/example_openai.py
b/providers/openai/tests/system/openai/example_openai.py
index e03bc639745..e14c89e6db0 100644
--- a/providers/openai/tests/system/openai/example_openai.py
+++ b/providers/openai/tests/system/openai/example_openai.py
@@ -107,6 +107,18 @@ def example_openai_dag():
input_text="Write a haiku about data pipelines.",
response_kwargs={"instructions": "You are a helpful assistant."},
)
+
+ # ``max_output_tokens`` is templated, so a ceiling can vary by Dag run
without hardcoding it.
+ # This Dag does not declare a ``tokens`` param, so ``params.tokens`` is
undefined at render
+ # time. ``| default('', true)`` renders undefined *and* None values to
``''`` (treated as "no
+ # ceiling"); the more common ``{{ params.tokens or '' }}`` idiom would
instead raise
+ # ``UndefinedError`` under Airflow's default ``StrictUndefined`` template
behavior.
+ OpenAIResponseOperator(
+ task_id="openai_response_with_token_ceiling",
+ conn_id="openai_default",
+ input_text="Write a haiku about data pipelines.",
+ max_output_tokens="{{ params.tokens | default('', true) }}",
+ )
# [END howto_operator_openai_response]
create_embeddings_using_hook()
diff --git a/providers/openai/tests/unit/openai/operators/test_openai.py
b/providers/openai/tests/unit/openai/operators/test_openai.py
index 0ec78bb182c..d4ba4403840 100644
--- a/providers/openai/tests/unit/openai/operators/test_openai.py
+++ b/providers/openai/tests/unit/openai/operators/test_openai.py
@@ -16,13 +16,17 @@
# under the License.
from __future__ import annotations
+from decimal import Decimal
+from fractions import Fraction
from unittest.mock import Mock
+import jinja2
import pytest
from openai.types.batch import Batch
from openai.types.responses import Response
+from openai.types.responses.response import IncompleteDetails
-from airflow.providers.common.compat.sdk import Context, TaskDeferred
+from airflow.providers.common.compat.sdk import DAG, BaseOperator, Context,
TaskDeferred, XComArg
from airflow.providers.openai.exceptions import OpenAIBatchJobException,
OpenAITriggerEventError
from airflow.providers.openai.hooks.openai import OpenAIHook
from airflow.providers.openai.operators.openai import (
@@ -105,6 +109,360 @@ def test_openai_response_operator_execute():
)
+def _build_completed_response(**overrides):
+ defaults = {"output_text": "haiku text", "id": "resp_123", "status":
"completed"}
+ return Mock(spec=Response, **{**defaults, **overrides})
+
+
+class TestOpenAIResponseOperatorTokenCeilings:
+ @pytest.mark.parametrize(
+ ("kwargs", "expected_extra"),
+ [
+ pytest.param({"max_output_tokens": 100}, {"max_output_tokens":
100}, id="max_output_tokens-int"),
+ pytest.param({"max_tool_calls": 5}, {"max_tool_calls": 5},
id="max_tool_calls-int"),
+ pytest.param(
+ {"max_output_tokens": "100"}, {"max_output_tokens": 100},
id="max_output_tokens-numeric-str"
+ ),
+ pytest.param({"max_tool_calls": "5"}, {"max_tool_calls": 5},
id="max_tool_calls-numeric-str"),
+ pytest.param(
+ {"max_output_tokens": 100, "max_tool_calls": 5},
+ {"max_output_tokens": 100, "max_tool_calls": 5},
+ id="both",
+ ),
+ ],
+ )
+ def test_valid_ceiling_forwarded_as_int(self, kwargs, expected_extra):
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID, conn_id=CONN_ID, input_text="Write a haiku.",
**kwargs
+ )
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_hook_instance.create_response.return_value =
_build_completed_response()
+ operator.hook = mock_hook_instance
+
+ operator.execute(Context())
+
+ mock_hook_instance.create_response.assert_called_once_with(
+ input="Write a haiku.", model="gpt-4o-mini", **expected_extra
+ )
+
+ @pytest.mark.parametrize(
+ "invalid_value",
+ [
+ pytest.param("not-a-number", id="non-integer-string"),
+ pytest.param("-5", id="negative-string"),
+ pytest.param("None", id="literal-none-string"),
+ pytest.param(Decimal("10.5"), id="decimal"),
+ pytest.param(Fraction(21, 2), id="fraction"),
+ ],
+ )
+ @pytest.mark.parametrize("param_name", ["max_output_tokens",
"max_tool_calls"])
+ def test_invalid_ceiling_raises_before_request(self, param_name,
invalid_value):
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID, conn_id=CONN_ID, input_text="Write a haiku.",
**{param_name: invalid_value}
+ )
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ operator.hook = mock_hook_instance
+
+ with pytest.raises(ValueError, match=param_name):
+ operator.execute(Context())
+
+ mock_hook_instance.create_response.assert_not_called()
+
+ @pytest.mark.parametrize(
+ "invalid_value",
+ [
+ pytest.param(0, id="zero"),
+ pytest.param(-1, id="negative"),
+ pytest.param(10.5, id="float"),
+ pytest.param(True, id="bool-true"),
+ pytest.param(False, id="bool-false"),
+ ],
+ )
+ @pytest.mark.parametrize("param_name", ["max_output_tokens",
"max_tool_calls"])
+ def test_non_string_invalid_ceiling_raises_at_construction(self,
param_name, invalid_value):
+ with pytest.raises(ValueError, match=param_name):
+ OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ **{param_name: invalid_value},
+ )
+
+ @pytest.mark.parametrize(
+ "operator_value",
+ [
+ pytest.param(100, id="int"),
+ pytest.param("", id="blank"),
+ ],
+ )
+ @pytest.mark.parametrize("param_name", ["max_output_tokens",
"max_tool_calls"])
+ def test_ceiling_conflicting_with_response_kwargs_raises(self, param_name,
operator_value):
+ # A blank operator_value must still conflict with response_kwargs;
that's checked at
+ # construction time, before rendering. pytest.raises() itself fails
with "DID NOT RAISE"
+ # if construction succeeded, so there's no operator instance
afterwards to assert
+ # anything further against.
+ with pytest.raises(ValueError, match=param_name):
+ OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ response_kwargs={param_name: 50},
+ **{param_name: operator_value},
+ )
+
+ def test_conflict_error_precedes_literal_type_error(self):
+ # 0 is both invalid on its own (not positive) and conflicting with
response_kwargs; the
+ # conflict message must win, since fixing the duplicate is the
actionable first step.
+ with pytest.raises(ValueError, match="was set both as an operator
argument"):
+ OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ response_kwargs={"max_output_tokens": 50},
+ max_output_tokens=0,
+ )
+
+ def test_xcom_arg_ceiling_does_not_fail_on_construction(self):
+ with DAG("test_dag", schedule=None) as dag:
+ upstream = BaseOperator(task_id="upstream")
+
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ max_output_tokens=upstream.output,
+ dag=dag,
+ )
+
+ assert isinstance(operator.max_output_tokens, XComArg)
+
+ @pytest.mark.parametrize(
+ "blank_value", [pytest.param("", id="empty"), pytest.param(" ",
id="whitespace")]
+ )
+ @pytest.mark.parametrize("param_name", ["max_output_tokens",
"max_tool_calls"])
+ def test_blank_ceiling_is_treated_as_unset(self, param_name, blank_value):
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID, conn_id=CONN_ID, input_text="Write a haiku.",
**{param_name: blank_value}
+ )
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_hook_instance.create_response.return_value =
_build_completed_response()
+ operator.hook = mock_hook_instance
+
+ operator.execute(Context())
+
+ call_kwargs = mock_hook_instance.create_response.call_args.kwargs
+ assert param_name not in call_kwargs
+
+ def test_max_output_tokens_and_max_tool_calls_are_templated(self):
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ max_output_tokens="{{ params.tokens }}",
+ max_tool_calls="{{ params.calls }}",
+ )
+
+ operator.render_template_fields(Context(params={"tokens": 100,
"calls": 5}))
+
+ assert operator.max_output_tokens == "100"
+ assert operator.max_tool_calls == "5"
+ assert "max_output_tokens" in operator.template_fields
+ assert "max_tool_calls" in operator.template_fields
+
+ # The rendered strings must still make it to the SDK as real ints, not
left as strings.
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_hook_instance.create_response.return_value =
_build_completed_response()
+ operator.hook = mock_hook_instance
+
+ operator.execute(Context())
+
+ call_kwargs = mock_hook_instance.create_response.call_args.kwargs
+ for key, expected in (
+ ("max_output_tokens", 100),
+ ("max_tool_calls", 5),
+ ):
+ assert call_kwargs[key] == expected
+ assert isinstance(call_kwargs[key], int)
+ assert not isinstance(call_kwargs[key], bool)
+
+ @pytest.mark.parametrize("param_name", ["max_output_tokens",
"max_tool_calls"])
+ def test_supplied_xcom_arg_ceiling_resolving_to_none_raises(self,
param_name):
+ # An XComArg bound at construction time (e.g. upstream.output) is not
"unset" -- if
+ # rendering it later resolves to None (no XCom was ever pushed), that
must raise instead
+ # of silently disabling the ceiling.
+ with DAG("test_dag", schedule=None) as dag:
+ upstream = BaseOperator(task_id="upstream")
+
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ dag=dag,
+ **{param_name: upstream.output},
+ )
+
+ mock_ti = Mock()
+ mock_ti.xcom_pull.return_value = None
+ # Airflow 2's XComArg.resolve() reads context["expanded_ti_count"]
unconditionally, so the
+ # key has to be present for this to render under the compatibility
test suite.
+ operator.render_template_fields(Context(ti=mock_ti,
expanded_ti_count=None))
+
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ operator.hook = mock_hook_instance
+
+ with pytest.raises(ValueError, match=param_name):
+ operator.execute(Context())
+
+ mock_hook_instance.create_response.assert_not_called()
+
+ @pytest.mark.parametrize("param_name", ["max_output_tokens",
"max_tool_calls"])
+ def test_native_rendered_none_ceiling_raises(self, param_name):
+ # render_template_as_native_obj=True can render a Jinja template
straight to a real
+ # None -- that is also "supplied but resolved to None", not "unset".
+ param_key = "tokens" if param_name == "max_output_tokens" else "calls"
+ with DAG("test_dag", schedule=None,
render_template_as_native_obj=True):
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ **{param_name: f"{{{{ params.{param_key} }}}}"},
+ )
+
+ operator.render_template_fields(Context(params={param_key: None}))
+
+ assert getattr(operator, param_name) is None
+
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ operator.hook = mock_hook_instance
+
+ with pytest.raises(ValueError, match=param_name):
+ operator.execute(Context())
+
+ mock_hook_instance.create_response.assert_not_called()
+
+ def test_or_fallback_idiom_raises_under_strict_undefined_dag_binding(self):
+ with DAG("test_dag", schedule=None) as dag:
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ max_output_tokens="{{ params.tokens or '' }}",
+ dag=dag,
+ )
+
+ with pytest.raises(jinja2.UndefinedError):
+ operator.render_template_fields(Context(params={}))
+
+ def
test_default_filter_idiom_renders_blank_under_strict_undefined_dag_binding(self):
+ with DAG("test_dag", schedule=None) as dag:
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ max_output_tokens="{{ params.tokens | default('', true) }}",
+ dag=dag,
+ )
+
+ assert operator.get_template_env().undefined is jinja2.StrictUndefined
+
+ operator.render_template_fields(Context(params={}))
+ assert operator.max_output_tokens == ""
+
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_hook_instance.create_response.return_value =
_build_completed_response()
+ operator.hook = mock_hook_instance
+
+ operator.execute(Context())
+
+ call_kwargs = mock_hook_instance.create_response.call_args.kwargs
+ assert "max_output_tokens" not in call_kwargs
+
+ @pytest.mark.parametrize(
+ ("reason", "output_text", "expected_fragment"),
+ [
+ pytest.param(
+ "max_output_tokens",
+ "Truncated hai",
+ "the returned output text is truncated, not empty.",
+ id="max_output_tokens-nonempty-output",
+ ),
+ pytest.param(
+ "content_filter",
+ "",
+ "the returned output text may be empty.",
+ id="content_filter-empty-output",
+ ),
+ pytest.param(
+ # A reasoning model can spend the entire max_output_tokens
ceiling on reasoning
+ # tokens and produce no visible output text -- the wording
must be decided by
+ # output_text, not by reason, even when reason is
"max_output_tokens".
+ "max_output_tokens",
+ "",
+ "the returned output text may be empty.",
+ id="max_output_tokens-empty-output",
+ ),
+ ],
+ )
+ def test_incomplete_details_reason_is_logged(self, caplog, reason,
output_text, expected_fragment):
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID, conn_id=CONN_ID, input_text="Write a haiku.",
max_output_tokens=10
+ )
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_hook_instance.create_response.return_value =
_build_completed_response(
+ status="incomplete",
+ incomplete_details=IncompleteDetails(reason=reason),
+ output_text=output_text,
+ )
+ operator.hook = mock_hook_instance
+
+ with caplog.at_level("WARNING"):
+ result = operator.execute(Context())
+
+ assert result == output_text
+ assert any(
+ f"incomplete_details.reason={reason}" in message and
expected_fragment in message
+ for message in caplog.messages
+ )
+ # The wording is decided by output_text, not by reason: a truthy
output_text always gets
+ # the "truncated, not empty" message and an empty one always gets "may
be empty",
+ # regardless of what reason is.
+ if output_text:
+ assert not any("may be empty" in message for message in
caplog.messages)
+ else:
+ assert not any("truncated, not empty" in message for message in
caplog.messages)
+
+ def test_incomplete_without_details_uses_truncated_or_empty_message(self,
caplog):
+ operator = OpenAIResponseOperator(task_id=TASK_ID, conn_id=CONN_ID,
input_text="Write a haiku.")
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_hook_instance.create_response.return_value =
_build_completed_response(
+ status="incomplete", incomplete_details=None, output_text=""
+ )
+ operator.hook = mock_hook_instance
+
+ with caplog.at_level("WARNING"):
+ result = operator.execute(Context())
+
+ assert result == ""
+ assert any("may be truncated or empty" in message for message in
caplog.messages)
+ assert not any("truncated, not empty" in message for message in
caplog.messages)
+ assert not any("may be empty" in message for message in
caplog.messages)
+
+ def
test_non_completed_without_incomplete_details_keeps_may_be_empty_message(self,
caplog):
+ operator = OpenAIResponseOperator(task_id=TASK_ID, conn_id=CONN_ID,
input_text="Write a haiku.")
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_hook_instance.create_response.return_value =
_build_completed_response(
+ status="failed", output_text=""
+ )
+ operator.hook = mock_hook_instance
+
+ with caplog.at_level("WARNING"):
+ operator.execute(Context())
+
+ assert any(
+ "ended with status failed" in message and "may be empty" in
message for message in caplog.messages
+ )
+
+
@pytest.mark.parametrize("wait_for_completion", [True, False])
def test_openai_trigger_batch_operator_not_deferred(mock_batch,
wait_for_completion):
operator = OpenAITriggerBatchOperator(
