Add a deferrable batch execution mode to common.ai

[Lee-W]

Provider batch APIs run at roughly half price with a 24h SLA, but a pipeline built on @task.llm had no way to use them. Switching meant the author owned JSONL construction, upload, chunking to provider limits, polling, and result retrieval by hand.

pydantic-ai exposes no batch abstraction, so there is no Agent path for batch work. Batch is submit/poll/fetch spanning hours rather than request/response in a loop, so the agent loop, tool execution, HITL approval and message_history have no meaning there. A dedicated deferrable operator and decorator fit better than a mode flag on @task.llm.

Summary

LLMBatchOperator / @task.llm_batch submit a list of prompts as one OpenAI or Anthropic batch job, defer while it runs, and land results as JSONL on object storage. The XCom value is a small manifest (counts per outcome, result URI, provenance); a batch can hold 100k items, so results never go through XCom. Structured output goes through a provider-neutral JSON Schema (OpenAI response_format, Anthropic forced tool call), and a response that does not match lands as an invalid_output row rather than failing the task.

Retry safety is the cost-critical part. A retry or clear computes the same identity key (dag, task, run, map index, never try number) and re-attaches to the recorded batch when the input fingerprint still matches. The state lives on the same object storage as the results, because the server clears a task instance's XCom at the start of every attempt. Intent is recorded before the paid submit call, so a crash between "request sent" and "response recorded" is recoverable on OpenAI through batch metadata; Anthropic offers no such lookup and falls through to an explicit on_orphaned_intent policy.

Design rationale

A cancelled batch fails the task instead of landing as a success. On Airflow 3.3+ a UI clear runs the trigger's on_kill, which cancels the batch under the default cancel_on_kill=True. If the next attempt re-attached and treated cancelled like completed, a clear would produce a green task with a manifest of mostly missing rows. Whatever finished is landed, the task raises LLMBatchCancelledError, and the state is cleared so the next attempt submits fresh. A provider-side failed (OpenAI rejected the input before running anything) clears the state for the same reason. Every other outcome keeps the state, so a clear after success re-lands the same results at no cost.

The vendor SDKs are called directly rather than through the vendor providers' hooks. OpenAIHook has no batch listing or file download, and a dependency from common.ai on apache-airflow-providers-openai would drag its Airflow 2 floor along. The cost is that this package tracks the openai and anthropic SDK versions in parallel with those providers; the pins match theirs (openai>=2.45.0, the floor pydantic-ai-slim[openai] already imposes, and anthropic>=1.0.0).

Adapters are pluggable. BatchAdapter is the contract; dispatch resolves the model_id prefix through the built-in table, register_adapter(), or an entry point in the airflow.providers.common.ai.batch_adapters group, and each adapter declares the connection types it can authenticate from. Bedrock and Vertex batch (S3/GCS based, different SDK shapes) can ship inside their own provider packages without touching this one.

OpenAI requests send max_completion_tokens. max_tokens is deprecated and rejected by reasoning models such as gpt-5, which is the model the example uses; with the default fail_on_partial_error=False an all-rejected batch would otherwise complete green with zero usable rows.

@task.llm_batch does not Jinja-render the returned prompts, unlike @task.llm. Batch inputs are bulk text the Dag author did not write, where a stray {{ or {% would either fail the whole batch or resolve var/conn accessors against Airflow secrets. Anything dynamic belongs in the callable, which receives the task context.

Gotchas

• cancel_on_kill in deferrable mode needs Airflow 3.3+, the first version whose triggerer calls a trigger's on_kill. On 3.0 to 3.2 a killed deferred task's batch keeps running and a clear re-attaches to it.
• cancel_on_timeout=False leaves the batch running and billing; a retry after the original budget elapsed gets a fresh timeout measured from the retry, so retries and retry_delay bound the total wait.
• Azure OpenAI, Bedrock and Vertex connection types are rejected with a message pointing at the two routes that work today: a pydanticai connection on the public endpoint, or an OpenAI-compatible gateway (LiteLLM exposes /v1/files and /v1/batches and routes to those backends behind it).

---

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)

---

• Read the 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) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

[kaxil]

End-to-end runs against the real OpenAI Batch API and the real Anthropic Message Batches API, on Airflow main in breeze with the deferrable path (worker submit, triggerer poll, worker landing). Same Dag shape for both providers: three batch tasks plus a downstream task that reads the JSONL rows back through ObjectStoragePath.

OpenAI

Task	Shape	Proves
classify_reviews	LLMBatchOperator, gpt-4.1-mini, output_type=Sentiment (Pydantic), request_params={"temperature": 0, "user": ...}	Structured output via response_format, OpenAI body params pass through
summarize_reviews	@task.llm_batch, model taken from the connection's Model field, per-request system_prompt / max_tokens / params overrides	Decorator path, per-request overrides (one request answered in French as instructed)
extract_keywords	gpt-5-mini, output_type=list[str], request_params={"reasoning_effort": "low"}, fail_on_partial_error=True	Reasoning-model parameter, max_completion_tokens translation, non-object schema wrapping/unwrapping

All 12 requests came back success; every manifest reconciled (request_count == sum(counts), terminal_reason: succeeded).

Task log for the operator: submit on the worker, one deferral, trigger fires, results landed with per-bucket counts.

XCom for the gpt-5-mini task: the batch_id key pushed at submit time and the manifest as the return value.

Clearing a finished task re-attaches to the recorded batch and re-lands the same file in a few seconds, with no new submission:

Sample rows from the landed JSONL (gpt-5-mini, list[str] output):

{"custom_id": "4fdf6648f4efda69-0", "index": 0, "status": "success", "output": ["loved", "would buy again", "positive"], "raw_output": null, "error": null, "model": "gpt-5-mini-2025-08-07", "usage": {"input_tokens": 61, "output_tokens": 89}, "finish_reason": "stop"}
{"custom_id": "4fdf6648f4efda69-1", "index": 1, "status": "success", "output": ["broke", "one-use", "disappointed"], "raw_output": null, "error": null, "model": "gpt-5-mini-2025-08-07", "usage": {"input_tokens": 62, "output_tokens": 89}, "finish_reason": "stop"}

Anthropic

Task	Shape	Proves
classify_reviews	LLMBatchOperator, claude-haiku-4-5, output_type=Sentiment, request_params={"temperature": 0.2, "metadata": {"user_id": ...}}	Structured output via a forced tool (finish_reason: tool_use), Anthropic Messages body params pass through
summarize_reviews	@task.llm_batch, model from the connection, one request overriding model to claude-sonnet-4-5 plus system_prompt / max_tokens / params overrides	Decorator path, per-request model override inside one batch (Anthropic allows it, OpenAI does not), French override honoured
extract_keywords	claude-haiku-4-5, output_type=list[str], fail_on_partial_error=True	Non-object schema wrapping/unwrapping through the tool input

All 12 requests success, manifests reconciled. Anthropic results arrive out of order; rows carry their index and the manifest says ordered: false, rejoin_key: index.

Triggerer polling and landing for the mixed-model decorator task:

Rows from the mixed-model batch, showing the per-request model override landing on Sonnet while the rest ran on Haiku:

{"index": 0, "status": "success", "output": "The reviewer highly enjoyed the product and plans to purchase it again in the future.", "model": "claude-haiku-4-5-20251001", "usage": {"input_tokens": 30, "output_tokens": 19}, "finish_reason": "end_turn"}
{"index": 1, "status": "success", "output": "Cet article s'est cassé après une seule utilisation, ce qui est très décevant.", "model": "claude-haiku-4-5-20251001", "usage": {"input_tokens": 33, "output_tokens": 27}, "finish_reason": "end_turn"}
{"index": 2, "status": "success", "output": "Product works as advertised but is unremarkable.", "model": "claude-sonnet-4-5-20250929", "usage": {"input_tokens": 32, "output_tokens": 14}, "finish_reason": "end_turn"}

Timing

Provider	Requests	Submit to landed
OpenAI	3 to 6	1 to 4 minutes
Anthropic	3	1 minute (keywords), 19 minutes (mixed-model summaries)
Anthropic	6	21 minutes

Changes that came out of the runs

Both are in the latest commit. The operator now logs a line when it re-attaches to a recorded batch and another when it lands results; before this a successful landing was silent in the task log. The operator guide gained a "Provider-specific parameters" section with a per-provider translation table for system_prompt, max_tokens, output_type, request_params and completion_window, and the example Dags were rewritten into four: basic operator with a downstream reader, decorator, provider params for both providers side by side, and per-request overrides.

One thing to note from the Anthropic log above: every status poll shows up as an INFO line (HTTP Request: GET .../batches/...), so a long batch fills the triggerer log with one line per poll_interval. That line is httpx's own request log at INFO; the OpenAI task logs above do not show it. Not changed here; worth a look as a follow-up.