mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(bedrock): use deployment credentials for AWS requests
This commit is contained in:
parent
281e52ac49
commit
9d69fdac72
3 changed files with 109 additions and 8 deletions
|
|
@ -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",
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
Loading…
Add table
Reference in a new issue