From 9d69fdac72b860627f627f9802ceb4758a7c372d Mon Sep 17 00:00:00 2001 From: daleselaji-dev <265319989+daleselaji-dev@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:14:56 +0800 Subject: [PATCH] fix(bedrock): use deployment credentials for AWS requests --- .../llms/bedrock/batches/transformation.py | 11 ++- litellm/llms/bedrock/files/transformation.py | 11 ++- .../test_bedrock_request_credentials.py | 95 +++++++++++++++++++ 3 files changed, 109 insertions(+), 8 deletions(-) create mode 100644 tests/litellm/llms/bedrock/test_bedrock_request_credentials.py diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 9e3dec26673..8b437af8581 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -130,7 +130,8 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): Get the complete URL for Bedrock batch creation. Bedrock batch jobs are created via the model invocation job API. """ - aws_region_name: Final = self._get_aws_region_name(optional_params, model) + request_params: Final = {**litellm_params, **optional_params} + aws_region_name: Final = self._get_aws_region_name(request_params, model) # Bedrock model invocation job endpoint # Format: https://bedrock.{region}.amazonaws.com/model-invocation-job @@ -232,14 +233,15 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): # For Bedrock, we need to return a pre-signed request with AWS auth headers # Use common utility for AWS signing + request_params: Final = {**litellm_params, **optional_params} endpoint_url: Final = ( - f"https://bedrock.{self._get_aws_region_name(optional_params, model)}.amazonaws.com/model-invocation-job" + f"https://bedrock.{self._get_aws_region_name(request_params, model)}.amazonaws.com/model-invocation-job" ) signed_headers, signed_data = self.common_utils.sign_aws_request( service_name="bedrock", data=bedrock_request, endpoint_url=endpoint_url, - optional_params=optional_params, + optional_params=request_params, method="POST", ) @@ -387,11 +389,12 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): endpoint_url: Final = f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{encoded_arn}" # Use common utility for AWS signing + request_params: Final = {**litellm_params, **optional_params} signed_headers, _ = self.common_utils.sign_aws_request( service_name="bedrock", data={}, # GET request has no body endpoint_url=endpoint_url, - optional_params=optional_params, + optional_params=request_params, method="GET", ) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 728fd02f001..072343ae936 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -257,6 +257,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """ Get the complete S3 URL for the file upload request """ + request_params: Final = {**litellm_params, **optional_params} bucket_name = litellm_params.get("s3_bucket_name") or os.getenv("AWS_S3_BUCKET_NAME") if not bucket_name: raise ValueError( @@ -265,7 +266,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name) s3_region_name: Final = litellm_params.get("s3_region_name") or optional_params.get("s3_region_name") - aws_region_name: Final = s3_region_name or self._get_aws_region_name(optional_params, model) + aws_region_name: Final = s3_region_name or self._get_aws_region_name(request_params, model) file_data: Final = data.get("file") purpose: Final = data.get("purpose") @@ -281,7 +282,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # S3 endpoint URL format s3_endpoint_url: Final = ( - optional_params.get("s3_endpoint_url") or f"https://s3.{aws_region_name}.amazonaws.com" + request_params.get("s3_endpoint_url") or f"https://s3.{aws_region_name}.amazonaws.com" ).rstrip("/") return f"{s3_endpoint_url}/{bucket_name}/{encoded_object_name}" @@ -728,14 +729,16 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): if s3_region_name: optional_params = {**optional_params, "aws_region_name": s3_region_name} + request_params: Final = {**litellm_params, **optional_params} + # Sign the request and return a pre-signed request object signed_headers, signed_body = self._sign_s3_request( content=file_content, api_base=api_base, - optional_params=optional_params, + optional_params=request_params, s3_encryption_key_id=resolve_s3_encryption_key_id( litellm_params=litellm_params, - optional_params=optional_params, + optional_params=request_params, ), ) diff --git a/tests/litellm/llms/bedrock/test_bedrock_request_credentials.py b/tests/litellm/llms/bedrock/test_bedrock_request_credentials.py new file mode 100644 index 00000000000..72da2daf9b5 --- /dev/null +++ b/tests/litellm/llms/bedrock/test_bedrock_request_credentials.py @@ -0,0 +1,95 @@ +from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig +from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + +def test_bedrock_file_upload_signing_uses_deployment_credentials(monkeypatch): + config = BedrockFilesConfig() + captured = {} + + def capture_signing(**kwargs): + captured.update(kwargs) + return {}, "" + + monkeypatch.setattr(config, "_sign_s3_request", capture_signing) + + result = config.transform_create_file_request( + model="", + create_file_data={ + "file": ( + "batch.jsonl", + b'{"custom_id":"req-1","body":{"model":"bedrock/model"}}\n', + "application/jsonl", + ), + "purpose": "batch", + }, + optional_params={}, + litellm_params={ + "s3_bucket_name": "deployment-bucket", + "aws_access_key_id": "deployment-access-key", + "aws_secret_access_key": "deployment-secret", + "aws_region_name": "eu-west-1", + }, + ) + + assert "eu-west-1" in result["url"] + assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key" + assert captured["optional_params"]["aws_secret_access_key"] == "deployment-secret" + assert captured["optional_params"]["aws_region_name"] == "eu-west-1" + + +def test_bedrock_batch_signing_uses_deployment_credentials(monkeypatch): + config = BedrockBatchesConfig() + captured = {} + + def capture_signing(**kwargs): + captured.update(kwargs) + return {}, b"{}" + + monkeypatch.setattr(config.common_utils, "sign_aws_request", capture_signing) + + result = config.transform_create_batch_request( + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + create_batch_data={ + "input_file_id": "s3://deployment-bucket/input.jsonl", + "completion_window": "24h", + "endpoint": "/v1/chat/completions", + }, + optional_params={}, + litellm_params={ + "aws_access_key_id": "deployment-access-key", + "aws_secret_access_key": "deployment-secret", + "aws_region_name": "eu-west-1", + "aws_batch_role_arn": "arn:aws:iam::123456789012:role/bedrock-batch", + }, + ) + + assert result["url"].startswith("https://bedrock.eu-west-1.amazonaws.com/") + assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key" + assert captured["optional_params"]["aws_secret_access_key"] == "deployment-secret" + assert captured["optional_params"]["aws_region_name"] == "eu-west-1" + + +def test_bedrock_batch_retrieval_signing_uses_deployment_credentials(monkeypatch): + config = BedrockBatchesConfig() + captured = {} + + def capture_signing(**kwargs): + captured.update(kwargs) + return {}, b"" + + monkeypatch.setattr(config.common_utils, "sign_aws_request", capture_signing) + + result = config.transform_retrieve_batch_request( + batch_id="arn:aws:bedrock:eu-west-1:123456789012:model-invocation-job/job-1", + optional_params={}, + litellm_params={ + "aws_access_key_id": "deployment-access-key", + "aws_secret_access_key": "deployment-secret", + "aws_region_name": "eu-west-1", + }, + ) + + assert result["url"].startswith("https://bedrock.eu-west-1.amazonaws.com/") + assert captured["optional_params"]["aws_access_key_id"] == "deployment-access-key" + assert captured["optional_params"]["aws_secret_access_key"] == "deployment-secret" + assert captured["optional_params"]["aws_region_name"] == "eu-west-1"