# 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,
AnthropicTriggerEventError,
)
from airflow.providers.common.compat.sdk import AirflowSkipException, BaseHook
[docs]
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.
[docs]
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.
[docs]
AnthropicClient = Anthropic | AnthropicBedrock | AnthropicVertex | AnthropicAWS | AnthropicFoundry
#: Consecutive failed polls tolerated in the synchronous wait helpers before giving up
#: (transient errors). Mirrors the deferrable triggers' tolerance so a single blip does
#: not fail (and cancel/archive) a still-healthy batch or session.
[docs]
MAX_CONSECUTIVE_POLL_FAILURES = 5
[docs]
class BatchStatus(str, Enum):
"""Top-level ``processing_status`` of an Anthropic Message Batch."""
[docs]
IN_PROGRESS = "in_progress"
[docs]
CANCELING = "canceling"
@classmethod
[docs]
def is_in_progress(cls, status: str) -> bool:
"""
Return ``True`` while the batch has not reached the terminal ``ended`` status.
This is broader than the ``in_progress`` value: a ``canceling`` batch is also
non-terminal (cancellation is in flight but the batch has not ended yet), so it
returns ``True`` too. Read the name as "not yet terminal", not "equals the
``in_progress`` status".
"""
return status != cls.ENDED
[docs]
class SessionStatus(str, Enum):
"""Status of a Managed Agents session."""
[docs]
RESCHEDULING = "rescheduling"
[docs]
TERMINATED = "terminated"
@classmethod
[docs]
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.
[docs]
OUTCOME_FAILURE_RESULTS = frozenset({"failed", "max_iterations_reached", "interrupted"})
[docs]
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
#: Statuses the provider's triggers emit in their terminal event.
[docs]
TRIGGER_EVENT_STATUSES = frozenset({"success", "error", "timeout"})
[docs]
def validate_execute_complete_event(event: dict[str, Any] | None = None) -> dict[str, Any]:
"""
Validate the event a deferred task resumes with, returning it if well-formed.
The event crosses the triggerer/worker boundary through the metadata DB, so a
resuming task can receive ``None`` or a status its handlers do not recognize
(version skew, a custom trigger). Both must fail loudly: the ``execute_complete``
handlers raise on ``timeout``/``error`` and treat everything else as success, so
an unrecognized status would otherwise silently succeed.
"""
if event is None:
raise AnthropicTriggerEventError("Trigger error: event is None")
if event.get("status") not in TRIGGER_EVENT_STATUSES:
raise AnthropicTriggerEventError(
f"Unexpected trigger event status {event.get('status')!r}: {event!r}"
)
return event
[docs]
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.batch.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)
[docs]
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`` and ``aws`` platforms.
- ``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>`.
"""
[docs]
conn_name_attr = "conn_id"
[docs]
default_conn_name = "anthropic_default"
[docs]
conn_type = "anthropic"
[docs]
hook_name = "Anthropic"
def __init__(self, conn_id: str = default_conn_name, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)
@cached_property
def _connection(self):
return self.get_connection(self.conn_id)
@cached_property
@cached_property
[docs]
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
[docs]
def conn(self) -> AnthropicClient:
"""Return the Anthropic client for the configured platform."""
return self.get_conn()
[docs]
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(aws_region=extras.get("aws_region"), **client_kwargs)
if platform == "foundry":
api_key = client_kwargs.pop("api_key", None) or conn.password
return AnthropicFoundry(api_key=api_key, resource=extras.get("resource"), **client_kwargs)
if platform != "anthropic":
raise AnthropicError(
f"Unknown Anthropic platform {platform!r}. "
"Expected one of: anthropic, bedrock, vertex, aws, foundry."
)
base_url = client_kwargs.pop("base_url", None) or conn.host or None
wif = extras.get("workload_identity")
if wif:
return Anthropic(
credentials=self._workload_identity_credentials(wif), base_url=base_url, **client_kwargs
)
api_key = client_kwargs.pop("api_key", None) or conn.password
if api_key:
return Anthropic(api_key=api_key, base_url=base_url, **client_kwargs)
# No static key and no explicit federation config: let the SDK resolve credentials
# from the environment, which supports env-driven Workload Identity Federation
# (ANTHROPIC_FEDERATION_RULE_ID etc.) and ``ant`` profiles.
return Anthropic(base_url=base_url, **client_kwargs)
@staticmethod
def _workload_identity_credentials(wif: dict[str, Any]) -> WorkloadIdentityCredentials:
"""
Build a WIF credential from the connection ``extra['workload_identity']`` mapping.
Exchanges a short-lived OIDC token (read from ``identity_token_file``) for an
Anthropic access token. See
https://platform.claude.com/docs/en/manage-claude/workload-identity-federation.
"""
kwargs: dict[str, Any] = {
"identity_token_provider": IdentityTokenFile(wif["identity_token_file"]),
"federation_rule_id": wif["federation_rule_id"],
"organization_id": wif["organization_id"],
"service_account_id": wif["service_account_id"],
}
if wif.get("workspace_id"):
kwargs["workspace_id"] = wif["workspace_id"]
if wif.get("scope"):
kwargs["scope"] = wif["scope"]
return WorkloadIdentityCredentials(**kwargs)
def _resolve_model(self, model: str | None) -> str:
"""Resolve the effective model id; Bedrock rejects a bare id, so require its prefix."""
resolved = model or self.default_model
# Valid Bedrock ids either start with the ``anthropic.`` provider prefix or carry a
# region/profile prefix as a dotted component (e.g. ``us.anthropic.``, ``global.anthropic.``).
is_bedrock_model_id = resolved.startswith("anthropic.") or ".anthropic." in resolved
if self.platform == "bedrock" and not is_bedrock_model_id:
raise AnthropicError(
f"Model {resolved!r} is not a valid Amazon Bedrock model id. Bedrock ids carry a "
"provider/region prefix (e.g. 'global.anthropic.claude-opus-4-6-v1'); set one via "
"the 'model' argument or the connection's extra['model']."
)
return resolved
def _require_first_party(self, feature: str) -> None:
if self.platform not in FIRST_PARTY_PLATFORMS:
raise AnthropicError(
f"{feature} is not available on the {self.platform!r} platform. "
"Use the first-party Anthropic API (platform='anthropic') or "
"Claude Platform on AWS (platform='aws')."
)
@property
def _first_party_conn(self) -> Anthropic:
"""
Client cast to the first-party type for endpoints only it exposes.
Callers must guard with :meth:`_require_first_party` first; the Bedrock/Vertex/
Foundry clients don't expose ``beta.agents``/``beta.sessions``/``models``.
"""
return cast("Anthropic", self.conn)
[docs]
def test_connection(self) -> tuple[bool, str]:
"""Test the Anthropic connection."""
try:
if self.platform in FIRST_PARTY_PLATFORMS:
# Narrowed by the platform guard: only the first-party / AWS clients,
# which expose the Models API, reach this branch.
self._first_party_conn.models.list()
return True, "Connection established!"
# models.list() is not served on bedrock/vertex/foundry; building the
# client validates the configuration without a paid request.
self.get_conn()
return True, f"Connection configured for platform {self.platform!r} (no live check available)."
except Exception as e:
return False, str(e)
[docs]
def create_message(
self,
messages: list[dict[str, Any]],
model: str | None = None,
max_tokens: int = 1024,
system: str | None = None,
**kwargs: Any,
) -> Message:
"""
Create a single message response (one-shot ``messages.create``).
:param messages: The conversation so far, as a list of message dicts.
:param model: Model ID to use. Defaults to :attr:`default_model` (the connection's
``extra['model']`` or :data:`DEFAULT_MODEL`).
:param max_tokens: Maximum number of tokens to generate.
:param system: Optional system prompt.
"""
params: dict[str, Any] = {
"model": self._resolve_model(model),
"max_tokens": max_tokens,
"messages": messages,
**kwargs,
}
if system is not None:
params["system"] = system
return self.conn.messages.create(**params)
[docs]
def count_tokens(
self,
messages: list[dict[str, Any]],
model: str | None = None,
system: str | None = None,
**kwargs: Any,
) -> int:
"""Return the number of input tokens the given request would consume."""
self._require_first_party("Token counting")
params: dict[str, Any] = {"model": model or self.default_model, "messages": messages, **kwargs}
if system is not None:
params["system"] = system
return self.conn.messages.count_tokens(**params).input_tokens
[docs]
def create_batch(self, requests: list[dict[str, Any]]) -> MessageBatch:
"""
Submit a Message Batch.
:param requests: A list of ``{"custom_id": str, "params": {...}}`` dicts, where
``params`` is a ``messages.create`` payload (``model``, ``max_tokens``,
``messages``, ...).
"""
self._require_first_party("The Message Batches API")
# ``Request`` is a TypedDict, so the plain dicts callers build match structurally.
return self.conn.messages.batches.create(requests=cast("Iterable[Request]", requests))
[docs]
def get_batch(self, batch_id: str) -> MessageBatch:
"""Retrieve a Message Batch by ID."""
self._require_first_party("The Message Batches API")
return self.conn.messages.batches.retrieve(batch_id)
[docs]
def cancel_batch(self, batch_id: str) -> MessageBatch:
"""Request cancellation of a Message Batch."""
self._require_first_party("The Message Batches API")
return self.conn.messages.batches.cancel(batch_id)
[docs]
def list_batches(self, **kwargs: Any) -> Any:
"""Return a (paginated) list of Message Batches."""
self._require_first_party("The Message Batches API")
return self.conn.messages.batches.list(**kwargs)
[docs]
def stream_batch_results(self, batch_id: str) -> Iterator[MessageBatchIndividualResponse]:
"""
Return a streaming iterator of per-request results, keyed by ``custom_id``.
Results stream from the API and arrive in **arbitrary order** — key them by
``result.custom_id``, never by position. Results are available for 29 days
after the batch is created. The result set can be very large: iterate and
persist to object storage; do not materialize it into XCom.
"""
# Return (don't ``yield``) so the platform guard fails fast at call time
# rather than only when the caller starts iterating.
self._require_first_party("The Message Batches API")
return self.conn.messages.batches.results(batch_id)
[docs]
def wait_for_batch(
self, batch_id: str, wait_seconds: float = 3, timeout: float = 24 * 60 * 60
) -> MessageBatch:
"""
Poll a batch synchronously until it reaches the terminal ``ended`` status.
:param batch_id: The batch to wait for.
:param wait_seconds: Seconds to sleep between polls.
:param timeout: Maximum seconds to wait before raising :class:`AnthropicBatchTimeout`.
:return: The terminal :class:`~anthropic.types.messages.MessageBatch`.
"""
start = time.monotonic()
consecutive_failures = 0
while True:
try:
batch = self.get_batch(batch_id)
except Exception as e:
# Tolerate transient poll errors (as the deferrable trigger does) so a
# single blip does not fail — and cancel — a still-running batch whose
# results remain recoverable for 29 days.
consecutive_failures += 1
if (
consecutive_failures >= MAX_CONSECUTIVE_POLL_FAILURES
or time.monotonic() - start > timeout
):
raise
self.log.warning("Polling batch %s failed (%s); retrying.", batch_id, e)
time.sleep(wait_seconds)
continue
consecutive_failures = 0
self.log.debug("Batch %s status=%s", batch_id, batch.processing_status)
if not BatchStatus.is_in_progress(batch.processing_status):
return batch
if time.monotonic() - start > timeout:
raise AnthropicBatchTimeout(
f"Batch {batch_id} did not reach a terminal status within {timeout} seconds."
)
time.sleep(wait_seconds)
# --- Managed Agents -------------------------------------------------------
# Agents and environments are persisted, reusable resources: create them once
# (these helpers, the ``ant`` CLI, or a setup script) and store the IDs. The
# operator references those IDs; it never creates an agent per run.
[docs]
def create_agent(self, name: str, model: str | None = None, **kwargs: Any) -> BetaManagedAgentsAgent:
"""
Create a (reusable, versioned) Managed Agents agent. One-time setup.
``model`` defaults to :attr:`default_model` (the connection's ``extra['model']``
or :data:`DEFAULT_MODEL`).
"""
self._require_first_party("Managed Agents")
agent = self._first_party_conn.beta.agents.create(
name=name, model=model or self.default_model, **kwargs
)
self.log.debug("Created agent %s (name=%r, model=%s)", agent.id, name, model or self.default_model)
return agent
[docs]
def create_environment(
self, name: str, config: dict[str, Any] | None = None, **kwargs: Any
) -> BetaEnvironment:
"""Create a (reusable) environment for agent sessions. One-time setup."""
self._require_first_party("Managed Agents")
if config is None:
config = {"type": "cloud", "networking": {"type": "unrestricted"}}
environment = self._first_party_conn.beta.environments.create(
name=name, config=cast("environment_create_params.Config", config), **kwargs
)
self.log.debug("Created environment %s (name=%r)", environment.id, name)
return environment
[docs]
def create_session(self, agent: str, environment_id: str, **kwargs: Any) -> BetaManagedAgentsSession:
"""Start a session against a pre-created agent + environment."""
self._require_first_party("Managed Agents")
return self._first_party_conn.beta.sessions.create(
agent=agent, environment_id=environment_id, **kwargs
)
[docs]
def get_session(self, session_id: str) -> BetaManagedAgentsSession:
"""Retrieve a session (carries its current ``status``)."""
self._require_first_party("Managed Agents")
return self._first_party_conn.beta.sessions.retrieve(session_id)
[docs]
def send_event(self, session_id: str, event: dict[str, Any]) -> Any:
"""Send a single event (e.g. a ``user.message`` or ``user.define_outcome``)."""
self._require_first_party("Managed Agents")
# Event dicts callers build match the SDK's TypedDict union structurally.
return self._first_party_conn.beta.sessions.events.send(
session_id, events=cast("list[BetaManagedAgentsEventParams]", [event])
)
[docs]
def archive_session(self, session_id: str) -> Any:
"""Archive a session (frees the server-side container). Best-effort teardown."""
self._require_first_party("Managed Agents")
return self._first_party_conn.beta.sessions.archive(session_id)
def _latest_idle_reason(self, session_id: str, kickoff_event_id: str | None) -> str | None:
"""
Return the ``stop_reason`` of the newest ``session.status_idle`` event, or ``None``.
Walks the event log newest-first. Returns ``None`` if the kickoff event is the most
recent event (the agent has not responded yet — defeats the start race) or no idle
event is found in the scan window.
"""
# The SDK cursor auto-paginates (page size 20); cap the walk at 100 events so a
# long event log can't make one poll iterate unboundedly.
examined = 0
for event in self._first_party_conn.beta.sessions.events.list(session_id, order="desc", limit=20):
if kickoff_event_id is not None and event.id == kickoff_event_id:
return None
if event.type == "session.status_idle":
return event.stop_reason.type
examined += 1
if examined >= 100:
break
return None
[docs]
def poll_session_completion(
self, session_id: str, *, expect_outcome: bool = False, kickoff_event_id: str | None = None
) -> tuple[bool, str | None]:
"""
Return ``(done, error_message)`` for one poll of a session.
Combines the session object (status / outcome verdict) with the event log
(``stop_reason`` of the latest idle) so a ``message`` run distinguishes genuine
``end_turn`` completion from ``requires_action`` / ``retries_exhausted``.
"""
session = self.get_session(session_id)
done, error_message, needs_event_check = evaluate_session_state(
session, expect_outcome=expect_outcome
)
self.log.debug(
"Session %s status=%s done=%s needs_event_check=%s",
session_id,
session.status,
done,
needs_event_check,
)
if not needs_event_check:
return done, error_message
reason = self._latest_idle_reason(session_id, kickoff_event_id)
if reason is None:
return False, None
if reason == "end_turn":
return True, None
return True, (
f"Session {session_id} is idle but did not complete ({reason}); "
"configure an autonomous agent or use an outcome run."
)
[docs]
def wait_for_session(
self,
session_id: str,
expect_outcome: bool = False,
kickoff_event_id: str | None = None,
poll_interval: float = 30,
timeout: float = 24 * 60 * 60,
) -> None:
"""
Poll a session synchronously until it completes.
:param session_id: The session to wait for.
:param expect_outcome: Whether the session is running a ``user.define_outcome`` loop
(completion judged from ``outcome_evaluations``).
:param kickoff_event_id: ID of the kickoff event, used to correlate the terminal
idle event on a ``message`` run (defeats the start race).
:param poll_interval: Seconds to sleep between polls.
:param timeout: Maximum seconds to wait before raising :class:`AnthropicAgentSessionTimeout`.
"""
start = time.monotonic()
consecutive_failures = 0
while True:
try:
done, error_message = self.poll_session_completion(
session_id, expect_outcome=expect_outcome, kickoff_event_id=kickoff_event_id
)
except Exception as e:
# Tolerate transient poll errors (as the deferrable trigger does) so a
# single blip does not fail — and archive — a still-running session.
consecutive_failures += 1
if (
consecutive_failures >= MAX_CONSECUTIVE_POLL_FAILURES
or time.monotonic() - start > timeout
):
raise
self.log.warning("Polling session %s failed (%s); retrying.", session_id, e)
time.sleep(poll_interval)
continue
consecutive_failures = 0
if done:
if error_message:
raise AnthropicAgentSessionError(error_message)
return
if time.monotonic() - start > timeout:
raise AnthropicAgentSessionTimeout(
f"Session {session_id} did not reach a terminal status within {timeout} seconds."
)
time.sleep(poll_interval)