Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions providers/anthropic/docs/operators/anthropic.rst
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,10 @@ Parameters

* ``requests`` — a list of ``{"custom_id": str, "params": {...}}`` dicts, where ``params`` is a
``messages.create`` payload (``model``, ``max_tokens``, ``messages``, ...).
* ``model`` — default model id applied to any request whose ``params`` omits ``model``. When
unset, those requests fall back to the connection's ``default_model`` (``extra['model']``). Set
it to choose the batch's model once instead of repeating it in every request; a request that
sets its own ``model`` always wins, so a batch can still mix models.
* ``conn_id`` — the Anthropic connection ID (default ``anthropic_default``).
* ``deferrable`` — run in deferrable mode (defaults to the ``operators.default_deferrable`` config).
* ``poll_interval`` — seconds between status checks, in both the synchronous and deferrable paths.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -409,17 +409,35 @@ def count_tokens(
params["system"] = system
return self.conn.messages.count_tokens(**params).input_tokens

def create_batch(self, requests: list[dict[str, Any]]) -> MessageBatch:
@staticmethod
def _apply_default_model(request: dict[str, Any], default_model: str) -> dict[str, Any]:
"""
Fill ``params['model']`` from ``default_model`` when the request omits it.

The input dict is never mutated, and a request that sets its own ``model`` is
returned unchanged, so a single batch can still mix models across requests.
"""
params = request.get("params")
if not isinstance(params, dict) or params.get("model"):
Comment thread
gopidesupavan marked this conversation as resolved.
return request
return {**request, "params": {**params, "model": default_model}}

def create_batch(self, requests: list[dict[str, Any]], model: str | None = None) -> 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``, ...).
``messages``, ...). A request that omits ``model`` inherits ``model`` below,
or the connection's ``default_model`` (``extra['model']``) when that is unset too.
:param model: Default model id for requests that do not set their own. Falls back
to the connection's :attr:`default_model`.
"""
self._require_first_party("The Message Batches API")
default_model = model or self.default_model
prepared = [self._apply_default_model(request, default_model) for request in requests]
# ``Request`` is a TypedDict, so the plain dicts callers build match structurally.
return self.conn.messages.batches.create(requests=cast("Iterable[Request]", requests))
return self.conn.messages.batches.create(requests=cast("Iterable[Request]", prepared))

def get_batch(self, batch_id: str) -> MessageBatch:
"""Retrieve a Message Batch by ID."""
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,10 @@ class AnthropicBatchOperator(BaseOperator):

:param requests: A list of ``{"custom_id": str, "params": {...}}`` dicts, where
``params`` is a ``messages.create`` payload (``model``, ``max_tokens``, ``messages``, ...).
A request that omits ``model`` inherits ``model`` below, or the connection's
``default_model`` (``extra['model']``) when that is unset too.
:param model: Default model id applied to requests that don't set their own. Lets you
pick the batch's model once instead of repeating it in every request.
: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
Expand All @@ -75,11 +79,12 @@ class AnthropicBatchOperator(BaseOperator):
results are not discarded).
"""

template_fields: Sequence[str] = ("requests",)
template_fields: Sequence[str] = ("requests", "model")

def __init__(
self,
requests: list[dict[str, Any]],
model: str | None = None,
conn_id: str = AnthropicHook.default_conn_name,
deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
poll_interval: float = 60,
Expand All @@ -90,6 +95,7 @@ def __init__(
) -> None:
super().__init__(**kwargs)
self.requests = requests
self.model = model
self.conn_id = conn_id
self.deferrable = deferrable
self.poll_interval = poll_interval
Expand All @@ -106,7 +112,7 @@ def hook(self) -> AnthropicHook:
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)
batch = self.hook.create_batch(self.requests, model=self.model)
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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -467,6 +467,34 @@ def test_create_batch_passes_requests(self):
hook.create_batch(reqs)
client.messages.batches.create.assert_called_once_with(requests=reqs)

def test_create_batch_defaults_missing_model_from_conn(self):
hook, client = self._hook_with_client(extra={"model": "claude-sonnet-4-6"})
hook.create_batch([{"custom_id": "a", "params": {"max_tokens": 1, "messages": []}}])
sent = client.messages.batches.create.call_args.kwargs["requests"]
assert sent[0]["params"]["model"] == "claude-sonnet-4-6"

def test_create_batch_preserves_explicit_request_model(self):
hook, client = self._hook_with_client(extra={"model": "claude-sonnet-4-6"})
hook.create_batch(
[{"custom_id": "a", "params": {"model": "claude-haiku-4-5", "max_tokens": 1, "messages": []}}]
)
sent = client.messages.batches.create.call_args.kwargs["requests"]
assert sent[0]["params"]["model"] == "claude-haiku-4-5"

def test_create_batch_model_arg_overrides_conn_default(self):
hook, client = self._hook_with_client(extra={"model": "claude-sonnet-4-6"})
hook.create_batch(
[{"custom_id": "a", "params": {"max_tokens": 1, "messages": []}}], model="claude-opus-4-8"
)
sent = client.messages.batches.create.call_args.kwargs["requests"]
assert sent[0]["params"]["model"] == "claude-opus-4-8"

def test_create_batch_does_not_mutate_caller_requests(self):
hook, _ = self._hook_with_client(extra={"model": "claude-sonnet-4-6"})
original = [{"custom_id": "a", "params": {"max_tokens": 1, "messages": []}}]
hook.create_batch(original)
assert "model" not in original[0]["params"]

def test_get_and_cancel_batch(self):
hook, client = self._hook_with_client()
hook.get_batch("batch_1")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,18 @@ def test_sync_non_timeout_error_cancels_and_raises(self, mock_hook_prop):
op.execute(_context())
hook.cancel_batch.assert_called_once_with("batch_1")

@mock.patch.object(AnthropicBatchOperator, "hook", new_callable=mock.PropertyMock)
def test_execute_forwards_model_to_hook(self, mock_hook_prop):
hook = mock.MagicMock(spec=AnthropicHook)
hook.create_batch.return_value.id = "batch_1"
mock_hook_prop.return_value = hook

op = AnthropicBatchOperator(
task_id="t", requests=REQUESTS, model="claude-haiku-4-5", wait_for_completion=False
)
op.execute(_context())
hook.create_batch.assert_called_once_with(REQUESTS, model="claude-haiku-4-5")

@mock.patch.object(AnthropicBatchOperator, "hook", new_callable=mock.PropertyMock)
def test_empty_requests_raises_before_any_api_call(self, mock_hook_prop):
hook = mock.MagicMock(spec=AnthropicHook)
Expand Down