airflow.providers.common.ai.batch.openai

OpenAI batch adapter.

Talks to the openai SDK directly rather than through apache-airflow-providers-openai: that provider’s hook has no batch listing or file-download methods, so reusing it would add a hard dependency without the capabilities this needs. The SDK import is deferred to _build_client() so importing this module never requires openai.

Attributes

log

Classes

OpenAIBatchAdapter

Batch adapter for OpenAI's Batch API (/v1/chat/completions).

Module Contents

airflow.providers.common.ai.batch.openai.log[source]
class airflow.providers.common.ai.batch.openai.OpenAIBatchAdapter(*, api_key=None, base_url=None, client=None)[source]

Bases: airflow.providers.common.ai.batch.base.BatchAdapter

Batch adapter for OpenAI’s Batch API (/v1/chat/completions).

Parameters:

  • api_key (str | None) – Passed straight to the openai.OpenAI client. None falls back to the SDK’s own env-var resolution (OPENAI_API_KEY).

  • base_url (str | None) – Passed straight to the openai.OpenAI client. Pointing it at an OpenAI-compatible gateway that exposes /v1/files and /v1/batches routes the batch through that gateway.

  • client (Any | None) – Inject a pre-built client (or a test double) instead of constructing one from api_key/base_url.

name = 'openai'[source]
max_requests = 50000[source]
max_payload_bytes = 200000000[source]
allows_per_request_model = False[source]
build_output_directive(spec)[source]

Translate spec.json_schema into this provider’s structured-output request fields.

Returns {} when spec.is_structured is False. The result is merged into every request’s body/params by the adapter itself.

validate_requests(requests, *, model, output_spec, idempotency_key, system_prompt='', max_tokens=1024, request_params=None, **kwargs)[source]

Reject an over-limit or otherwise invalid batch before any network call.

Must check len(requests) against max_requests and the serialized request size (including the per-request output directive built from output_spec, which a large schema repeats in every request) against max_payload_bytes, and call check_custom_id_length(). Raise a subclass of LLMBatchInputError naming the limit, the actual count/size, and the .expand() remedy.

**kwargs carries the same batch-level request-building context as submit() (system_prompt, max_tokens, request_params) so the serialized-size estimate matches what submit sends.

submit(requests, *, model, idempotency_key, input_fingerprint, output_spec, system_prompt='', max_tokens=1024, request_params=None, completion_window='24h', **kwargs)[source]

Upload/submit the batch and return the provider’s batch id.

Where the provider offers batch-level metadata (OpenAI), record both idempotency_key and input_fingerprint there so find_orphaned_batch() can read them back. The key alone identifies the task instance, not this submission of it: it is stable across a clear even when the prompts changed.

get_batch(batch_id)[source]

Return the current status of a submitted batch. Synchronous; the trigger wraps it in to_thread.

cancel_batch(batch_id)[source]

Request cancellation of a batch.

find_orphaned_batch(idempotency_key, input_fingerprint, not_before)[source]

Look for a batch whose metadata carries both identifiers, created at or after not_before.

submit() records idempotency_key and input_fingerprint in the batch’s metadata, so a batch whose submit response never made it back can be found and re-attached to. No server-side metadata filter exists, so this walks up to _ORPHAN_SCAN_LIMIT batches client-side; among every match the most recently created one wins.

iter_results(batch_id)[source]

Return (not yield) a streaming iterator of per-request results.

Returning a plain iterator rather than using yield in this method makes an invalid batch_id (or any other call-time failure) raise immediately, instead of only once the caller starts iterating.

extract_output(raw, spec)[source]

Pull the model’s output out of a provider-native result, in the shape ExtractedOutput describes.

Was this entry helpful?