diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 4fcf7cf91cb..a4ff1c78467 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/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': , 'value': } 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: diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 38900260c98..293bb74e211 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/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 diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index bdf6b8fefed..d9f8229dbed 100644 --- a/litellm/types/llms/bedrock.py +++ b/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"] diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index d1ad5943ae6..3681daffe5e 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/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 # --------------------------------------------------------------------------- # diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index b5d8727f7e6..72bd215b9be 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/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) ──────────────────── diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 1c175bf6f44..cffc3dd0aba 100644 --- a/tests/test_litellm/test_router.py +++ b/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