From babd7fe0240b4b6dfdf8394c6f0527dfc8cb6fd5 Mon Sep 17 00:00:00 2001
From: Kaxil Naik <kaxilnaik@gmail.com>
Date: Wed, 8 Jul 2026 21:56:00 +0530
Subject: [PATCH] Default the Anthropic batch model from the connection

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.
---
 .../anthropic/docs/operators/anthropic.rst    |  4 +++
 .../providers/anthropic/hooks/anthropic.py    | 24 ++++++++++++++--
 .../providers/anthropic/operators/batch.py    | 10 +++++--
 .../unit/anthropic/hooks/test_anthropic.py    | 28 +++++++++++++++++++
 .../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 f22d100ee1910..be0286d43f0ea 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 cdd344f847654..5f0cbd4088625 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 @@ 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"):
+            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 5ba8722d5568c..7e2fad8548be9 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 @@ 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
@@ -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)
diff --git a/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py b/providers/anthropic/tests/unit/anthropic/hooks/test_anthropic.py
index 05e39d5a80a95..2fe3f6fe1f7d3 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 @@ 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")
diff --git a/providers/anthropic/tests/unit/anthropic/operators/test_batch.py b/providers/anthropic/tests/unit/anthropic/operators/test_batch.py
index 656658b3362ea..28b0a40198b0a 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 @@ 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)
