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,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
             "roleArn": role_arn,
         }
 
+        bedrock_tags = litellm_params.get("bedrock_tags") or 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/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,93 @@ 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_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
 # --------------------------------------------------------------------------- #