diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 8211addaf95..ed57df788ed 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -26,6 +26,10 @@ from litellm.constants import ( BEDROCK_MAX_POLICY_SIZE, ) from litellm.litellm_core_utils.dd_tracing import tracer +from litellm.llms.bedrock.common_utils import ( + split_embedded_bedrock_region_prefix, + strip_bedrock_routing_prefix, +) from litellm.secret_managers.main import get_secret, get_secret_str if TYPE_CHECKING: @@ -390,6 +394,11 @@ class BaseAWSLLM: aws_region_name = self._get_aws_region_from_model_arn(model_id) else: aws_region_name = self._get_aws_region_from_model_arn(model) + if aws_region_name is None and model is not None: + _stripped = strip_bedrock_routing_prefix(model) + _embedded, _ = split_embedded_bedrock_region_prefix(_stripped) + if _embedded is not None: + aws_region_name = _embedded # check env # litellm_aws_region_name = get_secret("AWS_REGION_NAME", None) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 42cdb34fc1a..cb3ed402f36 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -56,7 +56,11 @@ from litellm.types.llms.openai import ( ChatCompletionToolCallFunctionChunk, ChatCompletionUsageBlock, ) -from litellm.types.utils import ChatCompletionMessageToolCall, Choices, Delta +from litellm.types.utils import ( + ChatCompletionMessageToolCall, + Choices, + Delta, +) from litellm.types.utils import GenericStreamingChunk as GChunk from litellm.types.utils import ( ModelResponse, @@ -67,7 +71,12 @@ from litellm.types.utils import ( from litellm.utils import CustomStreamWrapper, get_secret from ..base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockError, ModelResponseIterator, get_bedrock_tool_name +from ..common_utils import ( + BedrockError, + ModelResponseIterator, + apply_embedded_bedrock_region_from_model_path, + get_bedrock_tool_name, +) _response_stream_shape_cache = None bedrock_tool_name_mappings: InMemoryCache = InMemoryCache( @@ -1188,6 +1197,9 @@ class BedrockLLM(BaseAWSLLM): else: modelId = model + modelId = apply_embedded_bedrock_region_from_model_path( + modelId, optional_params + ) if provider == "llama" and "llama/" in modelId: modelId = self._get_model_id_for_llama_like_model(modelId) @@ -1448,7 +1460,7 @@ class AWSEventStreamDecoder: ######### /bedrock/invoke nova mappings ############### elif "contentBlockDelta" in chunk_data: # when using /bedrock/invoke/nova, the chunk_data is nested under "contentBlockDelta" - _chunk_data = chunk_data.get("contentBlockDelta", None) + _chunk_data = chunk_data["contentBlockDelta"] return self.converse_chunk_parser(chunk_data=_chunk_data) ######## bedrock.mistral mappings ############### elif "outputs" in chunk_data: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 08a0690716b..5b0f13ed7a0 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -19,7 +19,10 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.chat.invoke_handler import make_call, make_sync_call -from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.bedrock.common_utils import ( + BedrockError, + apply_embedded_bedrock_region_from_model_path, +) from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -562,6 +565,9 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): modelId = model modelId = modelId.replace("invoke/", "", 1) + modelId = apply_embedded_bedrock_region_from_model_path( + modelId, optional_params + ) if provider == "llama" and "llama/" in modelId: modelId = self._get_model_id_from_model_with_spec(modelId, spec="llama") elif provider == "deepseek_r1" and "deepseek_r1/" in modelId: diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index c7f7acf331d..1d620a9caf3 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -4,7 +4,7 @@ Common utilities used across bedrock chat/embedding/image generation import json import os -from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union +from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union if TYPE_CHECKING: from litellm.types.llms.bedrock import BedrockCreateBatchRequest @@ -102,6 +102,60 @@ class AmazonBedrockGlobalConfig: ] +_BEDROCK_INVOKE_EMBEDDED_REGIONS: Optional[frozenset] = None + + +def _bedrock_invoke_embedded_region_names() -> frozenset: + global _BEDROCK_INVOKE_EMBEDDED_REGIONS + if _BEDROCK_INVOKE_EMBEDDED_REGIONS is None: + _BEDROCK_INVOKE_EMBEDDED_REGIONS = frozenset( + AmazonBedrockGlobalConfig().get_all_regions() + ) + return _BEDROCK_INVOKE_EMBEDDED_REGIONS + + +def strip_bedrock_routing_prefix(model: str) -> str: + """Strip one LiteLLM routing prefix (bedrock/, converse/, etc.).""" + s = model + for prefix in ("bedrock/", "converse/", "invoke/", "openai/", "nova-2/", "nova/"): + if s.startswith(prefix): + s = s.split("/", 1)[1] + return s + + +def split_embedded_bedrock_region_prefix( + model_id: str, +) -> Tuple[Optional[str], str]: + """ + If model_id is ``{region}/{rest}`` and region is a Bedrock AWS region, return + (region, rest). Otherwise (None, model_id). + """ + if "/" not in model_id: + return None, model_id + prefix, remainder = model_id.split("/", 1) + if not remainder.strip(): + return None, model_id + if prefix not in _bedrock_invoke_embedded_region_names(): + return None, model_id + return prefix, remainder + + +def apply_embedded_bedrock_region_from_model_path( + model_id: str, optional_params: dict +) -> str: + """ + Strip routing prefixes, then if the id is ``region/bedrockModelId``, set + ``optional_params['aws_region_name']`` when unset and return ``bedrockModelId`` only. + """ + stripped = strip_bedrock_routing_prefix(model_id) + region, remainder = split_embedded_bedrock_region_prefix(stripped) + if region is None: + return stripped + if optional_params.get("aws_region_name") is None: + optional_params["aws_region_name"] = region + return remainder + + def add_custom_header(headers): """Closure to capture the headers and add them.""" @@ -191,7 +245,8 @@ def init_bedrock_client( config = boto3.session.Config(connect_timeout=timeout, read_timeout=timeout) # type: ignore elif isinstance(timeout, httpx.Timeout): config = boto3.session.Config( # type: ignore - connect_timeout=timeout.connect, read_timeout=timeout.read + connect_timeout=timeout.connect, # type: ignore[arg-type, union-attr] + read_timeout=timeout.read, # type: ignore[arg-type, union-attr] ) else: config = boto3.session.Config() # type: ignore diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index f66ed21cf78..94d6621cf0b 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -10,7 +10,28 @@ sys.path.insert( ) # Adds the parent directory to the system path -from litellm.llms.bedrock.common_utils import BedrockModelInfo +from litellm.llms.bedrock.common_utils import ( + BedrockModelInfo, + apply_embedded_bedrock_region_from_model_path, +) + + +def test_apply_embedded_bedrock_region_strips_prefix_and_sets_region(): + optional_params: dict = {} + out = apply_embedded_bedrock_region_from_model_path( + "bedrock/us-west-2/mistral.mistral-7b-instruct-v0:2", optional_params + ) + assert out == "mistral.mistral-7b-instruct-v0:2" + assert optional_params.get("aws_region_name") == "us-west-2" + + +def test_apply_embedded_bedrock_region_respects_explicit_aws_region_name(): + optional_params = {"aws_region_name": "us-east-1"} + out = apply_embedded_bedrock_region_from_model_path( + "bedrock/us-west-2/mistral.mistral-7b-instruct-v0:2", optional_params + ) + assert out == "mistral.mistral-7b-instruct-v0:2" + assert optional_params.get("aws_region_name") == "us-east-1" def test_deepseek_cris():