mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
fix(bedrock): honor model_id on the bedrock/openai invoke route
bedrock/openai/<model> ignored model_id when building the invoke URL and forwarded it in the request body, so a deployment that keeps the model name and points model_id at an application inference profile ARN got a 400 "Unknown parameter: 'model_id'" from Bedrock. Use model_id for the URL and the body model, and keep it out of the body, like the other Bedrock routes
This commit is contained in:
parent
e73f06ecdc
commit
b346e611bc
2 changed files with 62 additions and 8 deletions
|
|
@ -7,6 +7,7 @@ Model format: bedrock/openai/<model-id>
|
|||
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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue