Merge pull request #24897 from Sameerlite/litellm_bedrock-embedded-region-model-path

feat(bedrock): parse embedded region from invoke model path

litellm/llms/bedrock/base_aws_llm.py

@@ -488,6 +488,16 @@ 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:
+                from litellm.llms.bedrock.common_utils import (
+                    split_embedded_bedrock_region_prefix,
+                    strip_bedrock_routing_prefix,
+                )
+
+                _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)
 

litellm/llms/bedrock/chat/invoke_handler.py

@@ -67,7 +67,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(
@@ -1267,6 +1272,36 @@ class BedrockLLM(BaseAWSLLM):
                 return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, provider)
         return None
 
+    def get_bedrock_model_id(
+        self,
+        optional_params: dict,
+        provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL],
+        model: str,
+    ) -> str:
+        modelId = optional_params.pop("model_id", None)
+        if modelId is not None:
+            modelId = self.encode_model_id(model_id=modelId)
+        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)
+
+        return modelId
+
+    def _get_model_id_for_llama_like_model(
+        self,
+        model: str,
+    ) -> str:
+        """
+        Remove `llama` from modelID since `llama` is simply a spec to follow for custom bedrock models
+        """
+        model_id = model.replace("llama/", "")
+        return self.encode_model_id(model_id=model_id)
+
 
 def get_response_stream_shape():
     global _response_stream_shape_cache

litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py

@@ -18,7 +18,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,
@@ -564,6 +567,51 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
                 return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, provider)
         return None
 
+    def get_bedrock_model_id(
+        self,
+        optional_params: dict,
+        provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL],
+        model: str,
+    ) -> str:
+        modelId = optional_params.pop("model_id", None)
+        if modelId is not None:
+            modelId = self.encode_model_id(model_id=modelId)
+        else:
+            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:
+            modelId = self._get_model_id_from_model_with_spec(
+                modelId, spec="deepseek_r1"
+            )
+        return modelId
+
+    def _get_model_id_from_model_with_spec(
+        self,
+        model: str,
+        spec: str,
+    ) -> str:
+        """
+        Remove `llama` from modelID since `llama` is simply a spec to follow for custom bedrock models
+        """
+        model_id = model.replace(spec + "/", "")
+        return self.encode_model_id(model_id=model_id)
+
+    def encode_model_id(self, model_id: str) -> str:
+        """
+        Double encode the model ID to ensure it matches the expected double-encoded format.
+        Args:
+            model_id (str): The model ID to encode.
+        Returns:
+            str: The double-encoded model ID.
+        """
+        return urllib.parse.quote(model_id, safe="")
+
     def convert_messages_to_prompt(
         self, model, messages, provider, custom_prompt_dict
     ) -> Tuple[str, Optional[list]]:

litellm/llms/bedrock/common_utils.py

@@ -6,7 +6,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
@@ -145,6 +145,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."""
 
@@ -247,7 +301,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

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():