Lee-W commented on code in PR #72150:
URL: https://github.com/apache/airflow/pull/72150#discussion_r4015555340
##########
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:
The warning branches on `response.output_text` now. truncated-not-empty only
when there is text, may-be-empty otherwise, with the reason logged in both.
--
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]
