mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
Merge 84084d9a82 into 559247fa84
This commit is contained in:
Mateo Wang
GitHub
commit
baa60bea71
9 changed files with 363 additions and 0 deletions
|
|
@ -26,6 +26,7 @@ EXTRA_BOOLEAN_KEYS = frozenset(
|
|||
"gemini_audio_only_live",
|
||||
"uses_embed_content",
|
||||
"use_openai_responses_path",
|
||||
"use_bedrock_runtime_chat_completions",
|
||||
"bedrock_converse_supports_strict_tools",
|
||||
"thinking_always_on",
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1736,6 +1736,9 @@ if TYPE_CHECKING:
|
|||
from .llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import (
|
||||
AmazonBedrockOpenAIConfig as AmazonBedrockOpenAIConfig,
|
||||
)
|
||||
from .llms.bedrock.chat.chat_completions.transformation import (
|
||||
AmazonBedrockRuntimeChatCompletionsConfig as AmazonBedrockRuntimeChatCompletionsConfig,
|
||||
)
|
||||
from .llms.bedrock.image_generation.amazon_stability1_transformation import (
|
||||
AmazonStabilityConfig as AmazonStabilityConfig,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -204,6 +204,7 @@ LLM_CONFIG_NAMES: Final = (
|
|||
"AmazonTwelveLabsPegasusConfig",
|
||||
"AmazonInvokeConfig",
|
||||
"AmazonBedrockOpenAIConfig",
|
||||
"AmazonBedrockRuntimeChatCompletionsConfig",
|
||||
"AmazonStabilityConfig",
|
||||
"AmazonStability3Config",
|
||||
"AmazonNovaCanvasConfig",
|
||||
|
|
@ -847,6 +848,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
|
|||
".llms.bedrock.chat.invoke_transformations.amazon_openai_transformation",
|
||||
"AmazonBedrockOpenAIConfig",
|
||||
),
|
||||
"AmazonBedrockRuntimeChatCompletionsConfig": (
|
||||
".llms.bedrock.chat.chat_completions.transformation",
|
||||
"AmazonBedrockRuntimeChatCompletionsConfig",
|
||||
),
|
||||
"AmazonStabilityConfig": (
|
||||
".llms.bedrock.image_generation.amazon_stability1_transformation",
|
||||
"AmazonStabilityConfig",
|
||||
|
|
|
|||
174
litellm/llms/bedrock/chat/chat_completions/transformation.py
Normal file
174
litellm/llms/bedrock/chat/chat_completions/transformation.py
Normal file
|
|
@ -0,0 +1,174 @@
|
|||
"""
|
||||
Native OpenAI Chat Completions on Amazon Bedrock Runtime.
|
||||
|
||||
AWS serves this surface at
|
||||
``https://bedrock-runtime.{region}.amazonaws.com/openai/v1/chat/completions``.
|
||||
Grok 4.6 on runtime is one of the models that uses it: chat completions stay
|
||||
chat completions instead of being rewritten to Converse.
|
||||
|
||||
Usage: model="us.xai.grok-4.6" or model="bedrock/us.xai.grok-4.6"
|
||||
Explicit ``bedrock/converse/...`` still uses Converse.
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from typing import Any, Final
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock.common_utils import BedrockError, strip_bedrock_routing_prefix
|
||||
from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
||||
|
||||
class AmazonBedrockRuntimeChatCompletionsConfig(OpenAILikeChatConfig):
|
||||
def __init__(self, aws_signer: BaseAWSLLM | None = None):
|
||||
super().__init__()
|
||||
self._aws_signer: Final = aws_signer or BaseAWSLLM()
|
||||
|
||||
@property
|
||||
def custom_llm_provider(self) -> str | None:
|
||||
return "bedrock"
|
||||
|
||||
def get_error_class(
|
||||
self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers
|
||||
) -> BaseLLMException:
|
||||
return BedrockError(status_code=status_code, message=error_message, headers=headers)
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
stream: bool | None = None,
|
||||
) -> str:
|
||||
if api_base is not None and "chat/completions" in api_base:
|
||||
return api_base.rstrip("/")
|
||||
aws_region_name: Final = self._aws_signer._get_aws_region_name(optional_params=optional_params, model=model)
|
||||
endpoint_url, _ = self._aws_signer.get_runtime_endpoint(
|
||||
api_base=api_base,
|
||||
aws_bedrock_runtime_endpoint=optional_params.get("aws_bedrock_runtime_endpoint"),
|
||||
aws_region_name=aws_region_name,
|
||||
)
|
||||
base: Final = endpoint_url.rstrip("/")
|
||||
if base.endswith("/openai/v1/chat/completions"):
|
||||
return base
|
||||
if base.endswith("/openai/v1"):
|
||||
return f"{base}/chat/completions"
|
||||
return f"{base}/openai/v1/chat/completions"
|
||||
|
||||
def sign_request(
|
||||
self,
|
||||
headers: dict,
|
||||
optional_params: dict,
|
||||
request_data: dict,
|
||||
api_base: str,
|
||||
api_key: str | None = None,
|
||||
model: str | None = None,
|
||||
stream: bool | None = None,
|
||||
fake_stream: bool | None = None,
|
||||
) -> tuple[dict, bytes | None]:
|
||||
return self._aws_signer._sign_request(
|
||||
service_name="bedrock",
|
||||
headers=headers,
|
||||
optional_params=optional_params,
|
||||
request_data=request_data,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
stream=stream,
|
||||
fake_stream=fake_stream,
|
||||
)
|
||||
|
||||
def transform_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
inference_params: Final = {
|
||||
k: v for k, v in optional_params.items() if k not in self._aws_signer.aws_authentication_params
|
||||
}
|
||||
return super().transform_request(
|
||||
model=strip_bedrock_routing_prefix(model),
|
||||
messages=messages,
|
||||
optional_params=inference_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
async def async_transform_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
inference_params: Final = {
|
||||
k: v for k, v in optional_params.items() if k not in self._aws_signer.aws_authentication_params
|
||||
}
|
||||
return await super().async_transform_request(
|
||||
model=strip_bedrock_routing_prefix(model),
|
||||
messages=messages,
|
||||
optional_params=inference_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
model: str,
|
||||
messages: list[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict:
|
||||
headers = super().validate_environment(
|
||||
headers=headers,
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
)
|
||||
project_id: Final = litellm_params.get("aws_bedrock_project_id")
|
||||
if project_id:
|
||||
headers["OpenAI-Project"] = project_id
|
||||
return headers
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
base_params: Final = super().get_supported_openai_params(model)
|
||||
try:
|
||||
if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider):
|
||||
if "reasoning_effort" not in base_params:
|
||||
base_params.append("reasoning_effort")
|
||||
except Exception as e:
|
||||
verbose_logger.debug("AmazonBedrockRuntimeChatCompletionsConfig: error checking reasoning support: %s", e)
|
||||
return base_params
|
||||
|
||||
def get_model_response_iterator(
|
||||
self,
|
||||
streaming_response: Iterator[str] | AsyncIterator[str] | Any,
|
||||
sync_stream: bool,
|
||||
json_mode: bool | None = False,
|
||||
) -> Any:
|
||||
from litellm.llms.openai.chat.gpt_transformation import (
|
||||
OpenAIChatCompletionStreamingHandler,
|
||||
)
|
||||
|
||||
return OpenAIChatCompletionStreamingHandler(
|
||||
streaming_response=streaming_response,
|
||||
sync_stream=sync_stream,
|
||||
json_mode=json_mode,
|
||||
)
|
||||
|
|
@ -780,6 +780,20 @@ def strip_bedrock_routing_prefix(model: str) -> str:
|
|||
return model
|
||||
|
||||
|
||||
def uses_bedrock_runtime_chat_completions(model: str) -> bool:
|
||||
"""Whether this Bedrock model should use runtime native Chat Completions.
|
||||
|
||||
Data-driven from the price-map ``use_bedrock_runtime_chat_completions`` flag
|
||||
so onboarding a model is a JSON change. Explicit ``converse/`` still wins in
|
||||
``get_bedrock_route`` because prefix routes are checked first.
|
||||
"""
|
||||
stripped: Final = strip_bedrock_routing_prefix(model)
|
||||
return any(
|
||||
(litellm.model_cost.get(key) or {}).get("use_bedrock_runtime_chat_completions") is True
|
||||
for key in (model, stripped)
|
||||
)
|
||||
|
||||
|
||||
def strip_bedrock_throughput_suffix(model: str) -> str:
|
||||
"""Strip throughput tier suffixes and context window suffixes from Bedrock model names."""
|
||||
import re
|
||||
|
|
@ -1107,6 +1121,7 @@ class BedrockModelInfo(BaseLLMModelInfo):
|
|||
"async_invoke",
|
||||
"openai",
|
||||
"mantle",
|
||||
"chat_completions",
|
||||
]:
|
||||
"""
|
||||
Get the bedrock route for the given model.
|
||||
|
|
@ -1123,6 +1138,7 @@ class BedrockModelInfo(BaseLLMModelInfo):
|
|||
"async_invoke",
|
||||
"openai",
|
||||
"mantle",
|
||||
"chat_completions",
|
||||
],
|
||||
] = {
|
||||
"invoke/": "invoke",
|
||||
|
|
@ -1152,6 +1168,9 @@ class BedrockModelInfo(BaseLLMModelInfo):
|
|||
if is_bedrock_application_inference_profile_arn(model):
|
||||
return "converse"
|
||||
|
||||
if uses_bedrock_runtime_chat_completions(model):
|
||||
return "chat_completions"
|
||||
|
||||
base_model: Final = BedrockModelInfo.get_base_model(model)
|
||||
alt_model: Final = BedrockModelInfo.get_non_litellm_routing_model_name(model=model)
|
||||
if base_model in litellm.bedrock_converse_models or alt_model in litellm.bedrock_converse_models:
|
||||
|
|
@ -1328,6 +1347,8 @@ def get_bedrock_chat_config(model: str):
|
|||
return litellm.AmazonConverseConfig()
|
||||
elif bedrock_route == "openai":
|
||||
return litellm.AmazonBedrockOpenAIConfig()
|
||||
elif bedrock_route == "chat_completions":
|
||||
return litellm.AmazonBedrockRuntimeChatCompletionsConfig()
|
||||
elif bedrock_route == "agent":
|
||||
from litellm.llms.bedrock.chat.invoke_agent.transformation import (
|
||||
AmazonInvokeAgentConfig,
|
||||
|
|
|
|||
|
|
@ -44546,6 +44546,7 @@
|
|||
"input_cost_per_token": 2.64e-06,
|
||||
"output_cost_per_token": 7.92e-06,
|
||||
"cache_read_input_token_cost": 6.6e-07,
|
||||
"use_bedrock_runtime_chat_completions": true,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 500000,
|
||||
"max_output_tokens": 500000,
|
||||
|
|
@ -55974,6 +55975,7 @@
|
|||
"input_cost_per_token": 2.2e-06,
|
||||
"output_cost_per_token": 6.6e-06,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
"use_bedrock_runtime_chat_completions": true,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 500000,
|
||||
"max_output_tokens": 500000,
|
||||
|
|
@ -55989,6 +55991,7 @@
|
|||
"input_cost_per_token": 2e-06,
|
||||
"output_cost_per_token": 6e-06,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"use_bedrock_runtime_chat_completions": true,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 500000,
|
||||
"max_output_tokens": 500000,
|
||||
|
|
|
|||
|
|
@ -44546,6 +44546,7 @@
|
|||
"input_cost_per_token": 2.64e-06,
|
||||
"output_cost_per_token": 7.92e-06,
|
||||
"cache_read_input_token_cost": 6.6e-07,
|
||||
"use_bedrock_runtime_chat_completions": true,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 500000,
|
||||
"max_output_tokens": 500000,
|
||||
|
|
@ -55974,6 +55975,7 @@
|
|||
"input_cost_per_token": 2.2e-06,
|
||||
"output_cost_per_token": 6.6e-06,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
"use_bedrock_runtime_chat_completions": true,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 500000,
|
||||
"max_output_tokens": 500000,
|
||||
|
|
@ -55989,6 +55991,7 @@
|
|||
"input_cost_per_token": 2e-06,
|
||||
"output_cost_per_token": 6e-06,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"use_bedrock_runtime_chat_completions": true,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 500000,
|
||||
"max_output_tokens": 500000,
|
||||
|
|
|
|||
|
|
@ -855,6 +855,9 @@
|
|||
"minimum": 0,
|
||||
"description": "Provider default tokens-per-minute limit."
|
||||
},
|
||||
"use_bedrock_runtime_chat_completions": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"use_openai_responses_path": {
|
||||
"type": "boolean"
|
||||
},
|
||||
|
|
|
|||
|
|
@ -0,0 +1,150 @@
|
|||
"""Native Bedrock Runtime Chat Completions: Grok stays on /openai/v1/chat/completions."""
|
||||
|
||||
import json
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock.chat.chat_completions.transformation import (
|
||||
AmazonBedrockRuntimeChatCompletionsConfig,
|
||||
)
|
||||
from litellm.llms.bedrock.common_utils import (
|
||||
BedrockModelInfo,
|
||||
get_bedrock_chat_config,
|
||||
uses_bedrock_runtime_chat_completions,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def local_cost_map(monkeypatch):
|
||||
original_model_cost = litellm.model_cost
|
||||
try:
|
||||
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
|
||||
litellm.model_cost = litellm.get_model_cost_map(url="")
|
||||
litellm.get_model_info.cache_clear()
|
||||
yield
|
||||
finally:
|
||||
litellm.model_cost = original_model_cost
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"us.xai.grok-4.6",
|
||||
"global.xai.grok-4.6",
|
||||
"us-gov.xai.grok-4.6",
|
||||
"bedrock/us.xai.grok-4.6",
|
||||
],
|
||||
)
|
||||
def test_grok_runtime_models_use_chat_completions_route(local_cost_map, model):
|
||||
assert uses_bedrock_runtime_chat_completions(model) is True
|
||||
assert BedrockModelInfo.get_bedrock_route(model) == "chat_completions"
|
||||
assert isinstance(get_bedrock_chat_config(model), AmazonBedrockRuntimeChatCompletionsConfig)
|
||||
|
||||
|
||||
def test_explicit_converse_prefix_still_uses_converse(local_cost_map):
|
||||
assert BedrockModelInfo.get_bedrock_route("bedrock/converse/us.xai.grok-4.6") == "converse"
|
||||
assert BedrockModelInfo.get_bedrock_route("converse/us.xai.grok-4.6") == "converse"
|
||||
|
||||
|
||||
def test_claude_stays_on_converse(local_cost_map):
|
||||
assert uses_bedrock_runtime_chat_completions("us.anthropic.claude-3-sonnet-20240229-v1:0") is False
|
||||
assert BedrockModelInfo.get_bedrock_route("us.anthropic.claude-3-sonnet-20240229-v1:0") == "converse"
|
||||
|
||||
|
||||
def test_flag_absent_means_no_chat_completions_route(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "model_cost", {"us.xai.grok-4.6": {"litellm_provider": "bedrock_converse"}})
|
||||
assert uses_bedrock_runtime_chat_completions("us.xai.grok-4.6") is False
|
||||
|
||||
|
||||
def test_complete_url_is_runtime_openai_chat_completions(monkeypatch):
|
||||
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
|
||||
monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False)
|
||||
cfg = AmazonBedrockRuntimeChatCompletionsConfig()
|
||||
url = cfg.get_complete_url(
|
||||
api_base=None,
|
||||
api_key=None,
|
||||
model="us.xai.grok-4.6",
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
assert url == "https://bedrock-runtime.us-east-1.amazonaws.com/openai/v1/chat/completions"
|
||||
|
||||
|
||||
def test_complete_url_appends_to_openai_v1_base():
|
||||
cfg = AmazonBedrockRuntimeChatCompletionsConfig()
|
||||
url = cfg.get_complete_url(
|
||||
api_base="https://bedrock-runtime.us-west-2.amazonaws.com/openai/v1",
|
||||
api_key=None,
|
||||
model="us.xai.grok-4.6",
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
assert url == "https://bedrock-runtime.us-west-2.amazonaws.com/openai/v1/chat/completions"
|
||||
|
||||
|
||||
def test_transform_request_is_openai_chat_body_not_converse():
|
||||
cfg = AmazonBedrockRuntimeChatCompletionsConfig()
|
||||
body = cfg.transform_request(
|
||||
model="bedrock/us.xai.grok-4.6",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
optional_params={"temperature": 0.2, "aws_region_name": "us-east-1"},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
assert body["model"] == "us.xai.grok-4.6"
|
||||
assert body["messages"] == [{"role": "user", "content": "hello"}]
|
||||
assert body["temperature"] == 0.2
|
||||
assert "aws_region_name" not in body
|
||||
assert "inferenceConfig" not in body
|
||||
assert "messages" in body
|
||||
|
||||
|
||||
def test_completion_posts_runtime_chat_completions(local_cost_map, monkeypatch):
|
||||
monkeypatch.setenv("AWS_REGION_NAME", "us-west-2")
|
||||
monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False)
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "testing")
|
||||
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "testing")
|
||||
monkeypatch.setenv("AWS_SESSION_TOKEN", "testing")
|
||||
|
||||
requests: list[dict] = []
|
||||
|
||||
def mock_post(self, url, data=None, json=None, headers=None, **kwargs):
|
||||
requests.append({"url": url, "data": data, "json": json, "headers": headers or {}})
|
||||
return httpx.Response(
|
||||
status_code=200,
|
||||
json={
|
||||
"id": "chatcmpl-test",
|
||||
"object": "chat.completion",
|
||||
"created": 1733529600,
|
||||
"model": "us.xai.grok-4.6",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "ok"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
request=httpx.Request("POST", url),
|
||||
)
|
||||
|
||||
with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post):
|
||||
response = litellm.completion(
|
||||
model="us.xai.grok-4.6",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "ok"
|
||||
assert len(requests) == 1
|
||||
assert requests[0]["url"] == "https://bedrock-runtime.us-west-2.amazonaws.com/openai/v1/chat/completions"
|
||||
raw = requests[0]["data"]
|
||||
body = json.loads(raw) if isinstance(raw, (str, bytes, bytearray)) else (requests[0]["json"] or {})
|
||||
assert body["model"] == "us.xai.grok-4.6"
|
||||
assert body["messages"] == [{"role": "user", "content": "hello"}]
|
||||
assert "inferenceConfig" not in body
|
||||
Loading…
Add table
Reference in a new issue