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

Made-with: Cursor

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)
 

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:

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:

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

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