kaxil commented on code in PR #72150:
URL: https://github.com/apache/airflow/pull/72150#discussion_r3993071204
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -101,30 +123,117 @@ 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 cannot possibly still be an
unrendered template.
+
+ A ``str`` value might be a template awaiting
``render_template_fields()``, so it must wait
+ for ``_build_response_kwargs()`` at ``execute()`` time. Any other type
(``int``, ``bool``,
+ ``float``) is already final at construction -- it only reaches here as
a literal, never via
+ Jinja rendering -- so an invalid one is rejected at Dag-parse time
instead of surfacing only
+ when the task runs.
+ """
+ for param_name in self._TOKEN_CEILING_PARAM_NAMES:
+ value = getattr(self, param_name)
+ if value is not None and not isinstance(value, str):
Review Comment:
`not isinstance(value, str)` reads "final literal", but an `XComArg` is
neither a str nor final: `max_output_tokens=budget()` now raises `ValueError:
'max_output_tokens' must be an integer, got
XComArg(<Task(_PythonDecoratedOperator): budget>).` at import time, taking the
whole Dag file with it, even though `render_template` resolves anything
exposing `.resolve(context)` (`templater.py:264-265`) so an XComArg is a
supported value for a template field. The `(bool, float)` denylist in
`_coerce_token_ceiling` has the mirror hole: `Decimal('10.5')` and
`Fraction(21, 2)` both slip past it and reach the wire as `10`, which is
exactly the silent truncation the comment there gives as the reason float is
rejected. An allowlist on both sides closes both: fire the eager check only for
`isinstance(value, (bool, float, int))`, and have `_coerce_token_ceiling`
reject anything that is not `int`-or-`str` with `bool` still excluded.
##########
providers/openai/docs/operators/openai.rst:
##########
@@ -53,6 +53,22 @@ 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"`` and
truncated
+``output_text``, so ``return_value`` will be the truncated text, not an error.
+
+A rendered ``max_output_tokens`` or ``max_tool_calls`` that is blank or
whitespace-only -- for
+example ``max_output_tokens="{{ params.tokens or '' }}"`` when
``params.tokens`` is unset -- is
Review Comment:
This `{{ params.tokens or '' }}` spelling is mine from the last round and it
does not survive `StrictUndefined`, which is the Dag default (`dag.py:472`):
with `tokens` absent from `params` it raises `UndefinedError: 'dict object' has
no attribute 'tokens'`, and it renders `''` only when `tokens` is a declared
param whose default is `None` -- not the "unset" this sentence describes. `{{
params.tokens | default('', true) }}` renders `''` in both cases, and `|
default('')` alone is no good since it yields `'None'`, which the docstring
says raises. Worth carrying the working spelling into the `exampleinclude`
target as well, since `example_openai.py` still shows only the
`response_kwargs` dict that these 16 new lines of prose steer users away from.
##########
providers/openai/docs/operators/openai.rst:
##########
@@ -53,6 +53,22 @@ 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
Review Comment:
This paragraph introduces the two ceilings as a matched pair, then explains
the observable consequence for only one of them. Nothing in the new warning
branches can ever fire for `max_tool_calls`: `IncompleteDetails.reason` is
`Literal['max_output_tokens', 'content_filter']` at the 2.37.0 floor, and the
SDK documents the cap as silent -- "Any further attempts to call a tool by the
model will be ignored" -- with no status change and no `incomplete_details`. So
a run that dropped tool calls partway through looks identical to a clean one,
in the logs and in `return_value` alike. One sentence saying so here would save
someone a debugging session.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -101,30 +123,117 @@ 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 cannot possibly still be an
unrendered template.
+
+ A ``str`` value might be a template awaiting
``render_template_fields()``, so it must wait
+ for ``_build_response_kwargs()`` at ``execute()`` time. Any other type
(``int``, ``bool``,
+ ``float``) is already final at construction -- it only reaches here as
a literal, never via
+ Jinja rendering -- so an invalid one is rejected at Dag-parse time
instead of surfacing only
+ when the task runs.
+ """
+ for param_name in self._TOKEN_CEILING_PARAM_NAMES:
+ value = getattr(self, param_name)
+ if value is not None and not isinstance(value, str):
+ 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 | 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
+ # int check below; float must also be rejected explicitly since
int(10.5) == 10 silently
+ # truncates instead of raising -- both can reach here as real Python
objects, not just
+ # strings, when a Dag uses render_template_as_native_obj=True.
+ if isinstance(value, (bool, float)):
+ 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() ==
""):
+ continue
+ 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 == "max_output_tokens":
+ self.log.warning(
+ "Response %s is incomplete (incomplete_details.reason=%s);
the returned output "
+ "text is truncated, not empty.",
Review Comment:
Following up on my comment on this line from the last round: the split
landed as asked, but we both missed a case. openai 2.37.0 documents
`max_output_tokens` as bounding "visible output tokens **and** reasoning
tokens", so a reasoning model can spend the whole ceiling on reasoning and
return no message item at all. Built a real pydantic-validated `Response` with
`status="incomplete"`, `reason="max_output_tokens"` and a single
`ResponseReasoningItem`: `output_text` is `''`, `execute()` returns `''`, and
this warning still says the text is "truncated, not empty". Since `model` is a
plain operator arg, `model="gpt-5"` plus a modest ceiling is ordinary config.
Branching on `response.output_text`, which you already have in hand, would be
honest in every case, and the guide's stronger promise that `return_value`
"will be the truncated text, not an error" needs the same softening.
--
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]
