mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_decrease_anys_opus5_0826
# Conflicts: # basedpyright-code-budget.json # type-discipline-budget.json
This commit is contained in:
mateo-berri
commit
ee86b62b66
32 changed files with 864 additions and 81 deletions
|
|
@ -3,7 +3,7 @@
|
|||
"limit": 17271
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"limit": 2544
|
||||
"limit": 2539
|
||||
},
|
||||
"reportAssignmentType": {
|
||||
"limit": 319
|
||||
|
|
@ -111,7 +111,7 @@
|
|||
"limit": 19778
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 30292
|
||||
"limit": 30290
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 117
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
55
litellm/litellm_core_utils/aws_partition.py
Normal file
55
litellm/litellm_core_utils/aws_partition.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -5,6 +5,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:
|
||||
|
|
@ -337,7 +338,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}"
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, 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,
|
||||
)
|
||||
|
|
@ -141,8 +142,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
|
||||
|
||||
|
|
@ -241,8 +244,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",
|
||||
|
|
@ -374,7 +378,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
|
||||
|
|
@ -393,7 +397,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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
@ -99,7 +100,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=<value>
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -61,6 +61,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:
|
||||
|
|
@ -87,6 +88,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
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from typing import Final, Protocol
|
|||
from pydantic import JsonValue, 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 litellm.types.realtime import RealtimeResponseTransformInput
|
||||
|
||||
|
|
@ -128,7 +129,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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -34,6 +35,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 +62,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
|
||||
|
||||
|
|
@ -79,10 +82,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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
@ -57,6 +58,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 +85,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
|
||||
|
||||
|
|
@ -104,10 +107,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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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],
|
||||
},
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
{
|
||||
"ANN001": {
|
||||
"limit": 3013
|
||||
"limit": 3012
|
||||
},
|
||||
"ANN002": {
|
||||
"limit": 71
|
||||
|
|
@ -9,13 +9,13 @@
|
|||
"limit": 827
|
||||
},
|
||||
"ANN201": {
|
||||
"limit": 2007
|
||||
"limit": 2003
|
||||
},
|
||||
"ANN202": {
|
||||
"limit": 846
|
||||
"limit": 845
|
||||
},
|
||||
"ANN204": {
|
||||
"limit": 704
|
||||
"limit": 702
|
||||
},
|
||||
"ANN205": {
|
||||
"limit": 112
|
||||
|
|
@ -24,7 +24,7 @@
|
|||
"limit": 133
|
||||
},
|
||||
"ANN401": {
|
||||
"limit": 904
|
||||
"limit": 655
|
||||
},
|
||||
"ASYNC230": {
|
||||
"limit": 11
|
||||
|
|
@ -39,7 +39,7 @@
|
|||
"limit": 503
|
||||
},
|
||||
"B009": {
|
||||
"limit": 55
|
||||
"limit": 52
|
||||
},
|
||||
"B010": {
|
||||
"limit": 190
|
||||
|
|
@ -78,7 +78,7 @@
|
|||
"limit": 1
|
||||
},
|
||||
"C901": {
|
||||
"limit": 312
|
||||
"limit": 311
|
||||
},
|
||||
"D419": {
|
||||
"limit": 6
|
||||
|
|
@ -117,7 +117,7 @@
|
|||
"limit": 1
|
||||
},
|
||||
"PERF102": {
|
||||
"limit": 25
|
||||
"limit": 23
|
||||
},
|
||||
"PERF401": {
|
||||
"limit": 12
|
||||
|
|
@ -177,7 +177,7 @@
|
|||
"limit": 8
|
||||
},
|
||||
"RUF019": {
|
||||
"limit": 35
|
||||
"limit": 32
|
||||
},
|
||||
"RUF046": {
|
||||
"limit": 4
|
||||
|
|
@ -198,7 +198,7 @@
|
|||
"limit": 58
|
||||
},
|
||||
"SIM102": {
|
||||
"limit": 316
|
||||
"limit": 315
|
||||
},
|
||||
"SIM103": {
|
||||
"limit": 119
|
||||
|
|
@ -231,7 +231,7 @@
|
|||
"limit": 5
|
||||
},
|
||||
"TID251": {
|
||||
"limit": 1159
|
||||
"limit": 1117
|
||||
},
|
||||
"TRY002": {
|
||||
"limit": 524
|
||||
|
|
@ -246,7 +246,7 @@
|
|||
"limit": 113
|
||||
},
|
||||
"TRY300": {
|
||||
"limit": 858
|
||||
"limit": 857
|
||||
},
|
||||
"UP028": {
|
||||
"limit": 2
|
||||
|
|
|
|||
|
|
@ -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: str, expected_url: str) -> None:
|
||||
assert _s3_logger_for_region(region_name)._build_object_url("2025-01-01/key.json") == expected_url
|
||||
|
|
|
|||
202
tests/test_litellm/litellm_core_utils/test_aws_partition.py
Normal file
202
tests/test_litellm/litellm_core_utils/test_aws_partition.py
Normal file
|
|
@ -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: 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}:"
|
||||
|
||||
|
||||
@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: str, expected: bool) -> None:
|
||||
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: str, expected: bool) -> None:
|
||||
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: str, expected: bool) -> None:
|
||||
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: 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: str, region: str) -> None:
|
||||
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: str, region: str) -> None:
|
||||
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() -> None:
|
||||
assert _fstring_literal_offenders("amazonaws.com") == []
|
||||
|
||||
|
||||
def test_no_fstring_hardcodes_the_commercial_arn_prefix() -> None:
|
||||
assert _fstring_literal_offenders("arn:aws:") == []
|
||||
|
|
@ -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: BedrockBatchesConfig) -> None:
|
||||
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: 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(
|
||||
batch_id=arn, optional_params={}, litellm_params={}
|
||||
)
|
||||
assert result["url"].startswith(expected_prefix)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,135 @@ 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: 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(
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_region_name=aws_region_name,
|
||||
)
|
||||
== expected_region
|
||||
)
|
||||
|
||||
|
||||
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") == {
|
||||
"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() -> None:
|
||||
"""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: str, expected_region: str | None) -> None:
|
||||
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: str, region: str, expected: str) -> None:
|
||||
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 +1741,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 +1756,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):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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: str, stream: bool, expected_url: str) -> None:
|
||||
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() -> None:
|
||||
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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
|
|
@ -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() -> 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
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -83,3 +83,60 @@ async def test_write_and_read_json_secret():
|
|||
secret_name=test_secret_name
|
||||
)
|
||||
assert delete_resp is not 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(
|
||||
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: 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: pytest.MonkeyPatch) -> None:
|
||||
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: pytest.MonkeyPatch) -> None:
|
||||
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"
|
||||
|
|
|
|||
|
|
@ -27,10 +27,10 @@
|
|||
"limit": 0
|
||||
},
|
||||
"LIT010": {
|
||||
"limit": 16565
|
||||
"limit": 16564
|
||||
},
|
||||
"LIT011": {
|
||||
"limit": 5580
|
||||
"limit": 5577
|
||||
},
|
||||
"LIT012": {
|
||||
"limit": 4508
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue