kaxil commented on code in PR #72051:
URL: https://github.com/apache/airflow/pull/72051#discussion_r3968252182
##########
providers/openai/src/airflow/providers/openai/hooks/openai.py:
##########
@@ -619,21 +619,26 @@ def delete_vector_store_file(self, vector_store_id: str,
file_id: str) -> Vector
def create_batch(
Review Comment:
`create_batch` is the only one of this hook's eleven `create_*` methods
without a `**kwargs` passthrough. `create_chat_completion` (:245),
`create_response` (:266), `create_embeddings` (:497), `create_vector_store`
(:548) and the other six all forward arbitrary SDK options, and the operator
siblings expose the same hatch by name (`embedding_kwargs`, `response_kwargs`).
Wiring up the two params the hook already had still leaves
`output_expires_after` unreachable, and that one sets the expiry on a batch's
output and error files, which is what `cleanup_batch_output_file` in
`tests/system/openai/example_trigger_batch_operator.py:87` hand-rolls. So the
next person who wants a batch option writes this PR again.
`**kwargs` here plus `batch_kwargs` on the operator would cover it, and it
also settles `completion_window`: the SDK types it `Required[Literal["24h"]]`
in both the `openai>=2.37.0` floor and the current 2.54.0, so as a named
parameter it can only ever hold its own default, and behind a passthrough it
needs no named parameter at all.
##########
providers/openai/src/airflow/providers/openai/hooks/openai.py:
##########
@@ -619,21 +619,26 @@ def delete_vector_store_file(self, vector_store_id: str,
file_id: str) -> Vector
def create_batch(
self,
file_id: str,
- endpoint: Literal["/v1/chat/completions", "/v1/embeddings",
"/v1/completions"],
+ endpoint: str,
metadata: dict[str, str] | None = None,
completion_window: Literal["24h"] = "24h",
) -> Batch:
"""
Create a batch for a given model and files.
:param file_id: The ID of the file to be used for this batch.
- :param endpoint: The endpoint to use for this batch. Allowed values
include:
- '/v1/chat/completions', '/v1/embeddings', '/v1/completions'.
+ :param endpoint: The endpoint to use for this batch. Allowed values
are determined by the
+ OpenAI Batch API; see the OpenAI documentation for the current
list.
Review Comment:
The endpoint doc chain dead-ends here. The operator's docstring
(`operators/openai.py:142-144`) points at this method, and this one says to see
the OpenAI documentation with no link. Dropping the `Literal` is what took the
values out of the signature, so a reader has less to go on than before.
Both siblings in this module link their exact page (`/embeddings/create` at
`operators/openai.py:47`, `/responses/create` at :101), and
`https://platform.openai.com/docs/api-reference/batch/create` would close it.
Worth saying the widening did fix genuinely stale docs: the SDK is at eight
endpoints against the old three.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -138,43 +138,59 @@ class OpenAITriggerBatchOperator(BaseOperator):
"""
Operator that triggers an OpenAI Batch API endpoint and waits for the
batch to complete.
- :param file_id: Required. The ID of the batch file to trigger.
- :param endpoint: Required. The OpenAI Batch API endpoint to trigger.
+ :param file_id: Required. The ID of the batch file to trigger. (templated)
+ :param endpoint: Required. The OpenAI Batch API endpoint to trigger.
(templated) Allowed values
+ are determined by the OpenAI Batch API; see
+ :meth:`~airflow.providers.openai.hooks.openai.OpenAIHook.create_batch`.
:param conn_id: Optional. The OpenAI connection ID to use. Defaults to
'openai_default'.
:param deferrable: Optional. Run operator in the deferrable mode.
:param wait_seconds: Optional. Number of seconds between checks. Only used
when ``deferrable`` is False.
Defaults to 3 seconds.
:param timeout: Optional. The amount of time, in seconds, to wait for the
request to complete.
- Only used when ``deferrable`` is False. Defaults to 24 hour, which is
the SLA for OpenAI Batch API.
+ Applies in both deferrable and non-deferrable mode. Defaults to 24
hours, 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:
An observation rather than an ask. `metadata` is templated now, and on a DAG
with `render_template_as_native_obj=True` Jinja's native environment runs
`literal_eval` over every rendered leaf, so `metadata={"attempt": "1"}` reaches
`batches.create` as `{"attempt": 1}` against the SDK's `Dict[str, str]`.
`endpoint` is safe because an API path never parses as a literal.
Nothing in the tree guards this for any templated dict, so I am not asking
you to be the first here. The smaller half that does have precedent is a
`template_fields_renderers = {"metadata": "json"}` entry, as in
`s3_tables.py:66`, `s3_vectors.py:159` and `dataproc_metastore.py:144`.
##########
providers/openai/tests/unit/openai/operators/test_openai.py:
##########
@@ -125,6 +125,75 @@ 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):
Review Comment:
This and `..._create_batch_default_passthrough` (:128) differ only in the
`metadata` value; both assert `completion_window="24h"` (:148 and :174), so the
explicit one cannot distinguish itself from the default and adds no coverage
for that parameter.
Same shape for the deferred pair below, where the existing parametrized test
already picked up `assert trigger.poll_interval == 60` and the new function
repeats the whole body to assert 5. The file reaches for
`@pytest.mark.parametrize` five times already, so two parametrized functions
would cover strictly more in about half the lines.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -138,43 +138,59 @@ class OpenAITriggerBatchOperator(BaseOperator):
"""
Operator that triggers an OpenAI Batch API endpoint and waits for the
batch to complete.
- :param file_id: Required. The ID of the batch file to trigger.
- :param endpoint: Required. The OpenAI Batch API endpoint to trigger.
+ :param file_id: Required. The ID of the batch file to trigger. (templated)
+ :param endpoint: Required. The OpenAI Batch API endpoint to trigger.
(templated) Allowed values
+ are determined by the OpenAI Batch API; see
+ :meth:`~airflow.providers.openai.hooks.openai.OpenAIHook.create_batch`.
:param conn_id: Optional. The OpenAI connection ID to use. Defaults to
'openai_default'.
:param deferrable: Optional. Run operator in the deferrable mode.
:param wait_seconds: Optional. Number of seconds between checks. Only used
when ``deferrable`` is False.
Defaults to 3 seconds.
:param timeout: Optional. The amount of time, in seconds, to wait for the
request to complete.
- Only used when ``deferrable`` is False. Defaults to 24 hour, which is
the SLA for OpenAI Batch API.
+ Applies in both deferrable and non-deferrable mode. Defaults to 24
hours, 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")
def __init__(
self,
file_id: str,
- endpoint: Literal["/v1/chat/completions", "/v1/embeddings",
"/v1/completions"],
+ endpoint: str,
conn_id: str = OpenAIHook.default_conn_name,
deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
wait_seconds: float = 3,
timeout: float = 24 * 60 * 60,
wait_for_completion: bool = True,
+ *,
+ metadata: dict[str, str] | None = None,
+ completion_window: Literal["24h"] = "24h",
+ poll_interval: float = 60,
Review Comment:
A question rather than a request: `wait_seconds` (:173) and `poll_interval`
now carry the same sentence, 20x apart in default, with `deferrable` deciding
which one is live and `deferrable` itself defaulting from
`operators.default_deferrable` (:172). On a deployment with that turned on,
`wait_seconds=10` silently does nothing.
`AnthropicBatchOperator` went the other way, with a single `poll_interval`
covering both paths
(`providers/anthropic/src/airflow/providers/anthropic/operators/batch.py:90`)
and the keyword adapted at the sync call site (:143). Unifying here would move
the deferred default from 60 to 3, so it needs a compat story and may well not
be worth it. A log line naming which knob is in effect when the other is set
would at least remove the silent case.
--
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]
