scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165

fix(bedrock): keep cross-region inference-profile prefix on count_tokens URL

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
milan 2026-07-22 22:19:15 +00:00
parent 8d217a4d5f
commit b859a53c6d
2 changed files with 57 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
litellm/llms/bedrock/count_tokens/transformation.py
View file

@ -9,7 +9,11 @@ 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 (
extract_model_name_from_bedrock_arn,
strip_bedrock_routing_prefix,
strip_bedrock_throughput_suffix,
)
# Placeholder satisfying the Anthropic InvokeModel schema's required
# max_tokens field; CountTokens only counts input, so it has no effect
@ -207,12 +211,9 @@ 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 = strip_bedrock_routing_prefix(model)
model_id = extract_model_name_from_bedrock_arn(model_id)
model_id = strip_bedrock_throughput_suffix(model_id)
encoded_model_id = self.encode_model_id(model_id=model_id)
base_url, _ = self.get_runtime_endpoint(

49
tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py
View file

@ -258,3 +258,52 @@ def test_count_tokens_endpoint_encodes_model_id(monkeypatch):
endpoint
== "https://bedrock-runtime.us-east-1.amazonaws.com/model/..%2F..%2Fmodel%2Fother%3Fx%3D1%23frag/count-tokens"
)
def test_count_tokens_endpoint_keeps_cross_region_inference_profile_prefix(monkeypatch):
"""Cross-region inference-profile prefixes (global./us./eu./apac./...) must be
preserved in the count-tokens URL. Those models are inference-profile-only and
Bedrock rejects the bare foundation-model ID with a 400 on the count-tokens
route (issue #32683)."""
config = BedrockCountTokensConfig()
monkeypatch.setattr(
config,
"get_runtime_endpoint",
lambda **kwargs: ("https://bedrock-runtime.eu-central-1.amazonaws.com", None),
)
assert (
config.get_bedrock_count_tokens_endpoint(
model="bedrock/global.anthropic.claude-opus-4-8",
aws_region_name="eu-central-1",
)
== "https://bedrock-runtime.eu-central-1.amazonaws.com/model/global.anthropic.claude-opus-4-8/count-tokens"
)
assert (
config.get_bedrock_count_tokens_endpoint(
model="bedrock/eu.anthropic.claude-sonnet-4-6",
aws_region_name="eu-central-1",
)
== "https://bedrock-runtime.eu-central-1.amazonaws.com/model/eu.anthropic.claude-sonnet-4-6/count-tokens"
)
def test_count_tokens_endpoint_strips_throughput_suffix(monkeypatch):
"""Throughput / context-window suffixes are not part of the Bedrock model ID."""
config = BedrockCountTokensConfig()
monkeypatch.setattr(
config,
"get_runtime_endpoint",
lambda **kwargs: ("https://bedrock-runtime.us-east-1.amazonaws.com", None),
)
assert (
config.get_bedrock_count_tokens_endpoint(
model="anthropic.claude-3-5-sonnet-20241022-v2:0:51k",
aws_region_name="us-east-1",
)
== "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-5-sonnet-20241022-v2%3A0/count-tokens"
)