mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(bedrock): keep cross-region inference-profile prefix for count_tokens
get_bedrock_base_model strips the cross-region inference-profile prefix (global./us./eu./apac./...), so the count-tokens URL used the bare foundation-model ID. For models that are inference-profile-only in a region, Bedrock rejects that with a 400. Resolve the count-tokens/invocation model ID with a new get_bedrock_invocation_model_id helper that preserves the prefix while still stripping routing prefixes, ARNs, and throughput/context-window suffixes.
This commit is contained in:
parent
1fa200123f
commit
ba0caf1ad8
5 changed files with 91 additions and 9 deletions
|
|
@ -685,6 +685,41 @@ def get_bedrock_base_model(model: str) -> str:
|
|||
return model
|
||||
|
||||
|
||||
def get_bedrock_invocation_model_id(model: str) -> str:
|
||||
"""
|
||||
Resolve the Bedrock model ID for invocation URLs (e.g. ``/model/{id}/converse``
|
||||
or ``/model/{id}/count-tokens``).
|
||||
|
||||
Unlike ``get_bedrock_base_model``, this keeps the cross-region
|
||||
inference-profile prefix (``global.`` / ``us.`` / ``eu.`` / ``apac.`` / ...).
|
||||
Bedrock requires the inference-profile ID for models that are
|
||||
inference-profile-only in a region, so stripping the prefix would produce a
|
||||
bare foundation-model ID that Bedrock rejects with a 400 (see issue #32683).
|
||||
It still strips LiteLLM routing prefixes, resolves ARNs, drops
|
||||
throughput/context-window suffixes, and removes an embedded full-region path
|
||||
prefix (e.g. ``us-east-1/model``).
|
||||
"""
|
||||
stripped = model
|
||||
for rp in ["bedrock/converse/", "bedrock/", "converse/"]:
|
||||
if stripped.startswith(rp):
|
||||
stripped = stripped[len(rp) :]
|
||||
break
|
||||
if stripped.startswith("nova-2/"):
|
||||
return "amazon.nova-2-custom"
|
||||
elif stripped.startswith("nova/"):
|
||||
return "amazon.nova-custom"
|
||||
|
||||
model = strip_bedrock_routing_prefix(model)
|
||||
model = extract_model_name_from_bedrock_arn(model)
|
||||
model = strip_bedrock_throughput_suffix(model)
|
||||
|
||||
alt_potential_region = model.split("/", 1)[0]
|
||||
if alt_potential_region in _get_all_bedrock_regions() and len(model.split("/", 1)) > 1:
|
||||
return model.split("/", 1)[1]
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def bedrock_converse_supports_parallel_tool_use_config(model: str) -> bool:
|
||||
return any(
|
||||
(litellm.model_cost.get(candidate) or {}).get("supports_parallel_tool_use_config") is True
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from typing import Any, Dict, List, Optional
|
|||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.base_llm.base_utils import BaseTokenCounter
|
||||
from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model
|
||||
from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_invocation_model_id
|
||||
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
|
||||
from litellm.types.utils import LlmProviders, TokenCountResponse
|
||||
|
||||
|
|
@ -68,7 +68,7 @@ class BedrockTokenCounter(BaseTokenCounter):
|
|||
request_data["system"] = system
|
||||
|
||||
# Get the resolved model (strip prefixes like bedrock/, converse/, etc.)
|
||||
resolved_model = get_bedrock_base_model(model_to_use)
|
||||
resolved_model = get_bedrock_invocation_model_id(model_to_use)
|
||||
|
||||
try:
|
||||
handler = BedrockCountTokensHandler()
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ import re
|
|||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock.common_utils import get_bedrock_base_model
|
||||
from litellm.llms.bedrock.common_utils import get_bedrock_invocation_model_id
|
||||
|
||||
# Placeholder satisfying the Anthropic InvokeModel schema's required
|
||||
# max_tokens field; CountTokens only counts input, so it has no effect
|
||||
|
|
@ -207,12 +207,7 @@ class BedrockCountTokensConfig(BaseAWSLLM):
|
|||
Returns:
|
||||
Complete endpoint URL for CountTokens API
|
||||
"""
|
||||
# Use existing LiteLLM function to get the base model ID (removes region prefix)
|
||||
model_id = get_bedrock_base_model(model)
|
||||
|
||||
# Remove bedrock/ prefix if present
|
||||
if model_id.startswith("bedrock/"):
|
||||
model_id = model_id[8:] # Remove "bedrock/" prefix
|
||||
model_id = get_bedrock_invocation_model_id(model)
|
||||
encoded_model_id = self.encode_model_id(model_id=model_id)
|
||||
|
||||
base_url, _ = self.get_runtime_endpoint(
|
||||
|
|
|
|||
|
|
@ -167,6 +167,22 @@ class TestBedrockCountTokensEndpoint:
|
|||
)
|
||||
assert url == f"{api_base}/model/amazon.nova-lite-v1%3A0/count-tokens"
|
||||
|
||||
def test_cross_region_inference_profile_prefix_preserved(self):
|
||||
"""
|
||||
Regression test for #32683: the count-tokens URL must keep the
|
||||
cross-region inference-profile prefix (e.g. global.) so Bedrock does not
|
||||
reject inference-profile-only models with a 400.
|
||||
"""
|
||||
handler = self._make_handler()
|
||||
url = handler.get_bedrock_count_tokens_endpoint(
|
||||
model="bedrock/global.anthropic.claude-opus-4-8",
|
||||
aws_region_name="eu-central-1",
|
||||
)
|
||||
assert (
|
||||
url
|
||||
== "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-opus-4-8/count-tokens"
|
||||
)
|
||||
|
||||
def test_env_var_overrides_default(self, monkeypatch):
|
||||
monkeypatch.setenv(
|
||||
"AWS_BEDROCK_RUNTIME_ENDPOINT",
|
||||
|
|
|
|||
|
|
@ -285,6 +285,42 @@ def test_context_window_suffix_stripped_for_cost_lookup():
|
|||
)
|
||||
|
||||
|
||||
def test_get_bedrock_invocation_model_id_preserves_cross_region_prefix():
|
||||
"""
|
||||
Regression test for #32683: the count-tokens / invocation model ID must keep
|
||||
the cross-region inference-profile prefix (global./us./eu./apac./...) so
|
||||
Bedrock does not reject inference-profile-only models with a 400.
|
||||
"""
|
||||
from litellm.llms.bedrock.common_utils import get_bedrock_invocation_model_id
|
||||
|
||||
assert (
|
||||
get_bedrock_invocation_model_id("bedrock/global.anthropic.claude-opus-4-8")
|
||||
== "global.anthropic.claude-opus-4-8"
|
||||
)
|
||||
assert (
|
||||
get_bedrock_invocation_model_id("eu.anthropic.claude-sonnet-4-6")
|
||||
== "eu.anthropic.claude-sonnet-4-6"
|
||||
)
|
||||
assert (
|
||||
get_bedrock_invocation_model_id("us.meta.llama3-2-11b-instruct-v1:0")
|
||||
== "us.meta.llama3-2-11b-instruct-v1:0"
|
||||
)
|
||||
# Non-prefixed foundation models are unchanged
|
||||
assert (
|
||||
get_bedrock_invocation_model_id("anthropic.claude-opus-4-8")
|
||||
== "anthropic.claude-opus-4-8"
|
||||
)
|
||||
# Routing prefix and throughput/context suffixes are still stripped
|
||||
assert (
|
||||
get_bedrock_invocation_model_id("bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0:51k")
|
||||
== "anthropic.claude-3-5-sonnet-20241022-v2:0"
|
||||
)
|
||||
assert (
|
||||
get_bedrock_invocation_model_id("global.anthropic.claude-opus-4-5-20251101-v1:0[1m]")
|
||||
== "global.anthropic.claude-opus-4-5-20251101-v1:0"
|
||||
)
|
||||
|
||||
|
||||
def test_output_config_effort_normalization_uses_model_info_ceiling(monkeypatch):
|
||||
import litellm.llms.bedrock.common_utils as mod
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue