kaxil commented on code in PR #69003: URL: https://github.com/apache/airflow/pull/69003#discussion_r3520421592
########## providers/anthropic/src/airflow/providers/anthropic/triggers/batch.py: ########## @@ -0,0 +1,113 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import asyncio +import time +from collections.abc import AsyncIterator +from typing import Any + +from airflow.providers.anthropic.hooks.anthropic import AnthropicHook, BatchStatus +from airflow.triggers.base import BaseTrigger, TriggerEvent + +#: Consecutive failed polls tolerated before the trigger gives up (transient errors). +MAX_CONSECUTIVE_POLL_FAILURES = 5 + + +class AnthropicBatchTrigger(BaseTrigger): + """ + Poll an Anthropic Message Batch until it reaches the terminal ``ended`` status. + + :param conn_id: The Anthropic connection ID. + :param batch_id: The batch to poll. + :param poll_interval: Seconds to sleep between polls. + :param end_time: Wall-clock deadline (``time.time()`` epoch seconds) after which a + ``timeout`` event is emitted. Wall-clock is used deliberately: the trigger is + serialized to the metadata DB and may resume in a different triggerer process, + so a per-process ``time.monotonic()`` value would not survive serialization. + """ + + def __init__(self, conn_id: str, batch_id: str, poll_interval: float, end_time: float) -> None: + super().__init__() + self.conn_id = conn_id + self.batch_id = batch_id + self.poll_interval = poll_interval + self.end_time = end_time + + def serialize(self) -> tuple[str, dict[str, Any]]: + """Serialize AnthropicBatchTrigger arguments and class path.""" + return ( + "airflow.providers.anthropic.triggers.anthropic.AnthropicBatchTrigger", + { + "conn_id": self.conn_id, + "batch_id": self.batch_id, + "poll_interval": self.poll_interval, + "end_time": self.end_time, + }, + ) + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Poll the batch status and yield exactly one terminal event.""" + hook = AnthropicHook(conn_id=self.conn_id) + consecutive_failures = 0 + while True: + try: + # get_batch is a blocking SDK HTTP call; run it off the event loop so a + # single poll does not stall every other trigger on this triggerer. + batch = await asyncio.to_thread(hook.get_batch, self.batch_id) + except Exception as e: + # Tolerate transient polling errors rather than failing a (up to 24h) wait. + consecutive_failures += 1 + if consecutive_failures >= MAX_CONSECUTIVE_POLL_FAILURES or time.time() > self.end_time: + yield TriggerEvent({"status": "error", "batch_id": self.batch_id, "message": str(e)}) + return + self.log.warning("Polling batch %s failed (%s); retrying.", self.batch_id, e) + await asyncio.sleep(self.poll_interval) + continue + + consecutive_failures = 0 + self.log.debug("Batch %s status=%s", self.batch_id, batch.processing_status) + if not BatchStatus.is_in_progress(batch.processing_status): + counts = batch.request_counts + yield TriggerEvent( + { + "status": "success", + "batch_id": self.batch_id, + "message": f"Batch {self.batch_id} has ended.", + "request_counts": { + "succeeded": counts.succeeded, + "errored": counts.errored, + "canceled": counts.canceled, + "expired": counts.expired, + "processing": counts.processing, + }, + } + ) + return + if time.time() > self.end_time: + yield TriggerEvent( + { + "status": "timeout", + "batch_id": self.batch_id, + "message": ( + f"Batch {self.batch_id} did not reach a terminal status " + "before the configured timeout." + ), + } + ) + return + await asyncio.sleep(self.poll_interval) Review Comment: Done in `8a72eb5`. Added async `on_kill` to both triggers (batch cancels the batch, agent archives the session). It's a no-op override on Airflow <3.3.0 (where `BaseTrigger.on_kill` doesn't exist) and active on 3.3+, so the provider keeps its 3.0 floor; the operator docstrings note that older Airflow needs manual cleanup for a killed deferred task. ########## providers/anthropic/src/airflow/providers/anthropic/operators/agent.py: ########## @@ -0,0 +1,219 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import time +from collections.abc import Sequence +from datetime import timedelta +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.anthropic.exceptions import AnthropicAgentSessionError, AnthropicAgentSessionTimeout +from airflow.providers.anthropic.hooks.anthropic import AnthropicHook +from airflow.providers.anthropic.triggers.agent import AnthropicAgentSessionTrigger +from airflow.providers.common.compat.sdk import BaseOperator, conf + +if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context + + +class AnthropicAgentSessionOperator(BaseOperator): + """ + Run a Managed Agents session against a pre-created agent and environment. + + Anthropic runs the agent loop server-side; the worker creates a session, sends the + initial instruction, and waits for the session to reach a terminal status. In + deferrable mode it releases the worker slot while a trigger polls the session status. + + Provide exactly one of ``message`` (a single user turn) or ``outcome`` (a + ``user.define_outcome`` rubric that the agent iterates against until satisfied). + + .. important:: + This operator is for **autonomous** agents — configure the agent without + client-side custom tools or an ``always_ask`` permission policy. + + Completion is detected accurately for both modes. A ``message`` run reads the + terminal ``session.status_idle`` event's ``stop_reason`` (correlated against the + kickoff event to avoid a start-race false positive): ``end_turn`` succeeds, while + ``requires_action`` (the agent is blocked on input) and ``retries_exhausted`` + raise an error rather than silently passing. An ``outcome`` run is judged from the + session's ``outcome_evaluations`` verdict (``satisfied`` vs. + ``failed``/``max_iterations_reached``/``interrupted``). + + Agents and environments are created once (see + :meth:`~airflow.providers.anthropic.hooks.anthropic.AnthropicHook.create_agent`), + not per task run. + + Outputs the agent writes to ``/mnt/session/outputs/`` are retrieved afterwards via the + Files API (``scope_id=<session_id>``); the operator returns the **session ID only**. + + .. seealso:: + For more information, take a look at the guide: + :ref:`howto/operator:AnthropicAgentSessionOperator` + + :param agent_id: ID of a pre-created agent. + :param environment_id: ID of a pre-created environment. + :param message: A single user message to start the session. Mutually exclusive with ``outcome``. + :param outcome: A ``user.define_outcome`` payload (``description``, ``rubric``, + optional ``max_iterations``). Mutually exclusive with ``message``. + :param conn_id: The Anthropic connection ID to use. + :param deferrable: Run the operator in deferrable mode. + :param poll_interval: Seconds between session status checks (both paths). + :param timeout: Seconds to wait for a terminal status. Defaults to 24 hours. In + deferrable mode the trigger enforces this and tears the session down on timeout; + set ``execution_timeout`` only for a shorter hard cap (which preempts that + graceful teardown). + :param vault_ids: Vault IDs providing MCP/credential access to the session. + :param session_resources: Session resources (files, GitHub repos, memory stores). Named + ``session_resources`` to avoid colliding with the reserved ``BaseOperator.resources``; + forwarded to ``sessions.create`` as ``resources``. + :param session_kwargs: Extra keyword arguments forwarded to ``sessions.create``. + """ + + template_fields: Sequence[str] = ("agent_id", "environment_id", "message", "outcome") + + def __init__( + self, + *, + agent_id: str, + environment_id: str, + message: str | None = None, + outcome: dict[str, Any] | None = None, + conn_id: str = AnthropicHook.default_conn_name, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + poll_interval: float = 30, + timeout: float = 24 * 60 * 60, + vault_ids: list[str] | None = None, + session_resources: list[dict[str, Any]] | None = None, + session_kwargs: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + if (message is None) == (outcome is None): + raise ValueError("Provide exactly one of 'message' or 'outcome'.") + if outcome is not None and "rubric" not in outcome: + raise ValueError("'outcome' must include a 'rubric' (with 'description').") Review Comment: Good catch. The operator now requires both `description` and `rubric` in the outcome (the old message mentioned `description` but only checked `rubric`). Done in `8a72eb5`. ########## providers/anthropic/src/airflow/providers/anthropic/operators/batch.py: ########## @@ -0,0 +1,200 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import time +from collections.abc import Sequence +from datetime import timedelta +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.anthropic.exceptions import AnthropicBatchJobError, AnthropicBatchTimeout +from airflow.providers.anthropic.hooks.anthropic import AnthropicHook, evaluate_batch_counts +from airflow.providers.anthropic.triggers.anthropic import AnthropicBatchTrigger +from airflow.providers.common.compat.sdk import BaseOperator, conf + +if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context + + Review Comment: Agreed, done in `8a72eb5`: the `operators`/`sensors`/`triggers` `anthropic.py` modules are now `batch.py` (so it's `agent.py` + `batch.py`), with the tests renamed to match. `hooks.anthropic` stays since it's shared by both surfaces. ########## providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py: ########## @@ -0,0 +1,569 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import logging +import time +from enum import Enum +from functools import cached_property +from typing import TYPE_CHECKING, Any, cast + +from anthropic import ( + Anthropic, + AnthropicAWS, + AnthropicBedrock, + AnthropicFoundry, + AnthropicVertex, + IdentityTokenFile, + WorkloadIdentityCredentials, +) + +from airflow.providers.anthropic.exceptions import ( + AnthropicAgentSessionError, + AnthropicAgentSessionTimeout, + AnthropicBatchJobError, + AnthropicBatchTimeout, + AnthropicError, +) +from airflow.providers.common.compat.sdk import AirflowSkipException, BaseHook + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator + + from anthropic.types import Message + from anthropic.types.beta import ( + BetaEnvironment, + BetaManagedAgentsAgent, + BetaManagedAgentsSession, + environment_create_params, + ) + from anthropic.types.beta.sessions import BetaManagedAgentsEventParams + from anthropic.types.messages import MessageBatch, MessageBatchIndividualResponse + from anthropic.types.messages.batch_create_params import Request + +#: Default model used when an operator or hook caller does not specify one. +#: Prefer configuring the model on the connection so it can be updated without +#: a provider release when this model ID is retired. +DEFAULT_MODEL = "claude-opus-4-8" + +#: Platforms that serve the first-party-only endpoints (Message Batches, token +#: counting, the Models API). Amazon Bedrock, Google Vertex AI and Microsoft +#: Foundry do not serve these, so the hook fails fast rather than surfacing a +#: raw ``404`` from the SDK. +FIRST_PARTY_PLATFORMS = frozenset({"anthropic", "aws"}) + +AnthropicClient = Anthropic | AnthropicBedrock | AnthropicVertex | AnthropicAWS | AnthropicFoundry + + +class BatchStatus(str, Enum): + """Top-level ``processing_status`` of an Anthropic Message Batch.""" + + IN_PROGRESS = "in_progress" + CANCELING = "canceling" + ENDED = "ended" + + @classmethod + def is_in_progress(cls, status: str) -> bool: + """Return ``True`` while the batch has not reached the terminal ``ended`` status.""" + return status != cls.ENDED + + +class SessionStatus(str, Enum): + """Status of a Managed Agents session.""" + + RESCHEDULING = "rescheduling" + RUNNING = "running" + IDLE = "idle" + TERMINATED = "terminated" + + @classmethod + def is_terminal(cls, status: str) -> bool: + """ + Return ``True`` once the session has stopped working. + + ``idle`` means the agent finished its turn (done, for an autonomous run); + ``terminated`` is an unrecoverable failure. Both stop the wait. + """ + return status in (cls.IDLE, cls.TERMINATED) + + +#: ``outcome_evaluations[].result`` values that mean the outcome did NOT succeed. +OUTCOME_FAILURE_RESULTS = frozenset({"failed", "max_iterations_reached", "interrupted"}) + + +def evaluate_session_state( + session: BetaManagedAgentsSession, *, expect_outcome: bool +) -> tuple[bool, str | None, bool]: + """ + Judge a polled session from its object fields alone. + + Returns ``(done, error_message, needs_event_check)``. ``done=False`` means keep + polling. ``needs_event_check=True`` means the session is ``idle`` on a ``message`` + run and the object can't say *why* — the caller must inspect the event log (see + :meth:`AnthropicHook.poll_session_completion`). + + The ``status`` field can't distinguish a genuine ``end_turn`` from ``requires_action`` + or ``retries_exhausted``, nor a just-created ``idle``. For an outcome run the true + verdict is in ``outcome_evaluations`` (judged here, which also defeats the start race). + """ + if session.status == SessionStatus.TERMINATED: + return True, f"Session {session.id} terminated.", False + if session.status != SessionStatus.IDLE: + return False, None, False + if not expect_outcome: + return False, None, True + for evaluation in session.outcome_evaluations: + if evaluation.result == "satisfied": + return True, None, False + if evaluation.result in OUTCOME_FAILURE_RESULTS: + return True, f"Outcome not satisfied for session {session.id}: {evaluation.result}.", False + # idle but no terminal outcome verdict yet (e.g. the run has not started) + return False, None, False + + +def evaluate_batch_counts( + *, + batch_id: str | None, + canceled: int, + errored: int, + expired: int, + succeeded: int, + fail_on_partial_error: bool, +) -> None: + """ + Apply the success/skip/fail policy for a terminal batch's request counts. + + Lives in the hook module so both :class:`AnthropicBatchOperator` and + :class:`~airflow.providers.anthropic.sensors.anthropic.AnthropicBatchSensor` share it + without an operator/sensor cross-import. Raises ``AirflowSkipException`` for a + fully-cancelled batch, ``AnthropicBatchJobError`` when ``fail_on_partial_error`` and any + request failed, otherwise returns (logging a warning for partial failures). + """ + total = canceled + errored + expired + succeeded + if total and canceled == total: + raise AirflowSkipException(f"Batch {batch_id} was fully cancelled.") + failed = errored + expired + if failed: + message = ( + f"Batch {batch_id} ended with {failed} failed request(s) " + f"(errored={errored}, expired={expired}, succeeded={succeeded})." + ) + if fail_on_partial_error: + raise AnthropicBatchJobError(message) + logger.warning("%s Successful results are still available.", message) + + +class AnthropicHook(BaseHook): + """ + Use the Anthropic SDK to interact with the Claude API. + + The connection's ``password`` is used as the API key and ``host`` as an optional + base URL (for gateways/proxies). The ``extra`` field selects the platform client + and passes platform-specific configuration: + + - ``platform``: one of ``anthropic`` (default), ``bedrock``, ``vertex``, ``aws``, ``foundry``. + - ``model``: default model id used when an operator/hook call omits ``model`` (lets you + change the model without editing Dags); falls back to :data:`DEFAULT_MODEL`. + - ``aws_region``: region for the ``bedrock`` platform. + - ``project_id`` / ``region``: project and region for the ``vertex`` platform. + - ``resource``: Azure resource name for the ``foundry`` platform. + - ``anthropic_client_kwargs``: extra keyword arguments forwarded to the client + constructor (e.g. ``timeout``, ``max_retries``, ``default_headers``). + - ``workload_identity``: configure `Workload Identity Federation + <https://platform.claude.com/docs/en/manage-claude/workload-identity-federation>`__ + (keyless OIDC auth) with ``identity_token_file``, ``federation_rule_id``, + ``organization_id``, ``service_account_id`` and optional ``workspace_id`` / ``scope``. + + When the ``anthropic`` platform has no API Key and no ``workload_identity`` block, the + client is built with no static credential so the SDK resolves them from the environment + — supporting env-driven Workload Identity Federation and ``ant`` profiles. + + .. seealso:: https://docs.claude.com/en/api/client-sdks + + :param conn_id: :ref:`Anthropic connection id <howto/connection:anthropic>`. + """ + + conn_name_attr = "conn_id" + default_conn_name = "anthropic_default" + conn_type = "anthropic" + hook_name = "Anthropic" + + def __init__(self, conn_id: str = default_conn_name, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.conn_id = conn_id + + @cached_property + def _connection(self): + return self.get_connection(self.conn_id) + + @cached_property + def platform(self) -> str: + """Return the configured platform (defaults to ``anthropic``).""" + return (self._connection.extra_dejson.get("platform") or "anthropic").lower() + + @cached_property + def default_model(self) -> str: + """Default model id — connection ``extra['model']`` if set, else :data:`DEFAULT_MODEL`.""" + return self._connection.extra_dejson.get("model") or DEFAULT_MODEL + + @cached_property + def conn(self) -> AnthropicClient: + """Return the Anthropic client for the configured platform.""" + return self.get_conn() + + def get_conn(self) -> AnthropicClient: + """Build and return the Anthropic client for the configured platform.""" + conn = self._connection + extras = conn.extra_dejson + client_kwargs = dict(extras.get("anthropic_client_kwargs", {})) + platform = self.platform + self.log.debug("Building Anthropic client for platform %r (conn_id=%s)", platform, self.conn_id) + if platform == "bedrock": + return AnthropicBedrock(aws_region=extras.get("aws_region"), **client_kwargs) + if platform == "vertex": + return AnthropicVertex( + project_id=extras.get("project_id"), region=extras.get("region"), **client_kwargs + ) + if platform == "aws": + return AnthropicAWS(**client_kwargs) Review Comment: Applied in `8a72eb5`, thanks. ########## providers/anthropic/src/airflow/providers/anthropic/operators/batch.py: ########## @@ -0,0 +1,200 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import time +from collections.abc import Sequence +from datetime import timedelta +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.anthropic.exceptions import AnthropicBatchJobError, AnthropicBatchTimeout +from airflow.providers.anthropic.hooks.anthropic import AnthropicHook, evaluate_batch_counts +from airflow.providers.anthropic.triggers.anthropic import AnthropicBatchTrigger +from airflow.providers.common.compat.sdk import BaseOperator, conf + +if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context + + +class AnthropicBatchOperator(BaseOperator): + """ + Submit an Anthropic Message Batch and wait for it to complete. + + Message Batches process many ``messages.create`` requests asynchronously at 50% of + standard cost; most complete within an hour (24h SLA). This operator submits the + batch and, in deferrable mode, releases the worker slot while a trigger polls for + completion. + + The operator returns the **batch ID only** — never the results. Pull results with + :meth:`~airflow.providers.anthropic.hooks.anthropic.AnthropicHook.stream_batch_results` + and persist them to object storage; results can be very large and must not be pushed + to XCom. Results are retained for 29 days after the batch is created. + + .. note:: + A retry re-submits a brand-new batch. Prefer ``retries=0`` on this task (the + submitted ``batch_id`` is pushed to XCom under key ``batch_id`` immediately, so + a crashed run never loses track of an in-flight batch). + + .. seealso:: + For more information, take a look at the guide: + :ref:`howto/operator:AnthropicBatchOperator` + + :param requests: A list of ``{"custom_id": str, "params": {...}}`` dicts, where + ``params`` is a ``messages.create`` payload (``model``, ``max_tokens``, ``messages``, ...). + :param conn_id: The Anthropic connection ID to use. + :param deferrable: Run the operator in deferrable mode. + :param poll_interval: Seconds between status checks, in both the synchronous and + deferrable paths. + :param timeout: Seconds to wait for the batch to reach a terminal status. Defaults to + 24 hours (the Message Batches SLA). In deferrable mode this also bounds the + deferral; set ``execution_timeout`` only if you want a shorter hard cap (note a + shorter ``execution_timeout`` preempts the graceful cancel-on-timeout path). + :param wait_for_completion: Whether to wait for the batch to complete. If ``False``, + the operator returns the batch ID immediately after submission. + :param fail_on_partial_error: If ``True``, fail the task when any request errored or + expired. Defaults to ``False`` (succeed and log a warning so the successful + results are not discarded). + """ + + template_fields: Sequence[str] = ("requests",) + + def __init__( + self, + requests: list[dict[str, Any]], + conn_id: str = AnthropicHook.default_conn_name, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + poll_interval: float = 60, + timeout: float = 24 * 60 * 60, + wait_for_completion: bool = True, + fail_on_partial_error: bool = False, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.requests = requests + self.conn_id = conn_id + self.deferrable = deferrable + self.poll_interval = poll_interval + self.timeout = timeout + self.wait_for_completion = wait_for_completion + self.fail_on_partial_error = fail_on_partial_error + self.batch_id: str | None = None + + @cached_property + def hook(self) -> AnthropicHook: + """Return an instance of the AnthropicHook.""" + return AnthropicHook(conn_id=self.conn_id) + + def execute(self, context: Context) -> str | None: + if not self.requests: + raise ValueError("AnthropicBatchOperator requires at least one request; got an empty list.") + batch = self.hook.create_batch(self.requests) Review Comment: Thanks for the detailed dig. Agreed it's non-blocking: the Message Batches API has no documented server-side idempotency today, so neither form can be relied on for dedup. I implemented the mitigation you landed on: best-effort cancellation on every failure path after submission (the deferred `error` event and any synchronous error, not just timeout), so a task retry is cancel-old-then-submit-new rather than two concurrently billing batches. A Databricks-style `idempotency_token` param is a good follow-up if the API adds support. ########## providers/anthropic/src/airflow/providers/anthropic/operators/batch.py: ########## @@ -0,0 +1,200 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import time +from collections.abc import Sequence +from datetime import timedelta +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.anthropic.exceptions import AnthropicBatchJobError, AnthropicBatchTimeout +from airflow.providers.anthropic.hooks.anthropic import AnthropicHook, evaluate_batch_counts +from airflow.providers.anthropic.triggers.anthropic import AnthropicBatchTrigger +from airflow.providers.common.compat.sdk import BaseOperator, conf + +if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context + + +class AnthropicBatchOperator(BaseOperator): + """ + Submit an Anthropic Message Batch and wait for it to complete. + + Message Batches process many ``messages.create`` requests asynchronously at 50% of + standard cost; most complete within an hour (24h SLA). This operator submits the + batch and, in deferrable mode, releases the worker slot while a trigger polls for + completion. + + The operator returns the **batch ID only** — never the results. Pull results with + :meth:`~airflow.providers.anthropic.hooks.anthropic.AnthropicHook.stream_batch_results` + and persist them to object storage; results can be very large and must not be pushed + to XCom. Results are retained for 29 days after the batch is created. + + .. note:: + A retry re-submits a brand-new batch. Prefer ``retries=0`` on this task (the + submitted ``batch_id`` is pushed to XCom under key ``batch_id`` immediately, so + a crashed run never loses track of an in-flight batch). + + .. seealso:: + For more information, take a look at the guide: + :ref:`howto/operator:AnthropicBatchOperator` + + :param requests: A list of ``{"custom_id": str, "params": {...}}`` dicts, where + ``params`` is a ``messages.create`` payload (``model``, ``max_tokens``, ``messages``, ...). + :param conn_id: The Anthropic connection ID to use. + :param deferrable: Run the operator in deferrable mode. + :param poll_interval: Seconds between status checks, in both the synchronous and + deferrable paths. + :param timeout: Seconds to wait for the batch to reach a terminal status. Defaults to + 24 hours (the Message Batches SLA). In deferrable mode this also bounds the + deferral; set ``execution_timeout`` only if you want a shorter hard cap (note a + shorter ``execution_timeout`` preempts the graceful cancel-on-timeout path). + :param wait_for_completion: Whether to wait for the batch to complete. If ``False``, + the operator returns the batch ID immediately after submission. + :param fail_on_partial_error: If ``True``, fail the task when any request errored or + expired. Defaults to ``False`` (succeed and log a warning so the successful + results are not discarded). + """ + + template_fields: Sequence[str] = ("requests",) + + def __init__( + self, + requests: list[dict[str, Any]], + conn_id: str = AnthropicHook.default_conn_name, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + poll_interval: float = 60, + timeout: float = 24 * 60 * 60, + wait_for_completion: bool = True, + fail_on_partial_error: bool = False, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.requests = requests + self.conn_id = conn_id + self.deferrable = deferrable + self.poll_interval = poll_interval + self.timeout = timeout + self.wait_for_completion = wait_for_completion + self.fail_on_partial_error = fail_on_partial_error + self.batch_id: str | None = None + + @cached_property + def hook(self) -> AnthropicHook: + """Return an instance of the AnthropicHook.""" + return AnthropicHook(conn_id=self.conn_id) + + def execute(self, context: Context) -> str | None: + if not self.requests: + raise ValueError("AnthropicBatchOperator requires at least one request; got an empty list.") + batch = self.hook.create_batch(self.requests) + self.batch_id = batch.id + # Push immediately so a crash between submit and completion never loses the batch. + context["ti"].xcom_push(key="batch_id", value=batch.id) + self.log.info("Submitted Anthropic Message Batch %s (%d requests)", batch.id, len(self.requests)) + + if not self.wait_for_completion: + return self.batch_id + + if self.deferrable: + self.defer( + # Backstop the deferral slightly beyond the trigger's own end_time so the + # trigger's clean "timeout" event (which cancels the batch) wins over a + # generic AirflowTaskTimeout. A user-set execution_timeout still applies + # as a shorter hard cap. + timeout=self.execution_timeout or timedelta(seconds=self.timeout + self.poll_interval + 60), + trigger=AnthropicBatchTrigger( + conn_id=self.conn_id, + batch_id=self.batch_id, + poll_interval=self.poll_interval, + end_time=time.time() + self.timeout, + ), + method_name="execute_complete", + ) + + self.log.info("Waiting for batch %s to complete", self.batch_id) + try: + batch = self.hook.wait_for_batch( + self.batch_id, wait_seconds=self.poll_interval, timeout=self.timeout + ) + except AnthropicBatchTimeout: + # Mirror the deferrable execute_complete: tear down the still-running batch + # before the task fails, so a sync timeout does not leave it billing. + self.log.warning("Batch %s timed out; requesting cancellation.", self.batch_id) + self._cancel_batch_quietly() + raise + counts = batch.request_counts + self._apply_policy(counts.canceled, counts.errored, counts.expired, counts.succeeded) + return self.batch_id + + def execute_complete(self, context: Context, event: Any = None) -> str: + """ + Resume after the trigger fires. + + The deferred task is a fresh instance, so the batch ID is read from the event, + not ``self.batch_id``. + """ + self.batch_id = event["batch_id"] + status = event["status"] + if status == "timeout": + self.log.warning("Batch %s timed out; requesting cancellation.", self.batch_id) + self._cancel_batch_quietly() + raise AnthropicBatchTimeout(event["message"]) + if status == "error": Review Comment: Good catch, fixed in `8a72eb5`. The deferred `error` event now cancels the batch (and the agent operator archives the session) before raising, and the synchronous `execute()` except is broadened from timeout-only to any post-submission failure. I also added transient-poll tolerance to the synchronous wait helpers (matching the trigger's 5-failure tolerance) so a single blip doesn't cancel a still-healthy batch whose results are recoverable for 29 days. ########## providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py: ########## @@ -0,0 +1,531 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import time +from enum import Enum +from functools import cached_property +from typing import TYPE_CHECKING, Any, cast + +from anthropic import ( + Anthropic, + AnthropicAWS, + AnthropicBedrock, + AnthropicFoundry, + AnthropicVertex, + IdentityTokenFile, + WorkloadIdentityCredentials, +) + +from airflow.providers.anthropic.exceptions import ( + AnthropicAgentSessionError, + AnthropicAgentSessionTimeout, + AnthropicBatchTimeout, + AnthropicError, +) +from airflow.providers.common.compat.sdk import BaseHook + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator + + from anthropic.types import Message + from anthropic.types.beta import ( + BetaEnvironment, + BetaManagedAgentsAgent, + BetaManagedAgentsSession, + environment_create_params, + ) + from anthropic.types.beta.sessions import BetaManagedAgentsEventParams + from anthropic.types.messages import MessageBatch, MessageBatchIndividualResponse + from anthropic.types.messages.batch_create_params import Request + +#: Default model used when an operator or hook caller does not specify one. +DEFAULT_MODEL = "claude-opus-4-8" + +#: Platforms that serve the first-party-only endpoints (Message Batches, token +#: counting, the Models API). Amazon Bedrock, Google Vertex AI and Microsoft +#: Foundry do not serve these, so the hook fails fast rather than surfacing a +#: raw ``404`` from the SDK. +FIRST_PARTY_PLATFORMS = frozenset({"anthropic", "aws"}) + +AnthropicClient = Anthropic | AnthropicBedrock | AnthropicVertex | AnthropicAWS | AnthropicFoundry + + +class BatchStatus(str, Enum): + """Top-level ``processing_status`` of an Anthropic Message Batch.""" + + IN_PROGRESS = "in_progress" + CANCELING = "canceling" + ENDED = "ended" + + @classmethod + def is_in_progress(cls, status: str) -> bool: Review Comment: Kept the name for SDK-vocabulary consistency and extended the docstring to spell out that `canceling` is also non-terminal (so it returns True too), i.e. the method means "not yet terminal". Done in `8a72eb5`. ########## providers/anthropic/docs/index.rst: ########## @@ -0,0 +1,130 @@ + +.. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +``apache-airflow-providers-anthropic`` +====================================== + + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: Basics + + Home <self> + Changelog <changelog> + Security <security> + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: Guides + + Connection types <connections> + Operators <operators/anthropic> + + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: Resources + + Python API <_api/airflow/providers/anthropic/index> + PyPI Repository <https://pypi.org/project/apache-airflow-providers-anthropic/> + Installing from sources <installing-providers-from-sources> + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: System tests + + System Tests <_api/tests/system/anthropic/index> + +.. THE REMAINDER OF THE FILE IS AUTOMATICALLY GENERATED. IT WILL BE OVERWRITTEN AT RELEASE TIME! + + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: Commits + + Detailed list of commits <commits> + + +apache-airflow-providers-anthropic package +------------------------------------------------------ + +`Anthropic <https://docs.claude.com/>`__ provider for Apache Airflow. +Wraps the official Anthropic Python SDK to run the Claude Message Batches API +asynchronously from Airflow, plus direct message and token-counting helpers. + + +Release: 0.1.0 + +Provider package +---------------- + +This package is for the ``anthropic`` provider. +All classes for this package are included in the ``airflow.providers.anthropic`` python package. + +Installation +------------ + +You can install this package on top of an existing Airflow installation via +``pip install apache-airflow-providers-anthropic``. +For the minimum Airflow version supported, see ``Requirements`` below. + +Requirements +------------ + +The minimum Apache Airflow version supported by this provider distribution is ``3.0.0``. + +========================================== ================== +PIP package Version required +========================================== ================== +``apache-airflow`` ``>=3.0.0`` +``apache-airflow-providers-common-compat`` ``>=1.12.0`` Review Comment: Agreed, a shared common.ai interface for AI providers would be ideal. Happy to help move that way if/when it takes shape; for now this stays Anthropic-specific. Thanks! -- 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]
