kaxil commented on code in PR #72150:
URL: https://github.com/apache/airflow/pull/72150#discussion_r3963262704
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -101,30 +115,85 @@ 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")
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
@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``,
rejecting duplicates."""
+ response_kwargs = dict(self.response_kwargs)
+ for param_name, value in (
+ ("max_output_tokens", self.max_output_tokens),
+ ("max_tool_calls", self.max_tool_calls),
+ ):
+ if value is None:
+ continue
+ if param_name in response_kwargs:
+ raise ValueError(
+ f"{param_name!r} was set both as an operator argument and
in 'response_kwargs'; "
+ "set it in only one place."
+ )
+ 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:
+ self.log.warning(
+ "Response %s is incomplete (incomplete_details.reason=%s);
the returned output "
+ "text is truncated, not empty.",
Review Comment:
`incomplete_details.reason` can also be `content_filter`, where the filter
can fire before any output text is produced, so `output_text` may well be
empty. This message asserts truncated-not-empty for both reasons, and the new
test locks out the old "may be empty" wording for every incomplete response.
Scoping the truncation phrasing to `reason == "max_output_tokens"` and keeping
a may-be-empty hint for `content_filter` would keep the message honest in both
cases.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -101,30 +115,85 @@ 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")
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
@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}.")
Review Comment:
Now that this is templated, is there a way for a Dag to render "no ceiling"?
Under the default string rendering `{{ params.tokens }}` with `params.tokens =
None` renders to `'None'`, and `{{ params.tokens or '' }}` renders to `''`;
both land here as a ValueError. The `value is None` skip in
`_build_response_kwargs` isn't reachable from a template unless the Dag turns
on `render_template_as_native_obj=True`. Treating a blank rendered value as
unset would make the "vary by environment" case from the docs work without
native rendering.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -101,30 +115,85 @@ 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")
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
@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``,
rejecting duplicates."""
+ response_kwargs = dict(self.response_kwargs)
+ for param_name, value in (
+ ("max_output_tokens", self.max_output_tokens),
+ ("max_tool_calls", self.max_tool_calls),
+ ):
+ if value is None:
+ continue
+ if param_name in response_kwargs:
Review Comment:
This check doesn't depend on templating: `response_kwargs` is never
rendered, and the argument is non-None from `__init__` onward, so the conflict
is fully determined at Dag parse time. Raising it in `__init__` would surface a
misconfigured task as an import error rather than a failure on the first run.
##########
providers/openai/docs/operators/openai.rst:
##########
@@ -53,6 +53,14 @@ 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. These
are *token*-level
Review Comment:
`max_tool_calls` caps the number of built-in tool calls, not tokens, so
describing both as token-level ceilings isn't quite right. The same sentence is
in the operator docstring.
--
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]
