Lee-W commented on code in PR #72150:
URL: https://github.com/apache/airflow/pull/72150#discussion_r4015540814
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -101,30 +120,105 @@ 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
+ self._validate_no_response_kwargs_conflict()
+
+ def _validate_no_response_kwargs_conflict(self) -> None:
+ """Reject a ceiling set both as an operator argument and in
``response_kwargs``."""
+ for param_name, value in (
+ ("max_output_tokens", self.max_output_tokens),
+ ("max_tool_calls", self.max_tool_calls),
+ ):
+ 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."
+ )
@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]:
Review Comment:
`max_output_tokens=0` is rejected at construction now.
`_validate_literal_ceiling_values` runs from `__init__`. A plain literal
string like `"fifty"` still waits for `execute()`: a `str` cannot be told apart
from an unrendered template without guessing at Jinja delimiters, and getting
that wrong would break valid templates.
--
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]
