This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0e4b6600a45 Default the Anthropic batch model from the connection
(#69624)
0e4b6600a45 is described below
commit 0e4b6600a452a423848d00f3f737e9356456623a
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Jul 8 19:11:59 2026 +0100
Default the Anthropic batch model from the connection (#69624)
Batch requests that omit params['model'] now inherit it: from a new
AnthropicBatchOperator `model` argument, else the connection's
default_model (extra['model']). An explicit per-request model still
wins, so a batch can mix models. This aligns the batch path with the
hook's create_message/count_tokens, which already default the model
from the connection. The hook fills a missing model without mutating
the caller's request dicts.
---
providers/anthropic/docs/operators/anthropic.rst | 4 ++++
.../airflow/providers/anthropic/hooks/anthropic.py | 24 ++++++++++++++++---
.../airflow/providers/anthropic/operators/batch.py | 10 ++++++--
.../tests/unit/anthropic/hooks/test_anthropic.py | 28 ++++++++++++++++++++++
.../tests/unit/anthropic/operators/test_batch.py | 12 ++++++++++
5 files changed, 73 insertions(+), 5 deletions(-)
diff --git a/providers/anthropic/docs/operators/anthropic.rst
b/providers/anthropic/docs/operators/anthropic.rst
index f22d100ee19..be0286d43f0 100644
--- a/providers/anthropic/docs/operators/anthropic.rst
+++ b/providers/anthropic/docs/operators/anthropic.rst
@@ -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.
diff --git
a/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py
b/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py
index cdd344f8476..5f0cbd40886 100644
--- a/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py
+++ b/providers/anthropic/src/airflow/providers/anthropic/hooks/anthropic.py
@@ -409,17 +409,35 @@ class AnthropicHook(BaseHook):
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"):
+ 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."""
diff --git
a/providers/anthropic/src/airflow/providers/anthropic/operators/batch.py
b/providers/anthropic/src/airflow/providers/anthropic/operators/batch.py
index 5ba8722d556..7e2fad8548b 100644
--- a/providers/anthropic/src/airflow/providers/anthropic/operators/batch.py
+++ b/providers/anthropic/src/airflow/providers/anthropic/operators/batch.py
@@ -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
@@ -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,
@@ -90,6 +95,7 @@ class AnthropicBatchOperator(BaseOperator):
) -> None:
super().__init__(**kwargs)
self.requests = requests
+ self.model = model
self.conn_id = conn_id
self.deferrable = deferrable
self.poll_interval = poll_interval
@@ -106,7 +112,7 @@ class AnthropicBatchOperator(BaseOperator):
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)
diff --git a/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py
b/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py
index 05e39d5a80a..2fe3f6fe1f7 100644
--- a/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py
+++ b/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py
@@ -467,6 +467,34 @@ class TestAnthropicHookFeatures:
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")
diff --git a/providers/anthropic/tests/unit/anthropic/operators/test_batch.py
b/providers/anthropic/tests/unit/anthropic/operators/test_batch.py
index 656658b3362..28b0a40198b 100644
--- a/providers/anthropic/tests/unit/anthropic/operators/test_batch.py
+++ b/providers/anthropic/tests/unit/anthropic/operators/test_batch.py
@@ -140,6 +140,18 @@ class TestAnthropicBatchOperatorExecute:
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)
