From 44a5d7a47a949525ad013ecb35c9b2952faf5549 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 28 Aug 2026 18:25:14 -0700 Subject: [PATCH 1/3] fix(aws): forward aws_external_id in bedrock embeddings and sagemaker credential loading --- litellm/llms/bedrock/embed/embedding.py | 2 + litellm/llms/sagemaker/chat/handler.py | 2 + litellm/llms/sagemaker/completion/handler.py | 2 + .../bedrock/embed/test_bedrock_embedding.py | 48 +++++++++++++++++++ .../sagemaker/test_sagemaker_chat_handler.py | 48 +++++++++++++++++++ .../test_sagemaker_completion_handler.py | 47 ++++++++++++++++++ 6 files changed, 149 insertions(+) create mode 100644 tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 082bf7ee2d9..ddc93228b0a 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -58,6 +58,7 @@ class BedrockEmbedding(BaseAWSLLM): aws_profile_name: Final = optional_params.pop("aws_profile_name", None) aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None) + aws_external_id: Final = optional_params.pop("aws_external_id", None) ### SET REGION NAME ### if aws_region_name is None: @@ -84,6 +85,7 @@ class BedrockEmbedding(BaseAWSLLM): aws_role_name=aws_role_name, aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, ) return credentials, aws_region_name diff --git a/litellm/llms/sagemaker/chat/handler.py b/litellm/llms/sagemaker/chat/handler.py index b3e9ed671fc..79837170624 100644 --- a/litellm/llms/sagemaker/chat/handler.py +++ b/litellm/llms/sagemaker/chat/handler.py @@ -34,6 +34,7 @@ class SagemakerChatHandler(BaseAWSLLM): optional_params.pop("aws_bedrock_runtime_endpoint", None) # https://bedrock-runtime.{region_name}.amazonaws.com aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None) + aws_external_id: Final = optional_params.pop("aws_external_id", None) ### SET REGION NAME ### if aws_region_name is None: @@ -60,6 +61,7 @@ class SagemakerChatHandler(BaseAWSLLM): aws_role_name=aws_role_name, aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, ) return credentials, aws_region_name diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index 84cad56f0d4..91941075c4b 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -57,6 +57,7 @@ class SagemakerLLM(BaseAWSLLM): optional_params.pop("aws_bedrock_runtime_endpoint", None) # https://bedrock-runtime.{region_name}.amazonaws.com aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None) + aws_external_id: Final = optional_params.pop("aws_external_id", None) ### SET REGION NAME ### if aws_region_name is None: @@ -83,6 +84,7 @@ class SagemakerLLM(BaseAWSLLM): aws_role_name=aws_role_name, aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, ) return credentials, aws_region_name diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index 114e473be98..65ae719d021 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -985,3 +985,51 @@ def test_bedrock_cohere_embedding_types_wrapped_as_list( assert "embedding_types" in request_body assert request_body["embedding_types"] == expected_embedding_types assert isinstance(request_body["embedding_types"], list) + + +def test_load_credentials_assumes_role_with_external_id(monkeypatch): + """A trust policy requiring sts:ExternalId must be satisfied by the deployment's aws_external_id.""" + import datetime + + import boto3 + from botocore.exceptions import ClientError + + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + + monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False) + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if params.get("ExternalId") != "external-id-embed": + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIAEMBEDROLEKEY", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_access_key_id": "AKIAEMBEDCALLERKEY", + "aws_secret_access_key": "pod-caller-secret", + "aws_region_name": "us-east-1", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-embed-role", + "aws_session_name": "litellm-embed-session", + "aws_external_id": "external-id-embed", + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + credentials, aws_region_name = BedrockEmbedding()._load_credentials(optional_params) + + assert credentials.access_key == "ASIAEMBEDROLEKEY" + assert credentials.token == "assumed-session-token" + assert aws_region_name == "us-east-1" + assert "aws_external_id" not in optional_params diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py new file mode 100644 index 00000000000..aa1a59d0e5c --- /dev/null +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py @@ -0,0 +1,48 @@ +import datetime +from unittest.mock import patch + +import boto3 +from botocore.exceptions import ClientError + +from litellm.llms.sagemaker.chat.handler import SagemakerChatHandler + + +def test_load_credentials_assumes_role_with_external_id(monkeypatch): + """A trust policy requiring sts:ExternalId must be satisfied by the deployment's aws_external_id.""" + monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False) + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if params.get("ExternalId") != "external-id-sm-chat": + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIASMCHATROLEKEY", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_access_key_id": "AKIASMCHATCALLERKEY", + "aws_secret_access_key": "pod-caller-secret", + "aws_region_name": "us-east-1", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-sm-chat-role", + "aws_session_name": "litellm-sm-chat-session", + "aws_external_id": "external-id-sm-chat", + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + credentials, aws_region_name = SagemakerChatHandler()._load_credentials(optional_params) + + assert credentials.access_key == "ASIASMCHATROLEKEY" + assert credentials.token == "assumed-session-token" + assert aws_region_name == "us-east-1" + assert "aws_external_id" not in optional_params diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py index 1cb27b7cf5f..881bac096b1 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py @@ -172,3 +172,50 @@ async def test_async_native_streaming_forwards_each_frame_incrementally(): assert texts == [f"token{i} " for i in range(len(frames))] assert consumed_at_token == list(range(1, len(frames) + 1)) + + +def test_load_credentials_assumes_role_with_external_id(monkeypatch): + """A trust policy requiring sts:ExternalId must be satisfied by the deployment's aws_external_id.""" + import datetime + + import boto3 + from botocore.exceptions import ClientError + from unittest.mock import patch + + monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False) + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if params.get("ExternalId") != "external-id-sm-completion": + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIASMCOMPROLEKEY", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_access_key_id": "AKIASMCOMPCALLERKEY", + "aws_secret_access_key": "pod-caller-secret", + "aws_region_name": "us-east-1", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-sm-completion-role", + "aws_session_name": "litellm-sm-completion-session", + "aws_external_id": "external-id-sm-completion", + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + credentials, aws_region_name = SagemakerLLM()._load_credentials(optional_params) + + assert credentials.access_key == "ASIASMCOMPROLEKEY" + assert credentials.token == "assumed-session-token" + assert aws_region_name == "us-east-1" + assert "aws_external_id" not in optional_params From ad8c1457d18cbe6c069fbd25a565a4cdf943efb7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 01:21:59 -0700 Subject: [PATCH 2/3] fix(aws): build every AWS endpoint and ARN from the region partition Adds litellm/litellm_core_utils/aws_partition.py mapping a region to its AWS partition (aws, aws-cn, aws-us-gov, and the iso partitions), its DNS suffix, and its ARN prefix, and uses it at every AWS host and ARN build site: bedrock (runtime, agent, agentcore, legacy client, batches, files, realtime), sagemaker, polly, secrets manager, s3 log uploads, bedrock passthrough routes, and rag ingestion. ARN detection now accepts arn:aws-cn: and arn:aws-us-gov: prefixes. STS region resolution now falls back to the configured aws_region_name after the aws_sts_endpoint host and the AWS_REGION/AWS_DEFAULT_REGION env vars, so cn and gov role assumption no longer silently signs against us-west-2. A partition sweep test walks every endpoint builder with cn regions and asserts no amazonaws.com host or arn:aws: prefix comes out, plus an AST guard that fails on any new f-string hardcoding either literal. --- basedpyright-code-budget.json | 4 +- litellm/integrations/s3_v2.py | 6 +- litellm/litellm_core_utils/aws_partition.py | 55 +++++ .../text_to_speech/transformation.py | 3 +- litellm/llms/bedrock/base_aws_llm.py | 31 ++- litellm/llms/bedrock/batches/handler.py | 5 +- .../llms/bedrock/batches/transformation.py | 16 +- .../bedrock/chat/agentcore/transformation.py | 3 +- litellm/llms/bedrock/common_utils.py | 31 +-- litellm/llms/bedrock/files/transformation.py | 8 +- litellm/llms/bedrock/realtime/handler.py | 3 +- litellm/llms/sagemaker/chat/handler.py | 6 +- litellm/llms/sagemaker/chat/transformation.py | 6 +- litellm/llms/sagemaker/completion/handler.py | 6 +- litellm/passthrough/utils.py | 3 +- .../llm_passthrough_endpoints.py | 7 +- litellm/rag/ingestion/bedrock_ingestion.py | 29 ++- .../secret_managers/aws_secret_manager_v2.py | 16 +- ruff-strict-budget.json | 2 +- tests/test_litellm/integrations/test_s3_v2.py | 29 +++ .../litellm_core_utils/test_aws_partition.py | 202 ++++++++++++++++++ .../bedrock/batches/test_transformation.py | 34 +++ .../llms/bedrock/test_base_aws_llm.py | 133 +++++++++++- .../test_sagemaker_chat_transformation.py | 52 +++++ .../test_llm_pass_through_endpoints.py | 10 +- ...test_passthrough_endpoints_common_utils.py | 10 + .../test_aws_secret_manager_v2.py | 53 +++++ type-discipline-budget.json | 4 +- 28 files changed, 697 insertions(+), 70 deletions(-) create mode 100644 litellm/litellm_core_utils/aws_partition.py create mode 100644 tests/test_litellm/litellm_core_utils/test_aws_partition.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 122bd82c657..9609133a7d1 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,7 +3,7 @@ "limit": 18483 }, "reportArgumentType": { - "limit": 2557 + "limit": 2552 }, "reportAssignmentType": { "limit": 319 @@ -111,7 +111,7 @@ "limit": 19827 }, "reportUnknownVariableType": { - "limit": 30348 + "limit": 30346 }, "reportUnnecessaryCast": { "limit": 117 diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 10b4c0dd433..9f6ae72fb3a 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -17,6 +17,7 @@ import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS from litellm.integrations.s3 import get_s3_object_key, resolve_sse_params +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -222,7 +223,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): protocol: Final = "https://" if self.s3_endpoint_url.startswith("https://") else "http://" return f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{encoded_key}" return f"{self.s3_endpoint_url}/{self.s3_bucket_name}/{encoded_key}" - return f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{encoded_key}" + return ( + f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}." + f"{get_aws_dns_suffix(self.s3_region_name)}/{encoded_key}" + ) def _sse_headers(self) -> Mapping[str, str]: candidates: Final = { diff --git a/litellm/litellm_core_utils/aws_partition.py b/litellm/litellm_core_utils/aws_partition.py new file mode 100644 index 00000000000..f8ca3aa4473 --- /dev/null +++ b/litellm/litellm_core_utils/aws_partition.py @@ -0,0 +1,55 @@ +import re +from types import MappingProxyType +from typing import Final, NamedTuple + + +class AwsPartition(NamedTuple): + partition: str + dns_suffix: str + + +_COMMERCIAL_PARTITION: Final = AwsPartition(partition="aws", dns_suffix="amazonaws.com") + +_PARTITIONS_BY_REGION_PREFIX: Final = MappingProxyType( + { + "cn-": AwsPartition(partition="aws-cn", dns_suffix="amazonaws.com.cn"), + "us-gov-": AwsPartition(partition="aws-us-gov", dns_suffix="amazonaws.com"), + "us-isob-": AwsPartition(partition="aws-iso-b", dns_suffix="sc2s.sgov.gov"), + "us-isof-": AwsPartition(partition="aws-iso-f", dns_suffix="csp.hci.ic.gov"), + "us-iso-": AwsPartition(partition="aws-iso", dns_suffix="c2s.ic.gov"), + "eu-isoe-": AwsPartition(partition="aws-iso-e", dns_suffix="cloud.adc-e.uk"), + } +) + +_BEDROCK_ARN_PATTERN: Final = re.compile(r"arn:aws(?:-[a-z0-9-]+)?:bedrock") +_BEDROCK_ARN_PREFIX_PATTERN: Final = re.compile(r"\Aarn:aws(?:-[a-z0-9-]+)?:bedrock:") +_AWS_ARN_PATTERN: Final = re.compile(r"arn:aws(?:-[a-z0-9-]+)?:") + + +def get_aws_partition(aws_region_name: str | None) -> AwsPartition: + if not aws_region_name: + return _COMMERCIAL_PARTITION + return next( + (partition for prefix, partition in _PARTITIONS_BY_REGION_PREFIX.items() if aws_region_name.startswith(prefix)), + _COMMERCIAL_PARTITION, + ) + + +def get_aws_dns_suffix(aws_region_name: str | None) -> str: + return get_aws_partition(aws_region_name).dns_suffix + + +def get_aws_arn_prefix(aws_region_name: str | None) -> str: + return f"arn:{get_aws_partition(aws_region_name).partition}:" + + +def contains_bedrock_arn(value: str) -> bool: + return _BEDROCK_ARN_PATTERN.search(value) is not None + + +def is_bedrock_arn(value: str) -> bool: + return _BEDROCK_ARN_PREFIX_PATTERN.match(value) is not None + + +def contains_aws_arn(value: str) -> bool: + return _AWS_ARN_PATTERN.search(value) is not None diff --git a/litellm/llms/aws_polly/text_to_speech/transformation.py b/litellm/llms/aws_polly/text_to_speech/transformation.py index 68630335ca7..8f96f80d15e 100644 --- a/litellm/llms/aws_polly/text_to_speech/transformation.py +++ b/litellm/llms/aws_polly/text_to_speech/transformation.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final, Union import httpx +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, TextToSpeechRequestData, @@ -238,7 +239,7 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): return api_base.rstrip("/") + "/v1/speech" aws_region_name: Final = litellm_params.get("aws_region_name", self.DEFAULT_REGION) - return f"https://polly.{aws_region_name}.amazonaws.com/v1/speech" + return f"https://polly.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/v1/speech" def is_ssml_input(self, input: str) -> bool: """ diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 4332848e545..852cfaa24f2 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -23,6 +23,7 @@ from litellm.constants import ( BEDROCK_MAX_POLICY_SIZE, STS_CREDENTIAL_EXPIRY_SAFETY_MARGIN_SECONDS, ) +from litellm.litellm_core_utils.aws_partition import contains_bedrock_arn, get_aws_dns_suffix from litellm.litellm_core_utils.dd_tracing import tracer from litellm.secret_managers.main import get_secret, get_secret_str @@ -348,7 +349,7 @@ class BaseAWSLLM: def _get_aws_region_from_model_arn(self, model: str | None) -> str | None: try: # First check if the string contains the expected prefix - if not isinstance(model, str) or "arn:aws:bedrock" not in model: + if not isinstance(model, str) or not contains_bedrock_arn(model): return None # Split the ARN and check if we have enough parts @@ -625,24 +626,29 @@ class BaseAWSLLM: return match.group(1) if match else None @staticmethod - def _resolve_sts_region(aws_sts_endpoint: str | None = None) -> str | None: - """STS signing region: parsed from aws_sts_endpoint else AWS_REGION / AWS_DEFAULT_REGION.""" + def _resolve_sts_region( + aws_sts_endpoint: str | None = None, + aws_region_name: str | None = None, + ) -> str | None: + """STS signing region: parsed from aws_sts_endpoint, else AWS_REGION / AWS_DEFAULT_REGION, else the configured aws_region_name.""" return ( BaseAWSLLM._parse_sts_region_from_endpoint(aws_sts_endpoint) or os.getenv("AWS_REGION") or os.getenv("AWS_DEFAULT_REGION") + or aws_region_name ) def _build_sts_client_kwargs( self, aws_sts_endpoint: str | None = None, ssl_verify: bool | str | None = None, + aws_region_name: str | None = None, ) -> dict: """STS client kwargs with aligned endpoint_url and region_name (SigV4).""" kwargs: Final[dict] = {"verify": self._get_ssl_verify(ssl_verify)} if aws_sts_endpoint is not None: kwargs["endpoint_url"] = aws_sts_endpoint - sts_region: Final = self._resolve_sts_region(aws_sts_endpoint) + sts_region: Final = self._resolve_sts_region(aws_sts_endpoint, aws_region_name) if sts_region is not None: kwargs["region_name"] = sts_region return kwargs @@ -837,6 +843,7 @@ class BaseAWSLLM: sts_client_kwargs: Final = self._build_sts_client_kwargs( aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, + aws_region_name=aws_region_name, ) with tracer.trace("boto3.client(sts)"): @@ -948,6 +955,7 @@ class BaseAWSLLM: aws_external_id: str | None = None, aws_sts_endpoint: str | None = None, ssl_verify: bool | str | None = None, + aws_region_name: str | None = None, ) -> dict: """Handle cross-account role assumption for IRSA.""" import boto3 @@ -961,6 +969,7 @@ class BaseAWSLLM: irsa_sts_kwargs: Final = self._build_sts_client_kwargs( aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, + aws_region_name=aws_region_name, ) # Create an STS client without credentials @@ -1017,6 +1026,7 @@ class BaseAWSLLM: aws_external_id: str | None = None, aws_sts_endpoint: str | None = None, ssl_verify: bool | str | None = None, + aws_region_name: str | None = None, ) -> dict: """Handle same-account role assumption for IRSA.""" import boto3 @@ -1024,6 +1034,7 @@ class BaseAWSLLM: irsa_sts_kwargs: Final = self._build_sts_client_kwargs( aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, + aws_region_name=aws_region_name, ) verbose_logger.debug("Same account role assumption, using automatic IRSA") @@ -1153,6 +1164,7 @@ class BaseAWSLLM: aws_external_id, aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, + aws_region_name=aws_region_name, ) else: sts_response = self._handle_irsa_same_account( @@ -1161,6 +1173,7 @@ class BaseAWSLLM: aws_external_id, aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, + aws_region_name=aws_region_name, ) return self._extract_credentials_and_ttl(sts_response) @@ -1182,6 +1195,7 @@ class BaseAWSLLM: sts_client_kwargs: Final = self._build_sts_client_kwargs( aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, + aws_region_name=aws_region_name, ) if aws_access_key_id is None and aws_secret_access_key is None: with tracer.trace("boto3.client(sts)"): @@ -1363,14 +1377,15 @@ class BaseAWSLLM: """ Select the default endpoint url based on the endpoint type - Default endpoint url is https://bedrock-runtime.{aws_region_name}.amazonaws.com + Default endpoint url is https://bedrock-runtime.{aws_region_name}.{partition dns suffix} """ + dns_suffix: Final = get_aws_dns_suffix(aws_region_name) if endpoint_type == "agent": - return f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com" + return f"https://bedrock-agent-runtime.{aws_region_name}.{dns_suffix}" elif endpoint_type == "agentcore": - return f"https://bedrock-agentcore.{aws_region_name}.amazonaws.com" + return f"https://bedrock-agentcore.{aws_region_name}.{dns_suffix}" else: - return f"https://bedrock-runtime.{aws_region_name}.amazonaws.com" + return f"https://bedrock-runtime.{aws_region_name}.{dns_suffix}" def _get_boto_credentials_from_optional_params( self, optional_params: dict, model: str | None = None diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 6efdd17f98d..232f300b460 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final, cast from openai.types.batch import BatchRequestCounts from openai.types.batch import Metadata as OpenAIBatchMetadata +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.types.utils import LiteLLMBatch if TYPE_CHECKING: @@ -323,7 +324,9 @@ class BedrockBatchesHandler: api_key="", additional_args={ "complete_input_dict": {"jobIdentifier": batch_id}, - "api_base": (f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{url_path_id}"), + "api_base": ( + f"https://bedrock.{region}.{get_aws_dns_suffix(region)}/model-invocation-job/{url_path_id}" + ), }, ) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 04f395f2bf1..1bb3d6f5f5a 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -6,6 +6,7 @@ from typing import Any, Final, Literal, cast from httpx import Headers, Response from pydantic import TypeAdapter, ValidationError +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix, is_bedrock_arn from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_BATCH_PREFIX, ) @@ -138,8 +139,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): 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 - bedrock_endpoint: Final = f"https://bedrock.{aws_region_name}.amazonaws.com/model-invocation-job" + # Format: https://bedrock.{region}.{partition dns suffix}/model-invocation-job + bedrock_endpoint: Final = ( + f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/model-invocation-job" + ) return bedrock_endpoint @@ -238,8 +241,9 @@ 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 = merge_bedrock_aws_request_params(litellm_params, optional_params) + aws_region_name: Final = self._get_aws_region_name(request_params, model) endpoint_url: Final = ( - f"https://bedrock.{self._get_aws_region_name(request_params, model)}.amazonaws.com/model-invocation-job" + f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/model-invocation-job" ) signed_headers, signed_data = self.common_utils.sign_aws_request( service_name="bedrock", @@ -371,7 +375,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): """ # For Bedrock, batch_id should be the full job ARN # The GetModelInvocationJob API expects the full ARN as the identifier - if not batch_id.startswith("arn:aws:bedrock:"): + if not is_bedrock_arn(batch_id): raise ValueError(f"Invalid batch_id format. Expected ARN, got: {batch_id}") # Extract the job identifier from the ARN - use the full ARN path part @@ -390,7 +394,9 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): import urllib.parse as _ul encoded_arn: Final = _ul.quote(batch_id, safe="") - endpoint_url: Final = f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{encoded_arn}" + endpoint_url: Final = ( + f"https://bedrock.{region}.{get_aws_dns_suffix(region)}/model-invocation-job/{encoded_arn}" + ) # Use common utility for AWS signing request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 4a2db621421..c5ba0050be8 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -13,6 +13,7 @@ import httpx from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) @@ -97,7 +98,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): if aws_bedrock_runtime_endpoint: base_url = aws_bedrock_runtime_endpoint else: - base_url = f"https://bedrock-agentcore.{region}.amazonaws.com" + base_url = f"https://bedrock-agentcore.{region}.{get_aws_dns_suffix(region)}" # Based on boto3 client.invoke_agent_runtime, the path is: # /runtimes/{URL-ENCODED-ARN}/invocations?qualifier= diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 4ad20772ed0..72e3cc1b326 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -21,6 +21,7 @@ import httpx import litellm from litellm import verbose_logger +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -434,15 +435,15 @@ def init_bedrock_client( ssl_verify: Final = _get_bedrock_client_ssl_verify() ### SET REGION NAME - if region_name: - pass - elif aws_region_name: - region_name = aws_region_name - elif litellm_aws_region_name: - region_name = litellm_aws_region_name - elif standard_aws_region_name: - region_name = standard_aws_region_name - else: + resolved_region_name: Final = next( + ( + candidate + for candidate in (region_name, aws_region_name, litellm_aws_region_name, standard_aws_region_name) + if isinstance(candidate, str) and candidate + ), + None, + ) + if resolved_region_name is None: raise BedrockError( message="AWS region not set: set AWS_REGION_NAME or AWS_REGION env variable or in .env file", status_code=401, @@ -455,7 +456,7 @@ def init_bedrock_client( elif env_aws_bedrock_runtime_endpoint: endpoint_url = env_aws_bedrock_runtime_endpoint else: - endpoint_url = f"https://bedrock-runtime.{region_name}.amazonaws.com" + endpoint_url = f"https://bedrock-runtime.{resolved_region_name}.{get_aws_dns_suffix(resolved_region_name)}" import boto3 @@ -492,7 +493,7 @@ def init_bedrock_client( aws_access_key_id=sts_response["Credentials"]["AccessKeyId"], aws_secret_access_key=sts_response["Credentials"]["SecretAccessKey"], aws_session_token=sts_response["Credentials"]["SessionToken"], - region_name=region_name, + region_name=resolved_region_name, endpoint_url=endpoint_url, config=config, verify=ssl_verify, @@ -513,7 +514,7 @@ def init_bedrock_client( aws_access_key_id=sts_response["Credentials"]["AccessKeyId"], aws_secret_access_key=sts_response["Credentials"]["SecretAccessKey"], aws_session_token=sts_response["Credentials"]["SessionToken"], - region_name=region_name, + region_name=resolved_region_name, endpoint_url=endpoint_url, config=config, verify=ssl_verify, @@ -526,7 +527,7 @@ def init_bedrock_client( service_name="bedrock-runtime", aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, - region_name=region_name, + region_name=resolved_region_name, endpoint_url=endpoint_url, config=config, verify=ssl_verify, @@ -536,7 +537,7 @@ def init_bedrock_client( client = boto3.Session(profile_name=aws_profile_name).client( service_name="bedrock-runtime", - region_name=region_name, + region_name=resolved_region_name, endpoint_url=endpoint_url, config=config, verify=ssl_verify, @@ -547,7 +548,7 @@ def init_bedrock_client( client = boto3.client( service_name="bedrock-runtime", - region_name=region_name, + region_name=resolved_region_name, endpoint_url=endpoint_url, config=config, verify=ssl_verify, diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index b034696594a..f442608a288 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -20,6 +20,7 @@ from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.constants import BEDROCK_INVOKE_PROVIDERS_LITERAL from litellm.files.utils import FilesAPIUtils +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_BATCH_PREFIX, BEDROCK_MANAGED_S3_PREFIXES, @@ -413,7 +414,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # S3 endpoint URL format s3_endpoint_url: Final = ( - request_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}.{get_aws_dns_suffix(aws_region_name)}" ).rstrip("/") return f"{s3_endpoint_url}/{bucket_name}/{encoded_object_name}" @@ -1249,7 +1251,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): region_params: Final[dict[str, str | None]] = {"aws_region_name": region_preference} aws_region_name: Final = self._get_aws_region_name(optional_params=region_params, model="") - s3_endpoint_url = (request_params.s3_endpoint_url or f"https://s3.{aws_region_name}.amazonaws.com").rstrip("/") + s3_endpoint_url = ( + request_params.s3_endpoint_url or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" + ).rstrip("/") url: Final = f"{s3_endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}" litellm_params[S3_SIGNED_GET_HEADERS_PARAM] = self._sign_s3_get_request( diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 3eeb3cb9fc6..7b5de304b5e 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -12,6 +12,7 @@ from typing import Any, Final from pydantic import TypeAdapter from litellm._logging import _redact_string, verbose_proxy_logger +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from ..base_aws_llm import BaseAWSLLM @@ -81,7 +82,7 @@ class BedrockRealtime(BaseAWSLLM): elif aws_bedrock_runtime_endpoint is not None: endpoint_uri = aws_bedrock_runtime_endpoint else: - endpoint_uri = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com" + endpoint_uri = f"https://bedrock-runtime.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" verbose_proxy_logger.debug("Bedrock Realtime: Connecting to %s with model %s", endpoint_uri, model) diff --git a/litellm/llms/sagemaker/chat/handler.py b/litellm/llms/sagemaker/chat/handler.py index b3e9ed671fc..2b013b07af0 100644 --- a/litellm/llms/sagemaker/chat/handler.py +++ b/litellm/llms/sagemaker/chat/handler.py @@ -5,6 +5,7 @@ from typing import Final import httpx +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.utils import ModelResponse, get_secret @@ -79,10 +80,11 @@ class SagemakerChatHandler(BaseAWSLLM): raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") sigv4: Final = SigV4Auth(credentials, "sagemaker", aws_region_name) + dns_suffix: Final = get_aws_dns_suffix(aws_region_name) if optional_params.get("stream") is True: - api_base = f"https://runtime.sagemaker.{aws_region_name}.amazonaws.com/endpoints/{model}/invocations-response-stream" + api_base = f"https://runtime.sagemaker.{aws_region_name}.{dns_suffix}/endpoints/{model}/invocations-response-stream" else: - api_base = f"https://runtime.sagemaker.{aws_region_name}.amazonaws.com/endpoints/{model}/invocations" + api_base = f"https://runtime.sagemaker.{aws_region_name}.{dns_suffix}/endpoints/{model}/invocations" sagemaker_base_url: Final = optional_params.get("sagemaker_base_url", None) if sagemaker_base_url is not None: diff --git a/litellm/llms/sagemaker/chat/transformation.py b/litellm/llms/sagemaker/chat/transformation.py index 37ddd813d6f..04995f32d97 100644 --- a/litellm/llms/sagemaker/chat/transformation.py +++ b/litellm/llms/sagemaker/chat/transformation.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._models import Headers +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -93,10 +94,11 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): model=model, model_id=None, ) + dns_suffix: Final = get_aws_dns_suffix(aws_region_name) if stream is True: - api_base = f"https://runtime.sagemaker.{aws_region_name}.amazonaws.com/endpoints/{model}/invocations-response-stream" + api_base = f"https://runtime.sagemaker.{aws_region_name}.{dns_suffix}/endpoints/{model}/invocations-response-stream" else: - api_base = f"https://runtime.sagemaker.{aws_region_name}.amazonaws.com/endpoints/{model}/invocations" + api_base = f"https://runtime.sagemaker.{aws_region_name}.{dns_suffix}/endpoints/{model}/invocations" sagemaker_base_url: Final = cast(str | None, optional_params.get("sagemaker_base_url")) if sagemaker_base_url is not None: diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index 84cad56f0d4..4df9e701470 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -8,6 +8,7 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import ( @@ -104,10 +105,11 @@ class SagemakerLLM(BaseAWSLLM): raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") sigv4: Final = SigV4Auth(credentials, "sagemaker", aws_region_name) + dns_suffix: Final = get_aws_dns_suffix(aws_region_name) if optional_params.get("stream") is True: - api_base = f"https://runtime.sagemaker.{aws_region_name}.amazonaws.com/endpoints/{model}/invocations-response-stream" + api_base = f"https://runtime.sagemaker.{aws_region_name}.{dns_suffix}/endpoints/{model}/invocations-response-stream" else: - api_base = f"https://runtime.sagemaker.{aws_region_name}.amazonaws.com/endpoints/{model}/invocations" + api_base = f"https://runtime.sagemaker.{aws_region_name}.{dns_suffix}/endpoints/{model}/invocations" sagemaker_base_url: Final = optional_params.get("sagemaker_base_url", None) if sagemaker_base_url is not None: diff --git a/litellm/passthrough/utils.py b/litellm/passthrough/utils.py index df39b8fad48..7eb14fcc118 100644 --- a/litellm/passthrough/utils.py +++ b/litellm/passthrough/utils.py @@ -6,6 +6,7 @@ import httpx from litellm._logging import verbose_logger from litellm.constants import PASS_THROUGH_HEADER_PREFIX +from litellm.litellm_core_utils.aws_partition import contains_aws_arn # Headers that must not be overwritten via the x-pass- forwarding mechanism. # Includes standard credential/auth headers and protocol-level headers that @@ -126,7 +127,7 @@ class CommonUtils: import re # Early exit: if no ARN detected, return unchanged - if "arn:aws:" not in endpoint: + if not contains_aws_arn(endpoint): return endpoint # Handle all patterns in one go - more efficient and cleaner diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 9a3bc82c6fa..525e7099b89 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -26,6 +26,7 @@ from litellm.constants import ( ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS, BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES, ) +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._types import * @@ -1057,7 +1058,7 @@ async def bedrock_proxy_route( except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") - aws_region_name: Final = litellm.utils.get_secret(secret_name="AWS_REGION_NAME") + aws_region_name: Final = get_secret_str(secret_name="AWS_REGION_NAME") if not _is_bedrock_agent_runtime_route(endpoint=endpoint): return await bedrock_llm_proxy_route( endpoint=endpoint, @@ -1072,7 +1073,7 @@ async def bedrock_proxy_route( detail="bedrock-agent-runtime pass-through is disabled on this proxy.", ) - base_target_url: Final = f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com" + base_target_url: Final = f"https://bedrock-agent-runtime.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" encoded_endpoint = httpx.URL(endpoint).path # Ensure endpoint starts with '/' for proper URL construction @@ -1205,7 +1206,7 @@ async def comprehend_medical_proxy_route( "X-Amz-Target": f"{COMPREHEND_MEDICAL_TARGET_PREFIX}.{operation}", } ) - target_url: Final = f"https://comprehendmedical.{aws_region_name}.amazonaws.com/" + target_url: Final = f"https://comprehendmedical.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/" _request: Final = AWSRequest(method="POST", url=target_url, data=json.dumps(data), headers=headers) sigv4.add_auth(_request) prepped: Final = _request.prepare() diff --git a/litellm/rag/ingestion/bedrock_ingestion.py b/litellm/rag/ingestion/bedrock_ingestion.py index f721c204318..3d7056f8176 100644 --- a/litellm/rag/ingestion/bedrock_ingestion.py +++ b/litellm/rag/ingestion/bedrock_ingestion.py @@ -17,6 +17,7 @@ import uuid from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_logger +from litellm.litellm_core_utils.aws_partition import get_aws_arn_prefix from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion @@ -52,11 +53,12 @@ def _normalize_principal_arn(caller_arn: str, account_id: str) -> str: """ if ":assumed-role/" in caller_arn: # Extract role name from assumed-role ARN - # Format: arn:aws:sts::ACCOUNT:assumed-role/ROLE-NAME/SESSION-NAME + # Format: arn:PARTITION:sts::ACCOUNT:assumed-role/ROLE-NAME/SESSION-NAME + partition: Final = caller_arn.split(":")[1] parts: Final = caller_arn.split("/") if len(parts) >= 2: role_name: Final = parts[1] - return f"arn:aws:iam::{account_id}:role/{role_name}" + return f"arn:{partition}:iam::{account_id}:role/{role_name}" return caller_arn @@ -294,7 +296,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): normalized_caller_arn: Final = _normalize_principal_arn(caller_arn, account_id) verbose_logger.debug("Caller ARN: %s, Normalized: %s", caller_arn, normalized_caller_arn) - principals = [f"arn:aws:iam::{account_id}:root", normalized_caller_arn] + principals = [f"{get_aws_arn_prefix(self.aws_region_name)}iam::{account_id}:root", normalized_caller_arn] # Deduplicate in case caller is root principals = list(set(principals)) @@ -454,7 +456,10 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): "Condition": { "StringEquals": {"aws:SourceAccount": account_id}, "ArnLike": { - "aws:SourceArn": f"arn:aws:bedrock:{self.aws_region_name}:{account_id}:knowledge-base/*" + "aws:SourceArn": ( + f"{get_aws_arn_prefix(self.aws_region_name)}bedrock:" + f"{self.aws_region_name}:{account_id}:knowledge-base/*" + ) }, }, } @@ -475,7 +480,10 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): { "Effect": "Allow", "Action": ["bedrock:InvokeModel"], - "Resource": [f"arn:aws:bedrock:{self.aws_region_name}::foundation-model/{self.embedding_model}"], + "Resource": [ + f"{get_aws_arn_prefix(self.aws_region_name)}bedrock:" + f"{self.aws_region_name}::foundation-model/{self.embedding_model}" + ], }, { "Effect": "Allow", @@ -486,8 +494,8 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): "Effect": "Allow", "Action": ["s3:GetObject", "s3:ListBucket"], "Resource": [ - f"arn:aws:s3:::{self.s3_bucket}", - f"arn:aws:s3:::{self.s3_bucket}/*", + f"{get_aws_arn_prefix(self.aws_region_name)}s3:::{self.s3_bucket}", + f"{get_aws_arn_prefix(self.aws_region_name)}s3:::{self.s3_bucket}/*", ], }, ], @@ -517,7 +525,10 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): knowledgeBaseConfiguration={ "type": "VECTOR", "vectorKnowledgeBaseConfiguration": { - "embeddingModelArn": f"arn:aws:bedrock:{self.aws_region_name}::foundation-model/{self.embedding_model}", + "embeddingModelArn": ( + f"{get_aws_arn_prefix(self.aws_region_name)}bedrock:" + f"{self.aws_region_name}::foundation-model/{self.embedding_model}" + ), }, }, storageConfiguration={ @@ -562,7 +573,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): dataSourceConfiguration={ "type": "S3", "s3Configuration": { - "bucketArn": f"arn:aws:s3:::{self.s3_bucket}", + "bucketArn": f"{get_aws_arn_prefix(self.aws_region_name)}s3:::{self.s3_bucket}", "inclusionPrefixes": [self.s3_prefix], }, }, diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index 38a2ddd0bfc..2c7f1f8389d 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -22,12 +22,14 @@ import httpx import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, ) from litellm.proxy._types import KeyManagementSystem +from litellm.secret_managers.main import get_secret_str from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.secret_managers.main import KeyManagementSettings @@ -556,13 +558,15 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(optional_params) - # Get endpoint - _, endpoint_url = self.get_runtime_endpoint( - api_base=None, - aws_bedrock_runtime_endpoint=boto3_credentials_info.aws_bedrock_runtime_endpoint, - aws_region_name=boto3_credentials_info.aws_region_name, + region_name: Final = boto3_credentials_info.aws_region_name + explicit_runtime_endpoint: Final = boto3_credentials_info.aws_bedrock_runtime_endpoint or get_secret_str( + "AWS_BEDROCK_RUNTIME_ENDPOINT" + ) + endpoint_url: Final = ( + explicit_runtime_endpoint.replace("bedrock-runtime", "secretsmanager") + if explicit_runtime_endpoint + else f"https://secretsmanager.{region_name}.{get_aws_dns_suffix(region_name)}" ) - endpoint_url = endpoint_url.replace("bedrock-runtime", "secretsmanager") # Use provided request_data if available, otherwise build default data if request_data: diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 0418eeaac8f..8b7e1762126 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -78,7 +78,7 @@ "limit": 1 }, "C901": { - "limit": 312 + "limit": 311 }, "D419": { "limit": 6 diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 0f307450b50..e4286fc43fa 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1856,3 +1856,32 @@ async def test_download_percent_encodes_reserved_characters_in_object_key(s3_obj body=None, headers=call.kwargs["headers"], ) + + +def _s3_logger_for_region(region_name: str) -> S3Logger: + logger = S3Logger.__new__(S3Logger) + logger.s3_endpoint_url = None + logger.s3_bucket_name = "my-litellm-audit" + logger.s3_region_name = region_name + return logger + + +@pytest.mark.parametrize( + "region_name,expected_url", + [ + ( + "cn-northwest-1", + "https://my-litellm-audit.s3.cn-northwest-1.amazonaws.com.cn/2025-01-01/key.json", + ), + ( + "us-gov-west-1", + "https://my-litellm-audit.s3.us-gov-west-1.amazonaws.com/2025-01-01/key.json", + ), + ( + "us-east-1", + "https://my-litellm-audit.s3.us-east-1.amazonaws.com/2025-01-01/key.json", + ), + ], +) +def test_build_object_url_uses_partition_dns_suffix(region_name, expected_url): + assert _s3_logger_for_region(region_name)._build_object_url("2025-01-01/key.json") == expected_url diff --git a/tests/test_litellm/litellm_core_utils/test_aws_partition.py b/tests/test_litellm/litellm_core_utils/test_aws_partition.py new file mode 100644 index 00000000000..fcc3e370073 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_aws_partition.py @@ -0,0 +1,202 @@ +import ast +from pathlib import Path +from typing import Final +from urllib.parse import urlparse + +import pytest + +import litellm +from litellm.integrations.s3_v2 import S3Logger +from litellm.litellm_core_utils.aws_partition import ( + AwsPartition, + contains_aws_arn, + contains_bedrock_arn, + get_aws_arn_prefix, + get_aws_dns_suffix, + get_aws_partition, + is_bedrock_arn, +) +from litellm.llms.aws_polly.text_to_speech.transformation import AWSPollyTextToSpeechConfig +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig +from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig +from litellm.llms.bedrock.common_utils import init_bedrock_client +from litellm.llms.sagemaker.chat.transformation import SagemakerChatConfig + + +@pytest.mark.parametrize( + "region,partition,dns_suffix", + [ + ("us-east-1", "aws", "amazonaws.com"), + ("eu-central-1", "aws", "amazonaws.com"), + ("ap-southeast-1", "aws", "amazonaws.com"), + ("sa-east-1", "aws", "amazonaws.com"), + ("cn-north-1", "aws-cn", "amazonaws.com.cn"), + ("cn-northwest-1", "aws-cn", "amazonaws.com.cn"), + ("us-gov-west-1", "aws-us-gov", "amazonaws.com"), + ("us-gov-east-1", "aws-us-gov", "amazonaws.com"), + ("us-iso-east-1", "aws-iso", "c2s.ic.gov"), + ("us-isob-east-1", "aws-iso-b", "sc2s.sgov.gov"), + ("us-isof-south-1", "aws-iso-f", "csp.hci.ic.gov"), + ("eu-isoe-west-1", "aws-iso-e", "cloud.adc-e.uk"), + (None, "aws", "amazonaws.com"), + ("", "aws", "amazonaws.com"), + ], +) +def test_partition_lookup(region, partition, dns_suffix): + assert get_aws_partition(region) == AwsPartition(partition=partition, dns_suffix=dns_suffix) + assert get_aws_dns_suffix(region) == dns_suffix + assert get_aws_arn_prefix(region) == f"arn:{partition}:" + + +@pytest.mark.parametrize( + "value,expected", + [ + ("arn:aws:bedrock:us-west-2:123456789012:foundation-model/anthropic.claude-3", True), + ("arn:aws-cn:bedrock:cn-north-1:123456789012:inference-profile/p", True), + ("arn:aws-us-gov:bedrock:us-gov-west-1:123456789012:foundation-model/m", True), + ("bedrock/arn:aws-cn:bedrock:cn-north-1:123456789012:application-inference-profile/p", True), + ("arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/r", True), + ("anthropic.claude-3", False), + ("arn:aws:iam::123456789012:role/foo", False), + ], +) +def test_contains_bedrock_arn(value, expected): + assert contains_bedrock_arn(value) is expected + + +@pytest.mark.parametrize( + "value,expected", + [ + ("arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/j", True), + ("arn:aws-cn:bedrock:cn-north-1:123456789012:model-invocation-job/j", True), + ("arn:aws-us-gov:bedrock:us-gov-west-1:123456789012:model-invocation-job/j", True), + ("abc1234567", False), + ("bedrock/arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/j", False), + ("arn:aws:iam::123456789012:role/foo", False), + ], +) +def test_is_bedrock_arn(value, expected): + assert is_bedrock_arn(value) is expected + + +@pytest.mark.parametrize( + "value,expected", + [ + ("model/arn:aws:bedrock:us-east-1:123456789012:foundation-model/m/converse", True), + ("model/arn:aws-cn:bedrock:cn-north-1:123456789012:foundation-model/m/converse", True), + ("arn:aws-us-gov:bedrock:us-gov-west-1:123456789012:inference-profile/p", True), + ("model/anthropic.claude-3/converse", False), + ("arnaws:bedrock", False), + ], +) +def test_contains_aws_arn(value, expected): + assert contains_aws_arn(value) is expected + + +def _agentcore_model(region: str) -> str: + return f"agentcore/{get_aws_arn_prefix(region)}bedrock-agentcore:{region}:111122223333:runtime/my-agent" + + +def _s3_object_url(region: str) -> str: + logger = S3Logger.__new__(S3Logger) + logger.s3_endpoint_url = None + logger.s3_bucket_name = "audit-bucket" + logger.s3_region_name = region + return logger._build_object_url("2025-01-01/key.json") + + +ENDPOINT_BUILDERS: Final = { + "bedrock_runtime_default": lambda region: BaseAWSLLM()._select_default_endpoint_url("runtime", region), + "bedrock_agent_default": lambda region: BaseAWSLLM()._select_default_endpoint_url("agent", region), + "bedrock_agentcore_default": lambda region: BaseAWSLLM()._select_default_endpoint_url("agentcore", region), + "bedrock_get_runtime_endpoint": lambda region: BaseAWSLLM().get_runtime_endpoint(None, None, region)[0], + "bedrock_legacy_client": lambda region: init_bedrock_client( + region_name=region, + aws_access_key_id="test-key", + aws_secret_access_key="test-secret", + ).meta.endpoint_url, + "bedrock_batches": lambda region: BedrockBatchesConfig().get_complete_batch_url( + api_base=None, + api_key=None, + model="anthropic.claude-3", + optional_params={"aws_region_name": region}, + litellm_params={}, + data={"input_file_id": "s3://bucket/key.jsonl"}, + ), + "bedrock_agentcore_invoke": lambda region: AmazonAgentCoreConfig().get_complete_url( + api_base=None, + api_key=None, + model=_agentcore_model(region), + optional_params={}, + litellm_params={}, + ), + "polly": lambda region: AWSPollyTextToSpeechConfig().get_complete_url( + model="polly/neural", + api_base=None, + litellm_params={"aws_region_name": region}, + ), + "sagemaker_chat": lambda region: SagemakerChatConfig().get_complete_url( + api_base=None, + api_key=None, + model="my-endpoint", + optional_params={"aws_region_name": region}, + litellm_params={}, + stream=False, + ), + "sagemaker_chat_stream": lambda region: SagemakerChatConfig().get_complete_url( + api_base=None, + api_key=None, + model="my-endpoint", + optional_params={"aws_region_name": region}, + litellm_params={}, + stream=True, + ), + "s3_object_url": _s3_object_url, +} + + +@pytest.fixture(autouse=True) +def _clear_aws_env(monkeypatch): + for env_var in ("AWS_BEDROCK_RUNTIME_ENDPOINT", "AWS_REGION", "AWS_DEFAULT_REGION", "AWS_REGION_NAME"): + monkeypatch.delenv(env_var, raising=False) + + +@pytest.mark.parametrize("region", ["cn-north-1", "cn-northwest-1"]) +@pytest.mark.parametrize("builder_name", sorted(ENDPOINT_BUILDERS)) +def test_every_endpoint_builder_respects_cn_partition(builder_name, region): + url = ENDPOINT_BUILDERS[builder_name](region) + hostname = urlparse(url).hostname + assert hostname is not None + assert hostname.endswith(".amazonaws.com.cn"), url + assert not hostname.endswith("amazonaws.com"), url + assert "arn:aws:" not in url, url + + +@pytest.mark.parametrize("region", ["us-east-1", "us-gov-west-1"]) +@pytest.mark.parametrize("builder_name", sorted(ENDPOINT_BUILDERS)) +def test_every_endpoint_builder_keeps_amazonaws_com_outside_cn(builder_name, region): + url = ENDPOINT_BUILDERS[builder_name](region) + hostname = urlparse(url).hostname + assert hostname is not None + assert hostname.endswith(".amazonaws.com"), url + + +def _fstring_literal_offenders(needle: str) -> list[str]: + litellm_root = Path(litellm.__file__).parent + return [ + f"{path.relative_to(litellm_root)}: {part.value!r}" + for path in sorted(litellm_root.rglob("*.py")) + for node in ast.walk(ast.parse(path.read_text(encoding="utf-8"))) + if isinstance(node, ast.JoinedStr) + for part in node.values + if isinstance(part, ast.Constant) and isinstance(part.value, str) and needle in part.value + ] + + +def test_no_fstring_hardcodes_the_commercial_dns_suffix(): + assert _fstring_literal_offenders("amazonaws.com") == [] + + +def test_no_fstring_hardcodes_the_commercial_arn_prefix(): + assert _fstring_literal_offenders("arn:aws:") == [] diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index 87f9c506857..664cac6c145 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -785,3 +785,37 @@ class TestBedrockBatchesContract(BatchesConfigContractTests): expected_retrieve_batch_id = ARN expected_retrieve_status = "completed" + + +def test_get_complete_batch_url_cn_partition(config): + url = config.get_complete_batch_url( + api_base=None, + api_key=None, + model="anthropic.claude-3", + optional_params={"aws_region_name": "cn-north-1"}, + litellm_params={}, + data={"input_file_id": "s3://b/k"}, + ) + assert url == "https://bedrock.cn-north-1.amazonaws.com.cn/model-invocation-job" + + +@pytest.mark.parametrize( + "arn,expected_prefix", + [ + ( + "arn:aws-cn:bedrock:cn-north-1:123456789012:model-invocation-job/abc1234567", + "https://bedrock.cn-north-1.amazonaws.com.cn/model-invocation-job/", + ), + ( + "arn:aws-us-gov:bedrock:us-gov-west-1:123456789012:model-invocation-job/abc1234567", + "https://bedrock.us-gov-west-1.amazonaws.com/model-invocation-job/", + ), + ], +) +def test_retrieve_request_accepts_partition_arns(config, arn, expected_prefix): + with patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({"Authorization": "signed"}, b"") + result = config.transform_retrieve_batch_request( + batch_id=arn, optional_params={}, litellm_params={} + ) + assert result["url"].startswith(expected_prefix) diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 50e2b53c2b3..81321ce73ec 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -1223,7 +1223,7 @@ def test_different_roles_without_session_names_should_not_share_cache(): ({}, {"verify": True}), ( {"aws_region_name": "us-east-1"}, - {"verify": True}, + {"verify": True, "region_name": "us-east-1"}, ), ( {"aws_sts_endpoint": "https://sts.eu-west-1.amazonaws.com"}, @@ -1234,7 +1234,7 @@ def test_different_roles_without_session_names_should_not_share_cache(): }, ), ], - ids=["no_region_or_endpoint", "bedrock_region_ignored_for_sts", "explicit_sts_endpoint"], + ids=["no_region_or_endpoint", "configured_region_is_sts_fallback", "explicit_sts_endpoint"], ) def test_eks_irsa_ambient_credentials_used(role_kwargs, expected_client_kwargs): """ @@ -1418,6 +1418,132 @@ def test_build_sts_client_kwargs(env, aws_sts_endpoint, ssl_verify, expected): ) +@pytest.mark.parametrize( + "env,aws_sts_endpoint,aws_region_name,expected_region", + [ + ({}, None, "cn-north-1", "cn-north-1"), + ({"AWS_REGION": "eu-west-1"}, None, "cn-north-1", "eu-west-1"), + ({"AWS_DEFAULT_REGION": "ap-southeast-1"}, None, "cn-north-1", "ap-southeast-1"), + ({}, "https://sts.cn-north-1.amazonaws.com.cn", "us-east-1", "cn-north-1"), + ({}, None, None, None), + ], + ids=[ + "configured_region_fallback", + "env_region_beats_configured", + "env_default_region_beats_configured", + "cn_endpoint_beats_configured", + "nothing_configured", + ], +) +def test_resolve_sts_region_configured_region_fallback( + env, aws_sts_endpoint, aws_region_name, expected_region +): + with patch.dict(os.environ, env, clear=True): + assert ( + BaseAWSLLM._resolve_sts_region( + aws_sts_endpoint=aws_sts_endpoint, + aws_region_name=aws_region_name, + ) + == expected_region + ) + + +def test_build_sts_client_kwargs_configured_region_fallback(): + base_aws_llm = BaseAWSLLM() + with patch.dict(os.environ, {}, clear=True): + assert base_aws_llm._build_sts_client_kwargs(aws_region_name="cn-north-1") == { + "verify": True, + "region_name": "cn-north-1", + } + with patch.dict(os.environ, {"AWS_REGION": "eu-west-1"}, clear=True): + assert base_aws_llm._build_sts_client_kwargs(aws_region_name="cn-north-1") == { + "verify": True, + "region_name": "eu-west-1", + } + + +def test_assume_role_sts_client_uses_configured_cn_region(): + """arn:aws-cn roles must resolve against a cn STS endpoint, not the commercial default.""" + base_aws_llm = BaseAWSLLM() + mock_expiry = MagicMock() + mock_expiry.tzinfo = timezone.utc + time_diff = MagicMock() + time_diff.total_seconds.return_value = 3600 + mock_expiry.__sub__ = MagicMock(return_value=time_diff) + mock_sts_client = MagicMock() + mock_sts_client.assume_role.return_value = { + "Credentials": { + "AccessKeyId": "assumed-access-key", + "SecretAccessKey": "assumed-secret-key", + "SessionToken": "assumed-session-token", + "Expiration": mock_expiry, + } + } + + with patch.dict(os.environ, {}, clear=True): + with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client: + credentials, ttl = base_aws_llm._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name="arn:aws-cn:iam::2222222222222:role/LitellmBedrockRole", + aws_session_name="test-session", + aws_region_name="cn-north-1", + ) + mock_boto3_client.assert_called_with( + "sts", + region_name="cn-north-1", + verify=True, + ) + assert credentials.access_key == "assumed-access-key" + assert credentials.secret_key == "assumed-secret-key" + assert credentials.token == "assumed-session-token" + assert ttl is not None + + +@pytest.mark.parametrize( + "model,expected_region", + [ + ( + "arn:aws-cn:bedrock:cn-north-1:123456789012:application-inference-profile/p", + "cn-north-1", + ), + ( + "arn:aws-us-gov:bedrock:us-gov-west-1:123456789012:foundation-model/m", + "us-gov-west-1", + ), + ( + "bedrock/arn:aws-cn:bedrock:cn-northwest-1:123456789012:inference-profile/p", + "cn-northwest-1", + ), + ("anthropic.claude-3", None), + ], +) +def test_get_aws_region_from_model_arn_partition_arns(model, expected_region): + assert BaseAWSLLM()._get_aws_region_from_model_arn(model) == expected_region + + +@pytest.mark.parametrize( + "endpoint_type,region,expected", + [ + ("runtime", "cn-north-1", "https://bedrock-runtime.cn-north-1.amazonaws.com.cn"), + ("agent", "cn-north-1", "https://bedrock-agent-runtime.cn-north-1.amazonaws.com.cn"), + ("agentcore", "cn-north-1", "https://bedrock-agentcore.cn-north-1.amazonaws.com.cn"), + ("runtime", "us-east-1", "https://bedrock-runtime.us-east-1.amazonaws.com"), + ("agent", "us-east-1", "https://bedrock-agent-runtime.us-east-1.amazonaws.com"), + ("agentcore", "us-east-1", "https://bedrock-agentcore.us-east-1.amazonaws.com"), + ("runtime", "us-gov-west-1", "https://bedrock-runtime.us-gov-west-1.amazonaws.com"), + ], +) +def test_select_default_endpoint_url_partitions(endpoint_type, region, expected): + assert ( + BaseAWSLLM()._select_default_endpoint_url( + endpoint_type=endpoint_type, aws_region_name=region + ) + == expected + ) + + def test_irsa_cross_account_sts_client_uses_resolved_region(): """IRSA cross-account path must use _build_sts_client_kwargs (env region, not Bedrock).""" base_aws_llm = BaseAWSLLM() @@ -1612,6 +1738,7 @@ def test_sts_endpoint_region_matches_bedrock_region_param(): "aws_secret_access_key": "explicit-secret-key", "aws_session_token": "assumed-session-token", "verify": True, + "region_name": "us-east-1", }, ), ( @@ -1626,7 +1753,7 @@ def test_sts_endpoint_region_matches_bedrock_region_param(): }, ), ], - ids=["no_region_or_endpoint", "bedrock_region_ignored_for_sts", "explicit_sts_endpoint"], + ids=["no_region_or_endpoint", "configured_region_is_sts_fallback", "explicit_sts_endpoint"], ) def test_explicit_credentials_used_when_provided(role_kwargs, expected_client_kwargs): """ diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py index da6caca4f05..dd067aaa4fb 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py @@ -317,3 +317,55 @@ def test_body_model_stays_the_endpoint_name_when_hf_model_name_is_unset(monkeypa client = _invoke_sagemaker_chat(monkeypatch) assert client.request_body["model"] == "my-endpoint" + + +@pytest.mark.parametrize( + "region,stream,expected_url", + [ + ( + "cn-north-1", + False, + "https://runtime.sagemaker.cn-north-1.amazonaws.com.cn/endpoints/my-endpoint/invocations", + ), + ( + "cn-north-1", + True, + "https://runtime.sagemaker.cn-north-1.amazonaws.com.cn/endpoints/my-endpoint/invocations-response-stream", + ), + ( + "us-gov-west-1", + False, + "https://runtime.sagemaker.us-gov-west-1.amazonaws.com/endpoints/my-endpoint/invocations", + ), + ( + "us-west-2", + False, + "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/my-endpoint/invocations", + ), + ], +) +def test_get_complete_url_uses_partition_dns_suffix(region, stream, expected_url): + url = SagemakerChatConfig().get_complete_url( + api_base=None, + api_key=None, + model="my-endpoint", + optional_params={"aws_region_name": region}, + litellm_params={}, + stream=stream, + ) + assert url == expected_url + + +def test_get_complete_url_sagemaker_base_url_override_wins(): + url = SagemakerChatConfig().get_complete_url( + api_base=None, + api_key=None, + model="my-endpoint", + optional_params={ + "aws_region_name": "cn-north-1", + "sagemaker_base_url": "https://my-private-endpoint.example.com/invocations", + }, + litellm_params={}, + stream=False, + ) + assert url == "https://my-private-endpoint.example.com/invocations" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 3b506324ad7..09bab1dc416 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1870,7 +1870,10 @@ class TestBedrockAgentRuntimePassthroughToggle: with ( patch("litellm.proxy.proxy_server.general_settings", general_settings), - patch("litellm.utils.get_secret", return_value="us-east-1"), + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_secret_str", + return_value="us-east-1", + ), patch("litellm.llms.bedrock.chat.BedrockConverseLLM", return_value=bedrock_llm), patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_request_copy", @@ -1920,7 +1923,10 @@ class TestBedrockAgentRuntimePassthroughToggle: async def test_model_invoke_still_routed_when_agent_runtime_disabled(self): with ( patch("litellm.proxy.proxy_server.general_settings", self.DISABLED), - patch("litellm.utils.get_secret", return_value="us-east-1"), + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_secret_str", + return_value="us-east-1", + ), patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_request_copy", Mock(), diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py index 37d2141e460..df36326d064 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py @@ -92,3 +92,13 @@ def test_encode_bedrock_runtime_modelid_arn_edge_cases(): expected = "model/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile%2Ftest-profile.v1/invoke" result = CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) assert result == expected + + +def test_encode_bedrock_runtime_modelid_arn_partition_arns(): + endpoint = "model/arn:aws-cn:bedrock:cn-north-1:123456789012:application-inference-profile/r742sbn2zckd/converse" + expected = "model/arn:aws-cn:bedrock:cn-north-1:123456789012:application-inference-profile%2Fr742sbn2zckd/converse" + assert CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) == expected + + endpoint = "model/arn:aws-us-gov:bedrock:us-gov-west-1:123456789012:inference-profile/test-profile/invoke" + expected = "model/arn:aws-us-gov:bedrock:us-gov-west-1:123456789012:inference-profile%2Ftest-profile/invoke" + assert CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) == expected diff --git a/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py b/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py index 1e0e72c9ac6..874be247ec6 100644 --- a/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py +++ b/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py @@ -83,3 +83,56 @@ async def test_write_and_read_json_secret(): secret_name=test_secret_name ) assert delete_resp is not None + + +def _prepare_request_endpoint(monkeypatch, region_name, extra_optional_params=None): + monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) + secret_manager = AWSSecretsManagerV2(aws_region_name=region_name) + endpoint_url, _headers, _body = secret_manager._prepare_request( + action="GetSecretValue", + secret_name="my-secret", + optional_params={ + "aws_access_key_id": "test-key", + "aws_secret_access_key": "test-secret", + **(extra_optional_params or {}), + }, + ) + return endpoint_url + + +@pytest.mark.parametrize( + "region_name,expected_endpoint", + [ + ("cn-north-1", "https://secretsmanager.cn-north-1.amazonaws.com.cn"), + ("cn-northwest-1", "https://secretsmanager.cn-northwest-1.amazonaws.com.cn"), + ("us-gov-west-1", "https://secretsmanager.us-gov-west-1.amazonaws.com"), + ("us-east-1", "https://secretsmanager.us-east-1.amazonaws.com"), + ], +) +def test_prepare_request_builds_partition_endpoint(monkeypatch, region_name, expected_endpoint): + assert _prepare_request_endpoint(monkeypatch, region_name) == expected_endpoint + + +def test_prepare_request_explicit_bedrock_runtime_endpoint_param_still_wins(monkeypatch): + endpoint_url = _prepare_request_endpoint( + monkeypatch, + "cn-north-1", + {"aws_bedrock_runtime_endpoint": "https://bedrock-runtime.my-vpce.example.com"}, + ) + assert endpoint_url == "https://secretsmanager.my-vpce.example.com" + + +def test_prepare_request_env_bedrock_runtime_endpoint_still_wins(monkeypatch): + monkeypatch.setenv( + "AWS_BEDROCK_RUNTIME_ENDPOINT", "https://bedrock-runtime.eu-west-1.amazonaws.com" + ) + secret_manager = AWSSecretsManagerV2(aws_region_name="cn-north-1") + endpoint_url, _headers, _body = secret_manager._prepare_request( + action="GetSecretValue", + secret_name="my-secret", + optional_params={ + "aws_access_key_id": "test-key", + "aws_secret_access_key": "test-secret", + }, + ) + assert endpoint_url == "https://secretsmanager.eu-west-1.amazonaws.com" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index a7dec330a26..c2751b79adc 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16616 + "limit": 16615 }, "LIT011": { - "limit": 5583 + "limit": 5580 }, "LIT012": { "limit": 4509 From 1947c65081faa0732dfccece9b3fc544a1084742 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 02:27:18 -0700 Subject: [PATCH 3/3] test(aws): type the new partition test parameters --- tests/test_litellm/integrations/test_s3_v2.py | 2 +- .../litellm_core_utils/test_aws_partition.py | 18 +++++++++--------- .../bedrock/batches/test_transformation.py | 4 ++-- .../llms/bedrock/test_base_aws_llm.py | 15 +++++++++------ .../test_sagemaker_chat_transformation.py | 4 ++-- .../test_passthrough_endpoints_common_utils.py | 2 +- .../test_aws_secret_manager_v2.py | 12 ++++++++---- 7 files changed, 32 insertions(+), 25 deletions(-) diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index e4286fc43fa..51671d5101e 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1883,5 +1883,5 @@ def _s3_logger_for_region(region_name: str) -> S3Logger: ), ], ) -def test_build_object_url_uses_partition_dns_suffix(region_name, expected_url): +def test_build_object_url_uses_partition_dns_suffix(region_name: str, expected_url: str) -> None: assert _s3_logger_for_region(region_name)._build_object_url("2025-01-01/key.json") == expected_url diff --git a/tests/test_litellm/litellm_core_utils/test_aws_partition.py b/tests/test_litellm/litellm_core_utils/test_aws_partition.py index fcc3e370073..3594d3c354c 100644 --- a/tests/test_litellm/litellm_core_utils/test_aws_partition.py +++ b/tests/test_litellm/litellm_core_utils/test_aws_partition.py @@ -43,7 +43,7 @@ from litellm.llms.sagemaker.chat.transformation import SagemakerChatConfig ("", "aws", "amazonaws.com"), ], ) -def test_partition_lookup(region, partition, dns_suffix): +def test_partition_lookup(region: str | None, partition: str, dns_suffix: str) -> None: assert get_aws_partition(region) == AwsPartition(partition=partition, dns_suffix=dns_suffix) assert get_aws_dns_suffix(region) == dns_suffix assert get_aws_arn_prefix(region) == f"arn:{partition}:" @@ -61,7 +61,7 @@ def test_partition_lookup(region, partition, dns_suffix): ("arn:aws:iam::123456789012:role/foo", False), ], ) -def test_contains_bedrock_arn(value, expected): +def test_contains_bedrock_arn(value: str, expected: bool) -> None: assert contains_bedrock_arn(value) is expected @@ -76,7 +76,7 @@ def test_contains_bedrock_arn(value, expected): ("arn:aws:iam::123456789012:role/foo", False), ], ) -def test_is_bedrock_arn(value, expected): +def test_is_bedrock_arn(value: str, expected: bool) -> None: assert is_bedrock_arn(value) is expected @@ -90,7 +90,7 @@ def test_is_bedrock_arn(value, expected): ("arnaws:bedrock", False), ], ) -def test_contains_aws_arn(value, expected): +def test_contains_aws_arn(value: str, expected: bool) -> None: assert contains_aws_arn(value) is expected @@ -157,14 +157,14 @@ ENDPOINT_BUILDERS: Final = { @pytest.fixture(autouse=True) -def _clear_aws_env(monkeypatch): +def _clear_aws_env(monkeypatch: pytest.MonkeyPatch) -> None: for env_var in ("AWS_BEDROCK_RUNTIME_ENDPOINT", "AWS_REGION", "AWS_DEFAULT_REGION", "AWS_REGION_NAME"): monkeypatch.delenv(env_var, raising=False) @pytest.mark.parametrize("region", ["cn-north-1", "cn-northwest-1"]) @pytest.mark.parametrize("builder_name", sorted(ENDPOINT_BUILDERS)) -def test_every_endpoint_builder_respects_cn_partition(builder_name, region): +def test_every_endpoint_builder_respects_cn_partition(builder_name: str, region: str) -> None: url = ENDPOINT_BUILDERS[builder_name](region) hostname = urlparse(url).hostname assert hostname is not None @@ -175,7 +175,7 @@ def test_every_endpoint_builder_respects_cn_partition(builder_name, region): @pytest.mark.parametrize("region", ["us-east-1", "us-gov-west-1"]) @pytest.mark.parametrize("builder_name", sorted(ENDPOINT_BUILDERS)) -def test_every_endpoint_builder_keeps_amazonaws_com_outside_cn(builder_name, region): +def test_every_endpoint_builder_keeps_amazonaws_com_outside_cn(builder_name: str, region: str) -> None: url = ENDPOINT_BUILDERS[builder_name](region) hostname = urlparse(url).hostname assert hostname is not None @@ -194,9 +194,9 @@ def _fstring_literal_offenders(needle: str) -> list[str]: ] -def test_no_fstring_hardcodes_the_commercial_dns_suffix(): +def test_no_fstring_hardcodes_the_commercial_dns_suffix() -> None: assert _fstring_literal_offenders("amazonaws.com") == [] -def test_no_fstring_hardcodes_the_commercial_arn_prefix(): +def test_no_fstring_hardcodes_the_commercial_arn_prefix() -> None: assert _fstring_literal_offenders("arn:aws:") == [] diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index 664cac6c145..7e5716a7495 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -787,7 +787,7 @@ class TestBedrockBatchesContract(BatchesConfigContractTests): expected_retrieve_status = "completed" -def test_get_complete_batch_url_cn_partition(config): +def test_get_complete_batch_url_cn_partition(config: BedrockBatchesConfig) -> None: url = config.get_complete_batch_url( api_base=None, api_key=None, @@ -812,7 +812,7 @@ def test_get_complete_batch_url_cn_partition(config): ), ], ) -def test_retrieve_request_accepts_partition_arns(config, arn, expected_prefix): +def test_retrieve_request_accepts_partition_arns(config: BedrockBatchesConfig, arn: str, expected_prefix: str) -> None: with patch.object(config.common_utils, "sign_aws_request") as mock_sign: mock_sign.return_value = ({"Authorization": "signed"}, b"") result = config.transform_retrieve_batch_request( diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 81321ce73ec..7d07ac947b1 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -1436,8 +1436,11 @@ def test_build_sts_client_kwargs(env, aws_sts_endpoint, ssl_verify, expected): ], ) def test_resolve_sts_region_configured_region_fallback( - env, aws_sts_endpoint, aws_region_name, expected_region -): + env: dict[str, str], + aws_sts_endpoint: str | None, + aws_region_name: str | None, + expected_region: str | None, +) -> None: with patch.dict(os.environ, env, clear=True): assert ( BaseAWSLLM._resolve_sts_region( @@ -1448,7 +1451,7 @@ def test_resolve_sts_region_configured_region_fallback( ) -def test_build_sts_client_kwargs_configured_region_fallback(): +def test_build_sts_client_kwargs_configured_region_fallback() -> None: base_aws_llm = BaseAWSLLM() with patch.dict(os.environ, {}, clear=True): assert base_aws_llm._build_sts_client_kwargs(aws_region_name="cn-north-1") == { @@ -1462,7 +1465,7 @@ def test_build_sts_client_kwargs_configured_region_fallback(): } -def test_assume_role_sts_client_uses_configured_cn_region(): +def test_assume_role_sts_client_uses_configured_cn_region() -> None: """arn:aws-cn roles must resolve against a cn STS endpoint, not the commercial default.""" base_aws_llm = BaseAWSLLM() mock_expiry = MagicMock() @@ -1519,7 +1522,7 @@ def test_assume_role_sts_client_uses_configured_cn_region(): ("anthropic.claude-3", None), ], ) -def test_get_aws_region_from_model_arn_partition_arns(model, expected_region): +def test_get_aws_region_from_model_arn_partition_arns(model: str, expected_region: str | None) -> None: assert BaseAWSLLM()._get_aws_region_from_model_arn(model) == expected_region @@ -1535,7 +1538,7 @@ def test_get_aws_region_from_model_arn_partition_arns(model, expected_region): ("runtime", "us-gov-west-1", "https://bedrock-runtime.us-gov-west-1.amazonaws.com"), ], ) -def test_select_default_endpoint_url_partitions(endpoint_type, region, expected): +def test_select_default_endpoint_url_partitions(endpoint_type: str, region: str, expected: str) -> None: assert ( BaseAWSLLM()._select_default_endpoint_url( endpoint_type=endpoint_type, aws_region_name=region diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py index dd067aaa4fb..697f5a7ff59 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py @@ -344,7 +344,7 @@ def test_body_model_stays_the_endpoint_name_when_hf_model_name_is_unset(monkeypa ), ], ) -def test_get_complete_url_uses_partition_dns_suffix(region, stream, expected_url): +def test_get_complete_url_uses_partition_dns_suffix(region: str, stream: bool, expected_url: str) -> None: url = SagemakerChatConfig().get_complete_url( api_base=None, api_key=None, @@ -356,7 +356,7 @@ def test_get_complete_url_uses_partition_dns_suffix(region, stream, expected_url assert url == expected_url -def test_get_complete_url_sagemaker_base_url_override_wins(): +def test_get_complete_url_sagemaker_base_url_override_wins() -> None: url = SagemakerChatConfig().get_complete_url( api_base=None, api_key=None, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py index df36326d064..078bd4dd402 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py @@ -94,7 +94,7 @@ def test_encode_bedrock_runtime_modelid_arn_edge_cases(): assert result == expected -def test_encode_bedrock_runtime_modelid_arn_partition_arns(): +def test_encode_bedrock_runtime_modelid_arn_partition_arns() -> None: endpoint = "model/arn:aws-cn:bedrock:cn-north-1:123456789012:application-inference-profile/r742sbn2zckd/converse" expected = "model/arn:aws-cn:bedrock:cn-north-1:123456789012:application-inference-profile%2Fr742sbn2zckd/converse" assert CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) == expected diff --git a/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py b/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py index 874be247ec6..7e655b70756 100644 --- a/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py +++ b/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py @@ -85,7 +85,9 @@ async def test_write_and_read_json_secret(): assert delete_resp is not None -def _prepare_request_endpoint(monkeypatch, region_name, extra_optional_params=None): +def _prepare_request_endpoint( + monkeypatch: pytest.MonkeyPatch, region_name: str, extra_optional_params: dict[str, str] | None = None +) -> str: monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) secret_manager = AWSSecretsManagerV2(aws_region_name=region_name) endpoint_url, _headers, _body = secret_manager._prepare_request( @@ -109,11 +111,13 @@ def _prepare_request_endpoint(monkeypatch, region_name, extra_optional_params=No ("us-east-1", "https://secretsmanager.us-east-1.amazonaws.com"), ], ) -def test_prepare_request_builds_partition_endpoint(monkeypatch, region_name, expected_endpoint): +def test_prepare_request_builds_partition_endpoint( + monkeypatch: pytest.MonkeyPatch, region_name: str, expected_endpoint: str +) -> None: assert _prepare_request_endpoint(monkeypatch, region_name) == expected_endpoint -def test_prepare_request_explicit_bedrock_runtime_endpoint_param_still_wins(monkeypatch): +def test_prepare_request_explicit_bedrock_runtime_endpoint_param_still_wins(monkeypatch: pytest.MonkeyPatch) -> None: endpoint_url = _prepare_request_endpoint( monkeypatch, "cn-north-1", @@ -122,7 +126,7 @@ def test_prepare_request_explicit_bedrock_runtime_endpoint_param_still_wins(monk assert endpoint_url == "https://secretsmanager.my-vpce.example.com" -def test_prepare_request_env_bedrock_runtime_endpoint_still_wins(monkeypatch): +def test_prepare_request_env_bedrock_runtime_endpoint_still_wins(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv( "AWS_BEDROCK_RUNTIME_ENDPOINT", "https://bedrock-runtime.eu-west-1.amazonaws.com" )