• Home
• airflow.providers.common.ai
• airflow.providers.common.ai.batch
• airflow.providers.common.ai.batch.base

↑↓ Navigate ⏎ Select Esc Close

airflow.providers.common.ai.batch.base

Provider-agnostic types and the adapter contract for @task.llm_batch.

This module must never import a provider SDK (openai, anthropic, …): batch/dispatch.py and the provider-yaml validation check import every registered module, so a top-level SDK import here would make them fail in any environment without that SDK installed.

Attributes

BatchStatus
IN_PROGRESS_STATUSES
TERMINAL_STATUS_MAP

Classes

BatchRequest	One input item for @task.llm_batch.
SubmitResult	Returned by BatchAdapter.submit() once the provider has accepted the batch.
BatchState	The current status of a provider-native batch, as returned by BatchAdapter.get_batch().
RawResultItem	One provider-native per-request result, before output extraction or validation.
ExtractedOutput	The provider-agnostic shape adapters translate their native response into.
BatchAdapter	Adapter contract between the common.ai batch surface and one provider's batch API.

Functions

evaluate_batch_counts(counts)

Reduce a per-status count breakdown to a single terminal reason.

Module Contents

class airflow.providers.common.ai.batch.base.BatchRequest

Bases: TypedDict

One input item for @task.llm_batch.

A bare string "foo" is shorthand for {"prompt": "foo"}; the operator normalizes to this shape before anything else sees the input.

prompt: Required[str]

model: str | None

system_prompt: str | None

max_tokens: int | None

params: dict[str, Any]

class airflow.providers.common.ai.batch.base.SubmitResult

Returned by BatchAdapter.submit() once the provider has accepted the batch.

batch_id: str

provider_input_ref: str | None

airflow.providers.common.ai.batch.base.BatchStatus

class airflow.providers.common.ai.batch.base.BatchState

The current status of a provider-native batch, as returned by BatchAdapter.get_batch().

status: BatchStatus

counts: collections.abc.Mapping[str, int] | None

error_message: str | None

class airflow.providers.common.ai.batch.base.RawResultItem

One provider-native per-request result, before output extraction or validation.

iter_results yields these keyed by position (index), parsed back out of the custom_id each adapter wrote at submit time (f"{key16}-{index}"). raw is opaque here; only the adapter’s own BatchAdapter.extract_output() knows how to read it.

provider_status keeps "expired" and "cancelled" distinct from "errored": both providers report per-item expiry (Anthropic as a result type, OpenAI as an error-file line with code batch_expired), and folding it into "errored" would make an SLA lapse indistinguishable from a rate limit or a bad request.

custom_id: str

index: int

provider_status: Literal['success', 'errored', 'expired', 'cancelled']

model: str | None

usage: collections.abc.Mapping[str, int] | None

finish_reason: str | None

error: collections.abc.Mapping[str, Any] | None

raw: Any

class airflow.providers.common.ai.batch.base.ExtractedOutput

The provider-agnostic shape adapters translate their native response into.

extract_output picks exactly one kind:

• "text": plain-text output (output_type is str); text set.

• "json_text": a JSON document as a string that still needs parsing (e.g. OpenAI’s response_format content); text set.

• "json_value": an already-parsed JSON-native value (e.g. Anthropic’s tool-use input, which arrives as a dict); value set.

• "absent": the model did not produce structured output at all (e.g. no tool_use block came back even though one was requested); text may still carry leftover plain-text content for diagnostics.

kind: Literal['text', 'json_text', 'json_value', 'absent']

text: str | None = None

value: Any | None = None

class airflow.providers.common.ai.batch.base.BatchAdapter(*, api_key=None, base_url=None, client=None)

Bases: abc.ABC

Adapter contract between the common.ai batch surface and one provider’s batch API.

batch/dispatch.py selects a concrete subclass by the model_id prefix (request shape) and checks the connection type against conn_types (auth). Operator, trigger, and the state/results layers depend on this interface only, so another package can register its own adapter (see register_adapter()) without touching them.

name: ClassVar[str]

conn_types: ClassVar[frozenset[str]]

max_requests: ClassVar[int]

max_payload_bytes: ClassVar[int]

allows_per_request_model: ClassVar[bool]

abstract validate_requests(requests, *, model, output_spec, idempotency_key, **kwargs)

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.

abstract submit(requests, *, model, idempotency_key, input_fingerprint, output_spec, **kwargs)

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.

abstract get_batch(batch_id)

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

abstract cancel_batch(batch_id)

Request cancellation of a batch.

abstract iter_results(batch_id)

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.

abstract build_output_directive(spec)

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.

abstract extract_output(raw, spec)

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

abstract find_orphaned_batch(idempotency_key, input_fingerprint, not_before)

Best-effort recovery of a batch whose submit response was never recorded.

Return the provider’s batch id if a batch exists whose recorded idempotency_key and input_fingerprint both match and which was created at or after not_before (an ISO 8601 timestamp, the intent record’s own write time); prefer the most recently created candidate. All three conditions are needed: the key is stable across a clear with different prompts, and an unordered listing could otherwise return an older submission of identical content.

Return None if nothing matches. An adapter with no way to correlate an orphan back to a batch must return None unconditionally and say so in its docstring. Raise if the lookup itself failed (network, auth); the operator treats that as “unknown”, never as “absent”.

close()

Release the underlying SDK client’s connection pool.

The default closes self._client if the adapter set one and it has a close() method, which both the OpenAI and Anthropic clients do. Callers wrap adapter use in try/finally so the pool is freed deterministically rather than by garbage collection.

resolve_request_model(request_model, *, default_bare_model, request_index)

Resolve a per-request model override to this adapter’s bare model name.

request_model must be None (inherit the batch-level default, already bare) or a "<provider>:<model>" string whose prefix equals name. A request naming a different provider is rejected before any network call rather than sent to the wrong provider’s model namespace.

check_custom_id_length(idempotency_key, request_count, *, max_length=64)

Reject a batch whose worst-case custom_id would exceed max_length.

Anthropic documents custom_id as ^[a-zA-Z0-9_-]{1,64}$. A 16-character key plus a separator plus a 5-digit index is at most 22 characters, so this is not reachable today; it exists so a future change to the key length or max_requests fails clearly before submit instead of producing an invalid custom_id.

airflow.providers.common.ai.batch.base.evaluate_batch_counts(counts)

Reduce a per-status count breakdown to a single terminal reason.

Shared by the trigger (provider-reported counts) and the results layer’s manifest assembly (the full breakdown, including invalid_output and missing) so both describe “some requests did not produce a usable result” the same way.

airflow.providers.common.ai.batch.base.IN_PROGRESS_STATUSES

airflow.providers.common.ai.batch.base.TERMINAL_STATUS_MAP: dict[str, str]

Previous Next

Was this entry helpful?