This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 1b2708e875b Make OpenAI batch options reachable from
OpenAITriggerBatchOperator (#72051)
1b2708e875b is described below
commit 1b2708e875be58f3988245d16eb741193bb1f4e3
Author: Wei Lee <[email protected]>
AuthorDate: Wed Sep 16 07:16:23 2026 +0800
Make OpenAI batch options reachable from OpenAITriggerBatchOperator (#72051)
* Let OpenAITriggerBatchOperator set batch metadata, window and poll
interval
The operator could not reach the metadata and completion_window arguments
its
own hook accepts, could not template the endpoint, and polled every 60
seconds
in deferrable mode with no way to change that - wait_seconds only applies to
the synchronous path. All four are now settable, and every default matches
the
previous hard-coded behaviour.
* Let mypy-checked Dags template the OpenAI batch endpoint
The endpoint field entered template_fields but kept its three-value Literal
annotation, so a mypy-checked Dag passing "{{ var.value.batch_endpoint }}"
got an arg-type error - the new templating was unusable without a type
ignore at every call site.
The hard-coded list was stale anyway: the pinned SDK accepts eight
endpoints, and enumerating them here would only move the expiry date. The
allowed values are now left to the OpenAI Batch API, which validates them
server-side.
* Correct the timeout scope in OpenAITriggerBatchOperator docs
The timeout entry claimed it only applied when deferrable was False, but
the deferred path passes it to OpenAIBatchTrigger, so it bounds that wait
too - misleading now that the docstring names a mode for every other knob.
Dropping the template_fields equality assertion: it restated the class
attribute rather than exercising it, and the rendering test next to it
already covers the behaviour that matters.
* Let OpenAITriggerBatchOperator forward arbitrary OpenAI batch options
OpenAIHook.create_batch was the only one of the hook's create_* methods
without a keyword passthrough, so every option the OpenAI SDK exposes on
a batch needed another change here before a Dag author could reach it.
output_expires_after is the case in hand: it sets the expiry on a batch's
output and error files, and the system example Dag hand-rolls that
cleanup because the option is out of reach.
completion_window goes away as a named operator parameter for the same
reason. The SDK types it as a single-value Literal, so a named parameter
could only ever carry its own default, and the passthrough covers anyone
who wants to set it explicitly.
* Link the OpenAI Batch API reference from create_batch
Widening endpoint from a three-value Literal to str took the accepted
values out of the signature, so the docstring became the only place a
reader could learn them — and it pointed at "the OpenAI documentation"
without saying where that is. The sibling methods in this provider each
link their own API page.
* Log which OpenAI batch poll interval is actually in effect
wait_seconds and poll_interval mean the same thing and differ 20x in
default, and deferrable picks which one is live. Because deferrable
itself defaults from operators.default_deferrable, a Dag author on a
deployment with that turned on can set wait_seconds and have it do
nothing, with no way to tell from the task log.
* Declare the OpenAI batch metadata template field as JSON
metadata became templated in this PR, and every other operator with a
templated mapping declares how it should be rendered. Keeping the
declaration alongside them means this operator does not become the
odd one out if renderer-driven display returns.
* Assert the OpenAI batch poll-interval logs without caplog
The structlog-style membership assertion only works where the Task SDK
logging package is importable, so the provider's compatibility legs for
Airflow 2.11 and 3.0 hit pytest's stock LogCaptureFixture and failed with
a TypeError. Mocking the operator's logger holds across every version in
the matrix.
---
.../src/airflow/providers/openai/hooks/openai.py | 16 ++-
.../airflow/providers/openai/operators/openai.py | 54 +++++++--
.../openai/tests/unit/openai/hooks/test_openai.py | 13 ++-
.../tests/unit/openai/operators/test_openai.py | 130 ++++++++++++++++++++-
4 files changed, 196 insertions(+), 17 deletions(-)
diff --git a/providers/openai/src/airflow/providers/openai/hooks/openai.py
b/providers/openai/src/airflow/providers/openai/hooks/openai.py
index 97dcc4ceec8..e7af68ca38c 100644
--- a/providers/openai/src/airflow/providers/openai/hooks/openai.py
+++ b/providers/openai/src/airflow/providers/openai/hooks/openai.py
@@ -619,21 +619,29 @@ class OpenAIHook(BaseHook):
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",
+ **kwargs: Any,
) -> 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
https://platform.openai.com/docs/api-reference/batch/create for
+ the current list.
:param metadata: A set of key-value pairs that can be attached to an
object.
:param completion_window: The time window for the batch to complete.
Default is 24 hours.
"""
batch = self.conn.batches.create(
- input_file_id=file_id, endpoint=endpoint, metadata=metadata,
completion_window=completion_window
+ input_file_id=file_id,
+ # endpoint is intentionally str (not the SDK's Literal) so
templated values type-check;
+ # the OpenAI service validates the actual value.
+ endpoint=endpoint, # type: ignore[arg-type]
+ metadata=metadata,
+ completion_window=completion_window,
+ **kwargs,
)
return batch
diff --git a/providers/openai/src/airflow/providers/openai/operators/openai.py
b/providers/openai/src/airflow/providers/openai/operators/openai.py
index 0dad7792f18..2507dd075f2 100644
--- a/providers/openai/src/airflow/providers/openai/operators/openai.py
+++ b/providers/openai/src/airflow/providers/openai/operators/openai.py
@@ -19,7 +19,7 @@ from __future__ import annotations
from collections.abc import Sequence
from functools import cached_property
-from typing import TYPE_CHECKING, Any, ClassVar, Literal
+from typing import TYPE_CHECKING, Any, ClassVar
from airflow.providers.common.compat.sdk import BaseOperator, conf
from airflow.providers.openai.exceptions import OpenAIBatchJobException
@@ -279,43 +279,61 @@ 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 batch_kwargs: Optional. Additional keyword arguments to pass to the
OpenAI `create_batch`
+ method — for example `output_expires_after`, which sets the expiry on
the batch's output and
+ error files. Defaults to None.
+ :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")
+ template_fields_renderers = {"metadata": "json"}
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,
+ batch_kwargs: dict | None = None,
+ poll_interval: float = 60,
**kwargs: Any,
):
super().__init__(**kwargs)
- self.conn_id = conn_id
self.file_id = file_id
self.endpoint = endpoint
+ self.conn_id = conn_id
self.deferrable = deferrable
self.wait_seconds = wait_seconds
self.timeout = timeout
self.wait_for_completion = wait_for_completion
+ self.metadata = metadata
+ self.batch_kwargs = batch_kwargs or {}
+ self.poll_interval = poll_interval
+
self.batch_id: str | None = None
@cached_property
@@ -324,22 +342,38 @@ class OpenAITriggerBatchOperator(BaseOperator):
return OpenAIHook(conn_id=self.conn_id)
def execute(self, context: Context) -> str | None:
- batch = self.hook.create_batch(file_id=self.file_id,
endpoint=self.endpoint)
+ batch = self.hook.create_batch(
+ file_id=self.file_id,
+ endpoint=self.endpoint,
+ metadata=self.metadata,
+ **self.batch_kwargs,
+ )
self.batch_id = batch.id
if self.wait_for_completion:
if self.deferrable:
+ self.log.info(
+ "Deferring batch %s, polling every %s seconds via
poll_interval "
+ "(wait_seconds is not used in deferrable mode)",
+ self.batch_id,
+ self.poll_interval,
+ )
self.defer(
timeout=self.execution_timeout,
trigger=OpenAIBatchTrigger(
conn_id=self.conn_id,
batch_id=self.batch_id,
- poll_interval=60,
+ poll_interval=self.poll_interval,
timeout=self.timeout,
),
method_name="execute_complete",
)
else:
- self.log.info("Waiting for batch %s to complete",
self.batch_id)
+ self.log.info(
+ "Waiting for batch %s to complete, polling every %s
seconds via wait_seconds "
+ "(poll_interval is not used in non-deferrable mode)",
+ self.batch_id,
+ self.wait_seconds,
+ )
self.hook.wait_for_batch(self.batch_id,
wait_seconds=self.wait_seconds, timeout=self.timeout)
return self.batch_id
diff --git a/providers/openai/tests/unit/openai/hooks/test_openai.py
b/providers/openai/tests/unit/openai/hooks/test_openai.py
index d132a6a3d95..4115f0c2c69 100644
--- a/providers/openai/tests/unit/openai/hooks/test_openai.py
+++ b/providers/openai/tests/unit/openai/hooks/test_openai.py
@@ -608,8 +608,19 @@ def test_delete_vector_store_file(mock_openai_hook):
def test_create_batch(mock_openai_hook, mock_terminated_batch):
mock_openai_hook.conn.batches.create.return_value = mock_terminated_batch
- batch = mock_openai_hook.create_batch(endpoint="/v1/chat/completions",
file_id=FILE_ID)
+ batch = mock_openai_hook.create_batch(
+ endpoint="/v1/chat/completions",
+ file_id=FILE_ID,
+ output_expires_after={"anchor": "created_at", "seconds": 3600},
+ )
assert batch.id == mock_terminated_batch.id
+ mock_openai_hook.conn.batches.create.assert_called_once_with(
+ input_file_id=FILE_ID,
+ endpoint="/v1/chat/completions",
+ metadata=None,
+ completion_window="24h",
+ output_expires_after={"anchor": "created_at", "seconds": 3600},
+ )
def test_get_batch(mock_openai_hook, mock_terminated_batch):
diff --git a/providers/openai/tests/unit/openai/operators/test_openai.py
b/providers/openai/tests/unit/openai/operators/test_openai.py
index d4ba4403840..5842cfdec6c 100644
--- a/providers/openai/tests/unit/openai/operators/test_openai.py
+++ b/providers/openai/tests/unit/openai/operators/test_openai.py
@@ -18,6 +18,7 @@ from __future__ import annotations
from decimal import Decimal
from fractions import Fraction
+from unittest import mock
from unittest.mock import Mock
import jinja2
@@ -483,8 +484,76 @@ def
test_openai_trigger_batch_operator_not_deferred(mock_batch, wait_for_complet
assert batch_id == BATCH_ID
[email protected]("wait_for_completion", [True, False])
-def test_openai_trigger_batch_operator_with_deferred(mock_batch,
wait_for_completion):
[email protected](
+ ("metadata", "batch_kwargs", "expected_kwargs"),
+ [
+ pytest.param(None, None, {}, id="no-passthrough"),
+ pytest.param(
+ {"key": "value"},
+ {"output_expires_after": {"anchor": "created_at", "seconds":
3600}},
+ {"output_expires_after": {"anchor": "created_at", "seconds":
3600}},
+ id="metadata-and-batch-kwargs",
+ ),
+ ],
+)
+def test_openai_trigger_batch_operator_create_batch_passthrough(
+ mock_batch, metadata, batch_kwargs, expected_kwargs
+):
+ """metadata/batch_kwargs reach create_batch verbatim; unset batch_kwargs
forwards none."""
+ operator = OpenAITriggerBatchOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ file_id=FILE_ID,
+ endpoint=BATCH_ENDPOINT,
+ metadata=metadata,
+ batch_kwargs=batch_kwargs,
+ 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=metadata,
+ **expected_kwargs,
+ )
+
+
+def test_openai_trigger_batch_operator_templates_endpoint_and_metadata():
+ operator = OpenAITriggerBatchOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ file_id=FILE_ID,
+ endpoint="{{ ti.endpoint }}",
+ metadata={"run": "{{ ti.run_id }}"},
+ )
+
+ class FakeTaskInstance:
+ endpoint = BATCH_ENDPOINT
+ run_id = "run-123"
+
+ operator.render_template_fields(context={"ti": FakeTaskInstance()})
+
+ assert operator.endpoint == BATCH_ENDPOINT
+ assert operator.metadata == {"run": "run-123"}
+
+
[email protected](
+ ("wait_for_completion", "poll_interval_kwargs", "expected_poll_interval"),
+ [
+ pytest.param(False, {}, None, id="not-deferred"),
+ pytest.param(True, {}, 60, id="deferred-default-poll-interval"),
+ pytest.param(True, {"poll_interval": 5}, 5,
id="deferred-custom-poll-interval"),
+ ],
+)
+def test_openai_trigger_batch_operator_with_deferred(
+ mock_batch, wait_for_completion, poll_interval_kwargs,
expected_poll_interval
+):
operator = OpenAITriggerBatchOperator(
task_id=TASK_ID,
conn_id=CONN_ID,
@@ -492,6 +561,7 @@ def
test_openai_trigger_batch_operator_with_deferred(mock_batch, wait_for_comple
endpoint=BATCH_ENDPOINT,
deferrable=True,
wait_for_completion=wait_for_completion,
+ **poll_interval_kwargs,
)
mock_hook_instance = Mock(spec=OpenAIHook)
mock_hook_instance.get_batch.return_value = mock_batch
@@ -503,11 +573,67 @@ def
test_openai_trigger_batch_operator_with_deferred(mock_batch, wait_for_comple
with pytest.raises(TaskDeferred) as exc:
operator.execute(context)
assert isinstance(exc.value.trigger, OpenAIBatchTrigger)
+ assert exc.value.trigger.poll_interval == expected_poll_interval
else:
batch_id = operator.execute(context)
assert batch_id == BATCH_ID
[email protected](OpenAITriggerBatchOperator, "log")
+def test_openai_trigger_batch_operator_not_deferred_logs_active_knob(mock_log,
mock_batch):
+ """Non-deferred mode names wait_seconds' value and states poll_interval is
unused."""
+ operator = OpenAITriggerBatchOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ file_id=FILE_ID,
+ endpoint=BATCH_ENDPOINT,
+ deferrable=False,
+ wait_seconds=7,
+ poll_interval=99,
+ )
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_hook_instance.get_batch.return_value = mock_batch
+ mock_hook_instance.create_batch.return_value = mock_batch
+ operator.hook = mock_hook_instance
+
+ operator.execute(Context())
+
+ mock_log.info.assert_any_call(
+ "Waiting for batch %s to complete, polling every %s seconds via
wait_seconds "
+ "(poll_interval is not used in non-deferrable mode)",
+ BATCH_ID,
+ 7,
+ )
+
+
[email protected](OpenAITriggerBatchOperator, "log")
+def test_openai_trigger_batch_operator_deferred_logs_active_knob(mock_log,
mock_batch):
+ """Deferred mode names poll_interval's value and states wait_seconds is
unused."""
+ operator = OpenAITriggerBatchOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ file_id=FILE_ID,
+ endpoint=BATCH_ENDPOINT,
+ deferrable=True,
+ wait_seconds=71,
+ poll_interval=13,
+ )
+ mock_hook_instance = Mock(spec=OpenAIHook)
+ mock_hook_instance.get_batch.return_value = mock_batch
+ mock_hook_instance.create_batch.return_value = mock_batch
+ operator.hook = mock_hook_instance
+
+ with pytest.raises(TaskDeferred):
+ operator.execute(Context())
+
+ mock_log.info.assert_any_call(
+ "Deferring batch %s, polling every %s seconds via poll_interval "
+ "(wait_seconds is not used in deferrable mode)",
+ BATCH_ID,
+ 13,
+ )
+
+
class TestOpenAITriggerBatchOperatorExecuteComplete:
def _operator(self):
return OpenAITriggerBatchOperator(
