Lee-W commented on code in PR #72150:
URL: https://github.com/apache/airflow/pull/72150#discussion_r4015710128
##########
providers/openai/tests/unit/openai/operators/test_openai.py:
##########
@@ -105,6 +108,268 @@ 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("-5", id="negative-string"),
+ pytest.param("None", id="literal-none-string"),
+ pytest.param(Decimal("10.5"), id="decimal"),
+ pytest.param(Fraction(21, 2), id="fraction"),
+ ],
+ )
+ @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(
+ "invalid_value",
+ [
+ pytest.param(0, id="zero"),
+ pytest.param(-1, id="negative"),
+ pytest.param(10.5, id="float"),
+ pytest.param(True, id="bool-true"),
+ pytest.param(False, id="bool-false"),
+ ],
+ )
+ @pytest.mark.parametrize("param_name", ["max_output_tokens",
"max_tool_calls"])
+ def test_non_string_invalid_ceiling_raises_at_construction(self,
param_name, invalid_value):
+ with pytest.raises(ValueError, match=param_name):
+ OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ **{param_name: invalid_value},
+ )
+
+ @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):
+ # A blank operator_value must still conflict with response_kwargs;
that's checked at
+ # construction time, before rendering. pytest.raises() itself fails
with "DID NOT RAISE"
+ # if construction succeeded, so there's no operator instance
afterwards to assert
+ # anything further against.
+ with pytest.raises(ValueError, match=param_name):
+ OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ response_kwargs={param_name: 50},
+ **{param_name: operator_value},
+ )
+
+ def test_conflict_error_precedes_literal_type_error(self):
+ # 0 is both invalid on its own (not positive) and conflicting with
response_kwargs; the
+ # conflict message must win, since fixing the duplicate is the
actionable first step.
+ with pytest.raises(ValueError, match="was set both as an operator
argument"):
+ OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ response_kwargs={"max_output_tokens": 50},
+ max_output_tokens=0,
+ )
+
+ def test_xcom_arg_ceiling_does_not_fail_on_construction(self):
+ with DAG("test_dag", schedule=None) as dag:
+ upstream = BaseOperator(task_id="upstream")
+
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ max_output_tokens=upstream.output,
+ dag=dag,
+ )
+
+ assert isinstance(operator.max_output_tokens, XComArg)
+
+ @pytest.mark.parametrize(
+ "blank_value", [pytest.param("", id="empty"), pytest.param(" ",
id="whitespace")]
+ )
+ @pytest.mark.parametrize("param_name", ["max_output_tokens",
"max_tool_calls"])
+ def test_blank_ceiling_is_treated_as_unset(self, param_name, blank_value):
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID, conn_id=CONN_ID, input_text="Write a haiku.",
**{param_name: blank_value}
+ )
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_hook_instance.create_response.return_value =
_build_completed_response()
+ operator.hook = mock_hook_instance
+
+ operator.execute(Context())
+
+ call_kwargs = mock_hook_instance.create_response.call_args.kwargs
+ assert param_name not in call_kwargs
+
+ def test_max_output_tokens_and_max_tool_calls_are_templated(self):
+ operator = OpenAIResponseOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ input_text="Write a haiku.",
+ max_output_tokens="{{ params.tokens }}",
Review Comment:
Added
`test_default_filter_idiom_renders_blank_under_strict_undefined_dag_binding`
and `test_or_fallback_idiom_raises_under_strict_undefined_dag_binding`.
--
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]
