Merge pull request #33733 from BerriAI/litellm_lit_4162_bedrock_batch_tags

feat(bedrock): forward bedrock_tags to CreateModelInvocationJob for batch jobs

litellm/llms/bedrock/batches/transformation.py

@@ -4,6 +4,7 @@ import time
 from typing import Any, Dict, List, Literal, Optional, Union, cast
 
 from httpx import Headers, Response
+from pydantic import TypeAdapter, ValidationError
 
 from litellm.litellm_core_utils.cloud_storage_security import (
     BEDROCK_MANAGED_S3_BATCH_PREFIX,
@@ -19,6 +20,7 @@ from litellm.types.llms.bedrock import (
     BedrockOutputDataConfig,
     BedrockS3InputDataConfig,
     BedrockS3OutputDataConfig,
+    BedrockTag,
 )
 from litellm.types.llms.openai import (
     AllMessageValues,
@@ -38,6 +40,18 @@ _S3_BATCH_FILE_UUID_SUFFIX_PATTERN = re.compile(
     r"-[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\.jsonl$"
 )
 
+_BEDROCK_TAGS_ADAPTER: TypeAdapter[list[BedrockTag]] = TypeAdapter(list[BedrockTag])
+
+
+def _validate_bedrock_tags(raw_tags: object) -> list[BedrockTag]:
+    try:
+        return _BEDROCK_TAGS_ADAPTER.validate_python(raw_tags, strict=True)
+    except ValidationError as e:
+        raise ValueError(
+            "Invalid 'bedrock_tags' value. Expected a list of {'key': <str>, 'value': <str>} dicts, "
+            f"e.g. [{{'key': 'team', 'value': 'genai'}}]. Got: {raw_tags!r}"
+        ) from e
+
 
 class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
     """
@@ -201,6 +215,11 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
             "roleArn": role_arn,
         }
 
+        config_bedrock_tags = litellm_params.get("bedrock_tags")
+        bedrock_tags = config_bedrock_tags if config_bedrock_tags is not None else optional_params.get("bedrock_tags")
+        if bedrock_tags is not None:
+            bedrock_request["tags"] = _validate_bedrock_tags(bedrock_tags)
+
         # Add optional parameters if provided
         completion_window = create_batch_data.get("completion_window")
         if completion_window:

litellm/proxy/auth/auth_utils.py

@@ -273,6 +273,7 @@ _BANNED_REQUEST_BODY_PARAMS: Tuple[str, ...] = (
     # re-route the request's retention and accounting to any project
     # reachable with the deployment's shared AWS credentials.
     "aws_bedrock_project_id",
+    "bedrock_tags",
     # Provider-specific endpoint overrides that flow into the outbound
     # request via ``optional_params``. Same threat as ``api_base``:
     # ``s3_endpoint_url`` redirects Bedrock file uploads to attacker

litellm/types/llms/bedrock.py

@@ -985,6 +985,11 @@ class BedrockOutputDataConfig(TypedDict):
     s3OutputDataConfig: BedrockS3OutputDataConfig
 
 
+class BedrockTag(TypedDict):
+    key: str
+    value: str
+
+
 class BedrockCreateBatchRequest(TypedDict, total=False):
     """
     Request structure for creating a Bedrock batch inference job.
@@ -999,7 +1004,7 @@ class BedrockCreateBatchRequest(TypedDict, total=False):
     outputDataConfig: BedrockOutputDataConfig
     timeoutDurationInHours: Optional[int]
     clientRequestToken: Optional[str]
-    tags: Optional[List[dict]]
+    tags: Optional[List[BedrockTag]]
 
 
 BedrockBatchJobStatus = Literal["Submitted", "InProgress", "Completed", "Failed", "Stopping", "Stopped"]

tests/test_litellm/llms/bedrock/batches/test_transformation.py

@@ -258,6 +258,112 @@ def test_create_request_no_timeout_for_non_24h_window(config):
     assert "timeoutDurationInHours" not in mock_sign.call_args.kwargs["data"]
 
 
+def test_create_request_forwards_bedrock_tags_from_litellm_params(config):
+    tags = [
+        {"key": "application", "value": "genai-proxy"},
+        {"key": "team", "value": "ml-platform"},
+    ]
+    with patch.object(
+        config.common_utils,
+        "generate_unique_job_name",
+        return_value="litellm-batch-1",
+    ), patch.object(config.common_utils, "sign_aws_request") as mock_sign:
+        mock_sign.return_value = ({}, b"{}")
+        config.transform_create_batch_request(
+            model="m",
+            create_batch_data={"input_file_id": "s3://b/in.jsonl"},
+            optional_params={},
+            litellm_params={
+                "aws_batch_role_arn": "arn:aws:iam::1:role/r",
+                "bedrock_tags": tags,
+            },
+        )
+    assert mock_sign.call_args.kwargs["data"]["tags"] == tags
+
+
+def test_create_request_forwards_bedrock_tags_from_optional_params(config):
+    tags = [{"key": "env", "value": "prod"}]
+    with patch.object(
+        config.common_utils,
+        "generate_unique_job_name",
+        return_value="litellm-batch-1",
+    ), patch.object(config.common_utils, "sign_aws_request") as mock_sign:
+        mock_sign.return_value = ({}, b"{}")
+        config.transform_create_batch_request(
+            model="m",
+            create_batch_data={"input_file_id": "s3://b/in.jsonl"},
+            optional_params={"bedrock_tags": tags},
+            litellm_params={"aws_batch_role_arn": "arn:aws:iam::1:role/r"},
+        )
+    assert mock_sign.call_args.kwargs["data"]["tags"] == tags
+
+
+def test_create_request_empty_litellm_params_tags_do_not_fall_through(config):
+    with patch.object(
+        config.common_utils,
+        "generate_unique_job_name",
+        return_value="litellm-batch-1",
+    ), patch.object(config.common_utils, "sign_aws_request") as mock_sign:
+        mock_sign.return_value = ({}, b"{}")
+        config.transform_create_batch_request(
+            model="m",
+            create_batch_data={"input_file_id": "s3://b/in.jsonl"},
+            optional_params={"bedrock_tags": [{"key": "env", "value": "prod"}]},
+            litellm_params={
+                "aws_batch_role_arn": "arn:aws:iam::1:role/r",
+                "bedrock_tags": [],
+            },
+        )
+    assert mock_sign.call_args.kwargs["data"]["tags"] == []
+
+
+def test_create_request_omits_tags_when_bedrock_tags_absent(config):
+    with patch.object(
+        config.common_utils,
+        "generate_unique_job_name",
+        return_value="litellm-batch-1",
+    ), patch.object(config.common_utils, "sign_aws_request") as mock_sign:
+        mock_sign.return_value = ({}, b"{}")
+        config.transform_create_batch_request(
+            model="m",
+            create_batch_data={"input_file_id": "s3://b/in.jsonl"},
+            optional_params={},
+            litellm_params={"aws_batch_role_arn": "arn:aws:iam::1:role/r"},
+        )
+    assert "tags" not in mock_sign.call_args.kwargs["data"]
+
+
+@pytest.mark.parametrize(
+    "bad_tags",
+    [
+        ["application=genai-proxy"],
+        [{"key": "application"}],
+        [{"value": "genai-proxy"}],
+        [{"key": "application", "value": 42}],
+        {"key": "application", "value": "genai-proxy"},
+        "application=genai-proxy",
+    ],
+)
+def test_create_request_rejects_malformed_bedrock_tags(config, bad_tags):
+    with patch.object(
+        config.common_utils,
+        "generate_unique_job_name",
+        return_value="litellm-batch-1",
+    ), patch.object(config.common_utils, "sign_aws_request") as mock_sign:
+        mock_sign.return_value = ({}, b"{}")
+        with pytest.raises(ValueError, match="Invalid 'bedrock_tags' value"):
+            config.transform_create_batch_request(
+                model="m",
+                create_batch_data={"input_file_id": "s3://b/in.jsonl"},
+                optional_params={},
+                litellm_params={
+                    "aws_batch_role_arn": "arn:aws:iam::1:role/r",
+                    "bedrock_tags": bad_tags,
+                },
+            )
+    mock_sign.assert_not_called()
+
+
 # --------------------------------------------------------------------------- #
 # transform_create_batch_response - status mapping + LiteLLMBatch shape
 # --------------------------------------------------------------------------- #

tests/test_litellm/proxy/auth/test_auth_utils.py

@@ -1944,6 +1944,91 @@ class TestIsRequestBodySafeBlocksRivaUseSsl:
         )
 
 
+class TestIsRequestBodySafeBlocksBedrockTags:
+    """``bedrock_tags`` lands as AWS resource tags on Bedrock batch jobs
+    created with the proxy's AWS identity, so a caller-supplied value can
+    forge ownership or cost-allocation labels; like
+    ``aws_bedrock_project_id`` it is blocked without an admin opt-in."""
+
+    def test_bedrock_tags_in_request_body_is_rejected(self):
+        with pytest.raises(ValueError, match="bedrock_tags"):
+            is_request_body_safe(
+                request_body={
+                    "model": "bedrock-batch-opus",
+                    "bedrock_tags": [{"key": "application", "value": "genai-proxy"}],
+                },
+                general_settings={},
+                llm_router=None,
+                model="bedrock-batch-opus",
+            )
+
+    def test_admin_opt_in_proxy_wide_allows_bedrock_tags(self):
+        assert (
+            is_request_body_safe(
+                request_body={
+                    "model": "bedrock-batch-opus",
+                    "bedrock_tags": [{"key": "application", "value": "genai-proxy"}],
+                },
+                general_settings={"allow_client_side_credentials": True},
+                llm_router=None,
+                model="bedrock-batch-opus",
+            )
+            is True
+        )
+
+    def test_admin_opt_in_per_deployment_allows_bedrock_tags(self):
+        from litellm import Router
+
+        router = Router(
+            model_list=[
+                {
+                    "model_name": "bedrock-batch-opus",
+                    "litellm_params": {
+                        "model": "bedrock/us.anthropic.claude-opus-4-7",
+                        "configurable_clientside_auth_params": ["bedrock_tags"],
+                    },
+                }
+            ]
+        )
+        assert (
+            is_request_body_safe(
+                request_body={
+                    "model": "bedrock-batch-opus",
+                    "bedrock_tags": [{"key": "application", "value": "genai-proxy"}],
+                },
+                general_settings={},
+                llm_router=router,
+                model="bedrock-batch-opus",
+            )
+            is True
+        )
+
+    def test_per_deployment_opt_in_for_other_param_still_rejects_bedrock_tags(self):
+        from litellm import Router
+
+        router = Router(
+            model_list=[
+                {
+                    "model_name": "bedrock-batch-opus",
+                    "litellm_params": {
+                        "model": "bedrock/us.anthropic.claude-opus-4-7",
+                        "configurable_clientside_auth_params": ["api_base"],
+                    },
+                }
+            ]
+        )
+        with pytest.raises(ValueError, match="bedrock_tags"):
+            is_request_body_safe(
+                request_body={
+                    "model": "bedrock-batch-opus",
+                    "bedrock_tags": [{"key": "application", "value": "genai-proxy"}],
+                },
+                general_settings={},
+                llm_router=router,
+                model="bedrock-batch-opus",
+            )
+
+
 # ── is_request_body_safe nested-config recursion (VERIA-6) ────────────────────
 
 

tests/test_litellm/test_router.py

@@ -5806,3 +5806,63 @@ def test_get_configured_token_limits_coerces_numeric_strings():
     )
 
     assert router.get_configured_token_limits("quoted-limits-model") == (32000, 8000)
+
+
+@pytest.mark.asyncio
+async def test_acreate_batch_request_bedrock_tags_override_deployment_tags():
+    import httpx
+
+    from litellm.llms.bedrock.common_utils import CommonBatchFilesUtils
+
+    deployment_tags = [{"key": "application", "value": "config-level"}]
+    request_tags = [{"key": "application", "value": "request-level"}]
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "bedrock-batch-model",
+                "litellm_params": {
+                    "model": "bedrock/us.anthropic.claude-sonnet-5",
+                    "aws_batch_role_arn": "arn:aws:iam::123:role/batch-role",
+                    "aws_region_name": "us-west-2",
+                    "bedrock_tags": deployment_tags,
+                },
+            }
+        ]
+    )
+
+    def fake_response():
+        return httpx.Response(
+            status_code=200,
+            json={
+                "jobArn": "arn:aws:bedrock:us-west-2:123:model-invocation-job/abc1234567",
+                "status": "Submitted",
+            },
+        )
+
+    mock_client = MagicMock()
+    mock_client.post = AsyncMock(side_effect=lambda *args, **kwargs: fake_response())
+
+    with patch.object(
+        CommonBatchFilesUtils,
+        "sign_aws_request",
+        return_value=({"Authorization": "signed"}, b"{}"),
+    ) as mock_sign, patch(
+        "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client",
+        return_value=mock_client,
+    ):
+        await router.acreate_batch(
+            model="bedrock-batch-model",
+            input_file_id="s3://bucket/input.jsonl",
+            endpoint="/v1/chat/completions",
+            completion_window="24h",
+        )
+        assert mock_sign.call_args.kwargs["data"]["tags"] == deployment_tags
+
+        await router.acreate_batch(
+            model="bedrock-batch-model",
+            input_file_id="s3://bucket/input.jsonl",
+            endpoint="/v1/chat/completions",
+            completion_window="24h",
+            bedrock_tags=request_tags,
+        )
+        assert mock_sign.call_args.kwargs["data"]["tags"] == request_tags