You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "Lee-W (via GitHub)" <gi...@apache.org> on 2026/08/25 07:58:09 UTC
[PR] Let OpenAITriggerBatchOperator set batch metadata, window and poll interval [airflow]
Lee-W opened a new pull request, #72051:
URL: https://github.com/apache/airflow/pull/72051
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.
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: https://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
For user-facing UI changes, please attach before/after screenshots (or a short
screen recording) so reviewers can assess the visual impact.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
- [ ] Yes (please specify the tool below)
<!--
Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make OpenAI batch options reachable from OpenAITriggerBatchOperator [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #72051:
URL: https://github.com/apache/airflow/pull/72051#discussion_r4015796958
##########
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:
Added `template_fields_renderers = {"metadata": "json"}`.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Let OpenAITriggerBatchOperator set batch metadata, window and poll interval [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #72051:
URL: https://github.com/apache/airflow/pull/72051#discussion_r3966873161
##########
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:
Sure, update the `timeout` part as well
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Let OpenAITriggerBatchOperator set batch metadata, window and poll interval [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make OpenAI batch options reachable from OpenAITriggerBatchOperator [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil merged PR #72051:
URL: https://github.com/apache/airflow/pull/72051
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make OpenAI batch options reachable from OpenAITriggerBatchOperator [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #72051:
URL: https://github.com/apache/airflow/pull/72051#discussion_r4015814458
##########
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:
I didn't gate it on "the other one was explicitly set". `execute()` now logs the live knob in both branches — its value, and that the other one does nothing in that mode.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Let OpenAITriggerBatchOperator set batch metadata, window and poll interval [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make OpenAI batch options reachable from OpenAITriggerBatchOperator [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #72051:
URL: https://github.com/apache/airflow/pull/72051#discussion_r4015833247
##########
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:
`create_batch`'s `:param endpoint:` now links https://platform.openai.com/docs/api-reference/batch/create instead of pointing at "the OpenAI documentation".
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make OpenAI batch options reachable from OpenAITriggerBatchOperator [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #72051:
URL: https://github.com/apache/airflow/pull/72051#discussion_r4015838722
##########
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:
Both pairs are one parametrized function each now.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Make OpenAI batch options reachable from OpenAITriggerBatchOperator [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #72051:
URL: https://github.com/apache/airflow/pull/72051#discussion_r4015765294
##########
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` now takes `**kwargs` and forwards them to `batches.create`; the operator exposes them as `batch_kwargs`, so `output_expires_after` and whatever the SDK adds next are reachable without another named parameter.
Also dropped the operator's `completion_window`.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Let OpenAITriggerBatchOperator set batch metadata, window and poll interval [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #72051:
URL: https://github.com/apache/airflow/pull/72051#discussion_r3966885460
##########
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:
Dropped this assertion.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Let OpenAITriggerBatchOperator set batch metadata, window and poll interval [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #72051:
URL: https://github.com/apache/airflow/pull/72051#discussion_r3966869819
##########
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:
`endpoint` is `str`. Since widening it, we'll have to mark `# type: ignore[arg-type]` after `create_batch` on `batches.create` calls
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org