kaxil commented on code in PR #72150:
URL: https://github.com/apache/airflow/pull/72150#discussion_r4007655684
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -87,12 +87,38 @@ 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.
:param model: The OpenAI model to use.
:param response_kwargs: Additional keyword arguments to pass to the OpenAI
``create_response``
method (for example ``instructions``, ``tools``, ``conversation`` or
``previous_response_id``).
+ :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 (non-string) value is validated at task definition
(Dag-parse) time; a
Review Comment:
This promises parse-time validation more broadly than the check delivers.
The eager gate on line 177 is `isinstance(value, (bool, float, int))`, so
`Decimal("10.5")` and `Fraction(21, 2)` are non-string literals that construct
cleanly and only raise at `execute()`;
`test_invalid_ceiling_raises_before_request` demonstrates exactly that by
constructing with them and then calling `execute()`. Saying "validated when the
operator is constructed" instead would be true for both, and would also cover
`.partial()`/`.expand()`, where `unmap()` runs `__init__` on the worker rather
than at Dag-parse time.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -101,30 +127,123 @@ 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,
conn_id: str,
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)
self.conn_id = conn_id
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._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
-- they are already
+ final at construction and never arrive via Jinja rendering, so an
invalid one is rejected
+ at Dag-parse time 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) 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)
+ # Blank means unset; the response_kwargs conflict was already
rejected in __init__.
+ if value is None or (isinstance(value, str) and value.strip() ==
""):
Review Comment:
A ceiling that was supplied but resolves to `None` lands here as "unset", so
the request goes out with no cap at all and nothing logs that one was dropped.
Two paths reach it: `max_output_tokens=budget_task.output` where the upstream
pushed nothing, since `PlainXComArg.resolve` returns `None` for a missing
`XCOM_RETURN_KEY` (`xcom_arg.py:371-372`) and that is the input shape
`test_xcom_arg_ceiling_does_not_fail_on_construction` was added for; and
`render_template_as_native_obj=True` with `{{ params.tokens }}` where `tokens`
defaults to `None`, which renders a real `None` rather than the string `'None'`
(`dag.py:852-856` selects `NativeEnvironment`; I measured it). The same Dag
raises without native rendering, which is what line 115 promises, so flipping
that Dag flag turns a hard failure into a silent uncapped run. Recording which
params were non-`None` in `__init__` and rejecting a `None` resolution only for
those would keep blank-means-unset working.
##########
providers/openai/tests/unit/openai/operators/test_openai.py:
##########
@@ -105,6 +108,268 @@ 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 }}",
Review Comment:
The `StrictUndefined`-safe spelling from the last round has no test, and a
case in this shape could not give it one: `Templater.get_template_env` returns
the Dag's `StrictUndefined` environment only when the operator has a Dag and
otherwise hands back a permissive `SandboxedEnvironment`
(`templater.py:79-81`), so the rejected `{{ params.tokens or '' }}` also
renders `''` here and would pass. One case built inside `with DAG(...)` and
rendered with `Context(params={})` would actually pin the idiom
`example_openai.py:120` now teaches.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -101,30 +127,123 @@ 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,
conn_id: str,
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)
self.conn_id = conn_id
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._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
-- they are already
+ final at construction and never arrive via Jinja rendering, so an
invalid one is rejected
Review Comment:
This says `bool`/`float`/`int` never arrive via Jinja rendering, but the
comment on lines 192-193 says the opposite and is the one that holds: with
`render_template_as_native_obj=True` the Dag renders through a
`NativeEnvironment` (`dag.py:852-856`), so `{{ params.tokens }}` over an int
param lands here as a real `int`. Worth correcting because the false clause
reads as an argument for deleting the `not isinstance(value, (int, str))` guard
below, and that guard is what keeps `int(Decimal("10.5")) == 10` off the wire.
--
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]
