kaxil commented on code in PR #72051:
URL: https://github.com/apache/airflow/pull/72051#discussion_r3957916929
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -148,13 +148,18 @@ class OpenAITriggerBatchOperator(BaseOperator):
Only used when ``deferrable`` is False. Defaults to 24 hour, which is
the SLA for OpenAI Batch API.
:param wait_for_completion: Optional. Whether to wait for the batch to
complete. If set to False, the operator
will return immediately after triggering the batch. Defaults to True.
+ :param metadata: Optional. A set of key-value pairs that can be attached
to the batch. (templated)
+ :param completion_window: Optional. The time window for the batch to
complete. Defaults to 24 hours,
+ the only value OpenAI currently accepts.
+ :param poll_interval: Optional. Number of seconds between checks. Only
used when ``deferrable`` is True.
Review Comment:
Since this docstring now spells out which mode each knob applies to: the
`timeout` entry a few lines up says "Only used when ``deferrable`` is False",
but it goes into `OpenAIBatchTrigger(timeout=self.timeout)` below, so it bounds
the deferred wait as well.
##########
providers/openai/tests/unit/openai/operators/test_openai.py:
##########
@@ -125,6 +125,79 @@ def
test_openai_trigger_batch_operator_not_deferred(mock_batch, wait_for_complet
assert batch_id == BATCH_ID
+def
test_openai_trigger_batch_operator_create_batch_default_passthrough(mock_batch):
+ """No metadata/completion_window passed: create_batch must see the pre-PR
defaults."""
+ operator = OpenAITriggerBatchOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ file_id=FILE_ID,
+ endpoint=BATCH_ENDPOINT,
+ deferrable=False,
+ wait_for_completion=False,
+ )
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_hook_instance.create_batch.return_value = mock_batch
+ operator.hook = mock_hook_instance
+
+ operator.execute(Context())
+
+ mock_hook_instance.create_batch.assert_called_once_with(
+ file_id=FILE_ID,
+ endpoint=BATCH_ENDPOINT,
+ metadata=None,
+ completion_window="24h",
+ )
+
+
+def test_openai_trigger_batch_operator_create_batch_passthrough(mock_batch):
+ """metadata/completion_window are reachable through the operator and
forwarded as-is."""
+ operator = OpenAITriggerBatchOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ file_id=FILE_ID,
+ endpoint=BATCH_ENDPOINT,
+ metadata={"key": "value"},
+ completion_window="24h",
+ deferrable=False,
+ wait_for_completion=False,
+ )
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_hook_instance.create_batch.return_value = mock_batch
+ operator.hook = mock_hook_instance
+
+ operator.execute(Context())
+
+ mock_hook_instance.create_batch.assert_called_once_with(
+ file_id=FILE_ID,
+ endpoint=BATCH_ENDPOINT,
+ metadata={"key": "value"},
+ completion_window="24h",
+ )
+
+
+def test_openai_trigger_batch_operator_template_fields():
+ assert OpenAITriggerBatchOperator.template_fields == ("file_id",
"endpoint", "metadata")
Review Comment:
This can only fail if someone edits the assertion and the class attribute
apart, so it restates the source line rather than testing behaviour. The render
test right below already covers the part that matters.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -148,13 +148,18 @@ class OpenAITriggerBatchOperator(BaseOperator):
Only used when ``deferrable`` is False. Defaults to 24 hour, which is
the SLA for OpenAI Batch API.
:param wait_for_completion: Optional. Whether to wait for the batch to
complete. If set to False, the operator
will return immediately after triggering the batch. Defaults to True.
+ :param metadata: Optional. A set of key-value pairs that can be attached
to the batch. (templated)
+ :param completion_window: Optional. The time window for the batch to
complete. Defaults to 24 hours,
+ the only value OpenAI currently accepts.
+ :param poll_interval: Optional. Number of seconds between checks. Only
used when ``deferrable`` is True.
+ Defaults to 60 seconds.
.. seealso::
For more information on how to use this operator, please take a look
at the guide:
:ref:`howto/operator:OpenAITriggerBatchOperator`
"""
- template_fields: Sequence[str] = ("file_id",)
+ template_fields: Sequence[str] = ("file_id", "endpoint", "metadata")
Review Comment:
mypy-checked DAGs can't actually use this: `endpoint` is still annotated
`Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"]`, so
`endpoint="{{ var.value.batch_endpoint }}"` is an `arg-type` error and needs a
`# type: ignore`. Widening it to `str` would make the new templating usable,
but the hook's `create_batch` carries the same Literal, so that one needs
widening too or the provider's own mypy job breaks at the call site in
`execute`. The Literal is behind the pinned SDK anyway: `openai>=2.37.0` also
accepts `/v1/responses`, `/v1/moderations`, `/v1/images/generations`,
`/v1/images/edits` and `/v1/videos`.
--
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]
