pankajastro commented on code in PR #72150:
URL: https://github.com/apache/airflow/pull/72150#discussion_r3985710379
##########
providers/openai/tests/unit/openai/operators/test_openai.py:
##########
@@ -105,6 +108,213 @@ 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(0, id="zero"),
+ pytest.param(-1, id="negative"),
+ pytest.param("-5", id="negative-string"),
+ pytest.param(10.5, id="float"),
+ pytest.param(True, id="bool-true"),
+ pytest.param(False, id="bool-false"),
+ pytest.param("None", id="literal-none-string"),
+ ],
+ )
+ @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(
+ "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):
+ # The conflict is a construction-time (Dag-parse) failure now: a blank
operator_value
Review Comment:
"design decision C2's 'unset' case" and "the C3 check" -- these labels don't
appear anywhere else in this PR, the docstring, or the repo. Could this be
reworded to plain prose, e.g. `# A blank operator_value must still conflict
with response_kwargs; that's checked at construction time, before rendering.`,
dropping the C2/C3 shorthand?
##########
providers/openai/tests/unit/openai/operators/test_openai.py:
##########
@@ -97,6 +98,8 @@ def test_openai_response_operator_execute():
result = operator.execute(Context())
assert result == "haiku text"
+ # Backward-compat lock: without max_output_tokens/max_tool_calls,
create_response must be
Review Comment:
Given `test_openai_response_operator_execute` plus the exact-args assertion
immediately below already communicate this, does the "Backward-compat lock"
comment add anything the test name and assertion don't already show?
##########
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]:
+ """Merge the token-ceiling arguments into ``response_kwargs``,
skipping unset ceilings."""
+ 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),
+ ):
+ # A blank or whitespace-only rendered template means "no ceiling
this run"
Review Comment:
This "blank/whitespace-only rendered value means unset" rationale is already
spelled out in the class docstring and in `openai.rst` -- could this comment be
trimmed to a short pointer, e.g. `# Blank means unset; conflict with
response_kwargs was already checked in __init__.`, instead of restating the
full explanation a third time?
##########
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 (
Review Comment:
This `(("max_output_tokens", self.max_output_tokens), ("max_tool_calls",
self.max_tool_calls))` tuple is duplicated verbatim in `_build_response_kwargs`
below (line 182). Could both loops iterate over one shared class-level tuple of
ceiling-param names instead, so a future third ceiling parameter can't be added
to one loop and silently forgotten in the other?
##########
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:
`_validate_no_response_kwargs_conflict()` runs in `__init__`, so the
operator-arg/`response_kwargs` conflict is caught at Dag-parse time, but
`_coerce_token_ceiling()` is only reached from here, inside
`_build_response_kwargs()`, which runs at `execute()`. Doesn't that mean a
plain literal like `max_output_tokens=0` or `max_output_tokens="fifty"` still
passes Dag parsing silently and only fails once the task actually runs -- the
class of late failure this PR's parse-time conflict check was meant to avoid?
Should the type/positive-int check also move into `__init__` when the value
isn't a template?
--
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]
