kaxil opened a new pull request, #76974: URL: https://github.com/apache/airflow/pull/76974
`@task.llm_batch` talked to the Anthropic and OpenAI batch APIs through about 860 lines of adapters we wrote ourselves, because pydantic-ai had no batch abstraction when the operator shipped. pydantic-ai-slim 2.62 added one: every `Model` now has public `batch_encode`, `batch_create`, `batch_status`, `batch_results`, `batch_cancel` and `batch_decode` methods. This PR moves the Anthropic adapter onto them and adds Gemini batch (`google:<model>`, Gemini Developer API), which `@task.llm_batch` did not support. OpenAI follows in a stacked PR (#TBD); both should ship in the same provider release. The operator, trigger, state machine and `BatchAdapter` extension point stay as they are. `register_adapter()` and the `airflow.providers.common.ai.batch_adapters` entry point keep working, and a third-party adapter written for the old contract still runs because the new members have defaults. ## Design rationale **Why not the `BatchModelRequests` capability?** It pairs with `PendingModelResponse` to pool requests from many agent runs and resume a run from its message history. This operator owns one batch per task, with its own state file and trigger, and using the capability here would break two things: - A pool only cancels batches it submitted itself. A batch recorded in a message history from another process "is left running" ([`BatchPool.abandon`](https://github.com/pydantic/pydantic-ai/blob/v2.66.0/pydantic_ai_slim/pydantic_ai/models/_batch_model.py#L513-L530)), so `on_kill` and `cancel_on_timeout` would stop cancelling anything after a deferral. - The pending batch handle sits in a private `_Ticket`, which the trigger would have to poll through. Its tool-loop and output-retry features would also go unused, since the operator takes no toolsets. All the trigger does per poll is call `batch_status` on the stored handle, with a model rebuilt from the connection; no agent, prompt or tool reaches the triggerer. pydantic-ai checks a batch's group on every call, and Gemini (OpenAI too) puts the model in the group. So poll, collect and cancel build the model from the stored handle's `model_name` rather than the task's current `model_id`; otherwise changing a task's model would strand a batch in flight, or leave it running when the operator tries to cancel it. For structured output the adapter sets the mode itself. pydantic-ai's profile default for these models is `tool` mode, and with no output tool that sends no schema at all. Native JSON-schema output is used where the model supports it, a single forced output tool otherwise. `request_params` keeps its meaning. Typed sampling keys and `anthropic_*` settings become `ModelSettings`, and any other key goes through `extra_body`, which pydantic-ai merges into the Messages body. Because `extra_body` is merged last, keys the operator manages (`model`, `messages`, ...) are stripped from it with a warning. Gemini has no raw passthrough, so unknown keys fail before submission there. There is one private import: Gemini reports "this job wrote no results file" with `BatchJobError` from `pydantic_ai.models._batch_codec`. Nothing public tells a cancelled job with nothing to collect apart from a transient download error, and the second must propagate so the task's retry collects the results. A test runs pydantic-ai's real Gemini collect path, so a change on their side fails CI here. ## User-visible changes From the changelog entry: - `pydantic-ai-slim` floor rises to 2.66.0, `anthropic` to 1.15.0, and the `google` extra gains `google-genai`. 2.62.0 has the same `batch_*` signatures, but only 2.66.0 is tested. - A setting the model would drop silently (`seed` on Anthropic) now fails before submission instead. - Structured output uses native JSON-schema output instead of a forced tool call, except on models without it. - Results are read as one list per batch (that is pydantic-ai's API), so landing holds a batch's response bodies in memory at once. - An Anthropic batch submitted by the previous version and still running re-attaches after the upgrade. A trigger deferred by the previous version fails its poll with a message naming the cause, and the task's retry re-attaches. ## Known issues - An outdated batch recorded by the previous version under a different connection is not cancelled automatically; the task log says so. - `validate_requests` and `submit` each encode every request. I tried reusing validate's encoding in submit and backed it out: the encoded batch (up to 256 MB on Anthropic) would sit on the adapter across the state write, to save CPU work that makes no network call. ## Evidence Run against the real Anthropic Message Batches API on a local Airflow 3.5 stack (deferrable, Haiku 4.5 with one request on Sonnet 4.6), 3-4 requests per batch: | Scenario | Result | |---|---| | Structured output (`output_type=Sentiment`) | 4/4 rows parsed into the model; each batch took 1.5-2.5 min. The request carried `output_config: {"format": {"type": "json_schema", ...}}`. The triggerer logged a `status=in_progress` poll every 30 s, then `status=completed`. | | `@task.llm_batch` with a per-request model override | One batch; the overridden row was answered by `claude-sonnet-4-6`, the rest by Haiku. `temperature` (typed) and `metadata.user_id` (raw Messages field) were both accepted. A prompt containing `{{ x }}` came back as `"{{ x }}"`. | | Clear a deferred task | Try 2 logged `Re-attaching to batch msgbatch_... the input is unchanged, so nothing is resubmitted.` and landed the 3 results, with no second batch. | | `timeout=20, cancel_on_timeout=True` | `Cancelled batch msgbatch_...: timeout=20s exceeded`, then `LLMBatchTimeoutError`; the provider shows `cancel_initiated_at` set. Clearing the task re-attached to the same batch and landed what had finished. | | Mark a deferred task failed | `Trigger cancelled by user action, invoking on_kill`, then `on_kill: cancelled batch msgbatch_...`; the provider shows the cancel. | The structured-output run also passed on the OpenAI PR's branch, which carries this one. <!-- screenshots and the re-attach GIF are attached after the PR is created --> The full Common AI suite passes (4791 passed, 27 skipped) and mypy is clean on the changed modules. -- 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
