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)