Allow using ARNs when generation images via Bedrock (BerriAI#15789)

komarovd95 · web-flow · commit 1dfdcb076281 · 2025-10-28T19:41:35.000-07:00
* Use model_id in Bedrock's image_handler

* Fix MyPy for converse_handler and invoke_hanlder
diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py
@@ -1,6 +1,7 @@
 import hashlib
 import json
 import os
+import urllib.parse
 from datetime import datetime
 from typing import (
     TYPE_CHECKING,
@@ -331,16 +332,61 @@ def get_bedrock_invoke_provider(
                     return provider
         return None
 
+    @staticmethod
+    def get_bedrock_model_id(
+        optional_params: dict,
+        provider: Optional[BEDROCK_INVOKE_PROVIDERS_LITERAL],
+        model: str,
+    ) -> str:
+        model_id = optional_params.pop("model_id", None)
+        if model_id is not None:
+            model_id = BaseAWSLLM.encode_model_id(model_id=model_id)
+        else:
+            model_id = model
+
+        model_id = model_id.replace("invoke/", "", 1)
+        if provider == "llama" and "llama/" in model_id:
+            model_id = BaseAWSLLM._get_model_id_from_model_with_spec(
+                model_id, spec="llama"
+            )
+        elif provider == "deepseek_r1" and "deepseek_r1/" in model_id:
+            model_id = BaseAWSLLM._get_model_id_from_model_with_spec(
+                model_id, spec="deepseek_r1"
+            )
+        return model_id
+
+    @staticmethod
+    def _get_model_id_from_model_with_spec(
+        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 BaseAWSLLM.encode_model_id(model_id=model_id)
+
+    @staticmethod
+    def encode_model_id(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="")
+
     @staticmethod
     def get_bedrock_embedding_provider(
         model: str,
     ) -> Optional[BEDROCK_EMBEDDING_PROVIDERS_LITERAL]:
         """
         Helper function to get the bedrock embedding provider from the model
-        
+
         Handles scenarios like:
         1. model=cohere.embed-english-v3:0 -> Returns `cohere`
-        2. model=amazon.titan-embed-text-v1 -> Returns `amazon`  
+        2. model=amazon.titan-embed-text-v1 -> Returns `amazon`
         3. model=us.twelvelabs.marengo-embed-2-7-v1:0 -> Returns `twelvelabs`
         4. model=twelvelabs.marengo-embed-2-7-v1:0 -> Returns `twelvelabs`
         """
@@ -349,20 +395,24 @@ def get_bedrock_embedding_provider(
             parts = model.split(".")
             # Check if the second part (after potential region) is a known provider
             if len(parts) >= 2:
-                potential_provider = parts[1]  # e.g., "twelvelabs" from "us.twelvelabs.marengo-embed-2-7-v1:0"
+                potential_provider = parts[
+                    1
+                ]  # e.g., "twelvelabs" from "us.twelvelabs.marengo-embed-2-7-v1:0"
                 if potential_provider in get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL):
                     return cast(BEDROCK_EMBEDDING_PROVIDERS_LITERAL, potential_provider)
-            
+
             # Check if the first part is a known provider (standard format)
-            potential_provider = parts[0]  # e.g., "cohere" from "cohere.embed-english-v3:0"
+            potential_provider = parts[
+                0
+            ]  # e.g., "cohere" from "cohere.embed-english-v3:0"
             if potential_provider in get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL):
                 return cast(BEDROCK_EMBEDDING_PROVIDERS_LITERAL, potential_provider)
-        
+
         # Fallback: check if any provider name appears in the model string
         for provider in get_args(BEDROCK_EMBEDDING_PROVIDERS_LITERAL):
             if provider in model:
                 return cast(BEDROCK_EMBEDDING_PROVIDERS_LITERAL, provider)
-        
+
         return None
 
     def _get_aws_region_name(
@@ -984,20 +1034,23 @@ def get_request_headers(
                 raise ImportError(
                     "Missing boto3 to call bedrock. Run 'pip install boto3'."
                 )
-            
+
             # Filter headers for AWS signature calculation
             # AWS SigV4 only includes specific headers in signature calculation
             aws_signature_headers = self._filter_headers_for_aws_signature(headers)
             sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name)
             request = AWSRequest(
-                method="POST", url=endpoint_url, data=data, headers=aws_signature_headers
+                method="POST",
+                url=endpoint_url,
+                data=data,
+                headers=aws_signature_headers,
             )
             sigv4.add_auth(request)
-            
+
             # Add back all original headers (including forwarded ones) after signature calculation
             for header_name, header_value in headers.items():
                 request.headers[header_name] = header_value
-            
+
             if (
                 extra_headers is not None and "Authorization" in extra_headers
             ):  # prevent sigv4 from overwriting the auth header
@@ -1013,16 +1066,27 @@ def _filter_headers_for_aws_signature(self, headers: dict) -> dict:
         """
         aws_signature_headers = {}
         aws_headers = {
-            'host', 'content-type', 'date', 'x-amz-date', 'x-amz-security-token',
-            'x-amz-content-sha256', 'x-amz-algorithm', 'x-amz-credential',
-            'x-amz-signedheaders', 'x-amz-signature'
+            "host",
+            "content-type",
+            "date",
+            "x-amz-date",
+            "x-amz-security-token",
+            "x-amz-content-sha256",
+            "x-amz-algorithm",
+            "x-amz-credential",
+            "x-amz-signedheaders",
+            "x-amz-signature",
         }
-        
+
         for header_name, header_value in headers.items():
             header_lower = header_name.lower()
-            if header_lower in aws_headers or header_lower.startswith('x-amz-') or header_lower.startswith('x-amzn-'):
+            if (
+                header_lower in aws_headers
+                or header_lower.startswith("x-amz-")
+                or header_lower.startswith("x-amzn-")
+            ):
                 aws_signature_headers[header_name] = header_value
-        
+
         return aws_signature_headers
 
     def _sign_request(
diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py
@@ -1,5 +1,4 @@
 import json
-import urllib
 from typing import Any, Optional, Union
 
 import httpx
@@ -84,16 +83,6 @@ class BedrockConverseLLM(BaseAWSLLM):
     def __init__(self) -> None:
         super().__init__()
 
-    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="")  # type: ignore
-
     async def async_streaming(
         self,
         model: str,
diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py
@@ -3,20 +3,15 @@
 """
 
 import copy
-import json
 import time
 import types
-import urllib.parse
 from functools import partial
 from typing import (
-    Any,
     AsyncIterator,
     Callable,
     Iterator,
-    List,
     Optional,
     Tuple,
-    Union,
     cast,
     get_args,
 )
@@ -672,16 +667,6 @@ def process_response(  # noqa: PLR0915
 
         return model_response
 
-    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 completion(  # noqa: PLR0915
         self,
         model: str,
@@ -1176,33 +1161,6 @@ def _get_provider_from_model_path(
                 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
-
-        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
diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
@@ -1,7 +1,6 @@
 import copy
 import json
 import time
-import urllib.parse
 from functools import partial
 from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union, cast, get_args
 
@@ -190,14 +189,16 @@ def transform_request(
                     ] = True  # cohere requires stream = True in inference params
                 request_data = {"prompt": prompt, **inference_params}
         elif provider == "anthropic":
-            transformed_request = litellm.AmazonAnthropicClaudeConfig().transform_request(
-                model=model,
-                messages=messages,
-                optional_params=optional_params,
-                litellm_params=litellm_params,
-                headers=headers,
+            transformed_request = (
+                litellm.AmazonAnthropicClaudeConfig().transform_request(
+                    model=model,
+                    messages=messages,
+                    optional_params=optional_params,
+                    litellm_params=litellm_params,
+                    headers=headers,
+                )
             )
-            
+
             return transformed_request
         elif provider == "nova":
             return litellm.AmazonInvokeNovaConfig().transform_request(
@@ -327,7 +328,9 @@ def transform_response(  # noqa: PLR0915
             elif provider == "meta" or provider == "llama" or provider == "deepseek_r1":
                 outputText = completion_response["generation"]
             elif provider == "mistral":
-                outputText = litellm.AmazonMistralConfig.get_outputText(completion_response, model_response)
+                outputText = litellm.AmazonMistralConfig.get_outputText(
+                    completion_response, model_response
+                )
             else:  # amazon titan
                 outputText = completion_response.get("results")[0].get("outputText")
         except Exception as e:
@@ -549,48 +552,6 @@ def _get_provider_from_model_path(
                 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)
-        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]]:
diff --git a/litellm/llms/bedrock/image/image_handler.py b/litellm/llms/bedrock/image/image_handler.py
diff --git a/tests/test_litellm/llms/bedrock/image/test_bedrock_image_prepare_request.py b/tests/test_litellm/llms/bedrock/image/test_bedrock_image_prepare_request.py
