diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py index f1e1f3e8738..56779de293f 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py @@ -7,6 +7,7 @@ Model format: bedrock/openai/ Example: bedrock/openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/abc123 """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -68,12 +69,22 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): return model + def _get_invoke_model_id(self, model: str, optional_params: Mapping[str, object]) -> str: + """ + `model_id` (e.g. an application inference profile ARN) overrides the model name, like on the other + Bedrock routes, so a deployment can keep the real model name and still invoke the profile + """ + model_id: Final = optional_params.get("model_id") + if isinstance(model_id, str) and model_id: + return model_id + return self._get_openai_model_id(model) + def get_complete_url( self, api_base: str | None, api_key: str | None, model: str, - optional_params: dict, + optional_params: dict[str, object], litellm_params: dict, stream: bool | None = None, ) -> str: @@ -82,7 +93,7 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): Uses the standard Bedrock invoke endpoint format. """ - model_id = self._get_openai_model_id(model) + model_id = self._get_invoke_model_id(model, optional_params) # Get AWS region aws_region_name: Final = self._get_aws_region_name(optional_params=optional_params, model=model) @@ -91,7 +102,9 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): aws_bedrock_runtime_endpoint: Final = optional_params.get("aws_bedrock_runtime_endpoint", None) endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint( api_base=api_base, - aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, + aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint + if isinstance(aws_bedrock_runtime_endpoint, str) + else None, aws_region_name=aws_region_name, ) @@ -131,7 +144,7 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): self, model: str, messages: list[AllMessageValues], - optional_params: dict, + optional_params: dict[str, object], litellm_params: dict, headers: dict, ) -> dict: @@ -144,12 +157,14 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): # Remove stream from optional_params as it's handled separately in URL optional_params.pop("stream", None) - # Remove AWS-specific params that shouldn't be in the request body - inference_params: Final = {k: v for k, v in optional_params.items() if k not in self.aws_authentication_params} + # Remove AWS-specific params and model_id (it goes in the URL) that shouldn't be in the request body + inference_params: Final = { + k: v for k, v in optional_params.items() if k not in self.aws_authentication_params and k != "model_id" + } - # Use parent class transform_request for OpenAI format + # Use parent class transform_request for OpenAI format. Bedrock rejects a body "model" that differs from the URL return super().transform_request( - model=self._get_openai_model_id(model), + model=self._get_invoke_model_id(model, optional_params), messages=messages, optional_params=inference_params, litellm_params=litellm_params, diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_openai_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_openai_transformation.py index 147501f01c2..72ef4538df1 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_openai_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_openai_transformation.py @@ -51,3 +51,42 @@ def test_openai_imported_model_preserves_message_types_and_request_parameters() "https://bedrock-runtime.us-east-1.amazonaws.com/model/" "arn:aws:bedrock:us-east-1:117159858402:imported-model%2Fm4gc1mrfuddy/invoke" ) + + +def test_model_id_overrides_model_name_in_url_and_body() -> None: + profile_arn: Final = "arn:aws:bedrock:ap-northeast-2:123456789012:application-inference-profile/abc123def456" + model: Final = "bedrock/openai/global.openai.gpt-6-luna" + messages: Final = [{"role": "user", "content": "Reply with the single word pong"}] + optional_params: Final = { + "model_id": profile_arn, + "max_completion_tokens": 200, + "aws_access_key_id": "AKIAUNITTESTKEY", + "aws_secret_access_key": "unit-test-secret", + "aws_region_name": "ap-northeast-2", + } + config: Final = AmazonBedrockOpenAIConfig() + url: Final = config.get_complete_url( + api_base=None, + api_key=None, + model=model, + optional_params=dict(optional_params), + litellm_params={}, + stream=False, + ) + request: Final = config.transform_request( + model=model, + messages=messages, + optional_params=dict(optional_params), + litellm_params={}, + headers={}, + ) + + assert url == ( + "https://bedrock-runtime.ap-northeast-2.amazonaws.com/model/" + "arn:aws:bedrock:ap-northeast-2:123456789012:application-inference-profile%2Fabc123def456/invoke" + ) + assert request == { + "model": profile_arn, + "messages": messages, + "max_completion_tokens": 200, + }