fix(bedrock): use deployment credentials for AWS requests

This commit is contained in:
daleselaji-dev 2026-08-07 13:14:56 +08:00
parent 281e52ac49
commit 9d69fdac72
3 changed files with 109 additions and 8 deletions

View file

@ -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",
)

View file

@ -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,
),
)

View file

@ -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"