ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add support for bedrock in token counting api

Browse Source
This commit is contained in:
Raghav Jhavar 2026-01-09 17:08:07 +07:00
parent d9b275e62a
commit ba78194ff1
6 changed files with 363 additions and 64 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

142
litellm/llms/bedrock/common_utils.py
View File

@ -15,7 +15,7 @@ import litellm
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.secret_managers.main import get_secret
@ -359,6 +359,74 @@ def get_bedrock_tool_name(response_tool_name: str) -> str:
return response_tool_name
# ============================================================================
# Standalone model name utility functions
# ============================================================================
# Cache the global regions list at module level
_BEDROCK_GLOBAL_REGIONS: Optional[List[str]] = None
def _get_all_bedrock_regions() -> List[str]:
"""Get all Bedrock regions, cached at module level."""
global _BEDROCK_GLOBAL_REGIONS
if _BEDROCK_GLOBAL_REGIONS is None:
_BEDROCK_GLOBAL_REGIONS = AmazonBedrockGlobalConfig().get_all_regions()
return _BEDROCK_GLOBAL_REGIONS
def get_bedrock_cross_region_inference_regions() -> List[str]:
"""Abbreviations of regions AWS Bedrock supports for cross region inference."""
return ["global", "us", "eu", "apac", "jp", "au", "us-gov"]
def extract_model_name_from_bedrock_arn(model: str) -> str:
"""
Extract the model name from an AWS Bedrock ARN.
Returns the string after the last '/' if 'arn' is in the input string.
"""
if "arn" in model.lower():
return model.split("/")[-1]
return model
def strip_bedrock_routing_prefix(model: str) -> str:
"""Strip LiteLLM routing prefixes from model name."""
for prefix in ["bedrock/", "converse/", "invoke/", "openai/"]:
if model.startswith(prefix):
model = model.split("/", 1)[1]
return model
def get_bedrock_base_model(model: str) -> str:
"""
Get the base model from the given model name.
Handle model names like:
- "us.meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1"
- "bedrock/converse/model" -> "model"
"""
model = strip_bedrock_routing_prefix(model)
model = extract_model_name_from_bedrock_arn(model)
potential_region = model.split(".", 1)[0]
alt_potential_region = model.split("/", 1)[0]
if potential_region in get_bedrock_cross_region_inference_regions():
return model.split(".", 1)[1]
elif (
alt_potential_region in _get_all_bedrock_regions()
and len(model.split("/", 1)) > 1
):
return model.split("/", 1)[1]
return model
# Import after standalone functions to avoid circular imports
from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter
class BedrockModelInfo(BaseLLMModelInfo):
global_config = AmazonBedrockGlobalConfig()
all_global_regions = global_config.get_all_regions()
@ -394,76 +462,34 @@ class BedrockModelInfo(BaseLLMModelInfo):
) -> List[str]:
return []
@staticmethod
def extract_model_name_from_arn(model: str) -> str:
def get_token_counter(self) -> Optional[BaseTokenCounter]:
"""
Extract the model name from an AWS Bedrock ARN.
Returns the string after the last '/' if 'arn' is in the input string.
Args:
arn (str): The ARN string to parse
Factory method to create a Bedrock token counter.
Returns:
str: The extracted model name if 'arn' is in the string,
otherwise returns the original string
BedrockTokenCounter instance for this provider.
"""
if "arn" in model.lower():
return model.split("/")[-1]
return model
return BedrockTokenCounter()
@staticmethod
def extract_model_name_from_arn(model: str) -> str:
"""Wrapper for standalone function. See extract_model_name_from_bedrock_arn()."""
return extract_model_name_from_bedrock_arn(model)
@staticmethod
def get_non_litellm_routing_model_name(model: str) -> str:
if model.startswith("bedrock/"):
model = model.split("/", 1)[1]
if model.startswith("converse/"):
model = model.split("/", 1)[1]
if model.startswith("invoke/"):
model = model.split("/", 1)[1]
if model.startswith("openai/"):
model = model.split("/", 1)[1]
return model
"""Wrapper for standalone function. See strip_bedrock_routing_prefix()."""
return strip_bedrock_routing_prefix(model)
@staticmethod
def get_base_model(model: str) -> str:
"""
Get the base model from the given model name.
Handle model names like - "us.meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1"
AND "meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1"
"""
model = BedrockModelInfo.get_non_litellm_routing_model_name(model=model)
model = BedrockModelInfo.extract_model_name_from_arn(model)
potential_region = model.split(".", 1)[0]
alt_potential_region = model.split("/", 1)[
0
] # in model cost map we store regional information like `/us-west-2/bedrock-model`
if (
potential_region
in BedrockModelInfo._supported_cross_region_inference_region()
):
return model.split(".", 1)[1]
elif (
alt_potential_region in BedrockModelInfo.all_global_regions
and len(model.split("/", 1)) > 1
):
return model.split("/", 1)[1]
return model
"""Wrapper for standalone function. See get_bedrock_base_model()."""
return get_bedrock_base_model(model)
@staticmethod
def _supported_cross_region_inference_region() -> List[str]:
"""
Abbreviations of regions AWS Bedrock supports for cross region inference
"""
return ["global", "us", "eu", "apac", "jp", "au", "us-gov"]
"""Wrapper for standalone function. See get_bedrock_cross_region_inference_regions()."""
return get_bedrock_cross_region_inference_regions()
@staticmethod
def get_bedrock_route(

87
litellm/llms/bedrock/count_tokens/bedrock_token_counter.py Normal file
View File

@ -0,0 +1,87 @@
"""
Bedrock Token Counter implementation using the CountTokens API.
"""
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 get_bedrock_base_model
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
from litellm.types.utils import LlmProviders, TokenCountResponse
class BedrockTokenCounter(BaseTokenCounter):
"""Token counter implementation for AWS Bedrock provider using the CountTokens API."""
def should_use_token_counting_api(
self,
custom_llm_provider: Optional[str] = None,
) -> bool:
"""
Returns True if we should use the Bedrock CountTokens API for token counting.
"""
return custom_llm_provider == LlmProviders.BEDROCK.value
async def count_tokens(
self,
model_to_use: str,
messages: Optional[List[Dict[str, Any]]],
contents: Optional[List[Dict[str, Any]]],
deployment: Optional[Dict[str, Any]] = None,
request_model: str = "",
) -> Optional[TokenCountResponse]:
"""
Count tokens using AWS Bedrock's CountTokens API.
This method calls the existing BedrockCountTokensHandler to make an API call
to Bedrock's token counting endpoint, bypassing the local tiktoken-based counting.
Args:
model_to_use: The model identifier
messages: The messages to count tokens for
contents: Alternative content format (not used for Bedrock)
deployment: Deployment configuration containing litellm_params
request_model: The original request model name
Returns:
TokenCountResponse with token count, or None if counting fails
"""
if not messages:
return None
deployment = deployment or {}
litellm_params = deployment.get("litellm_params", {})
# Build request data in the format expected by BedrockCountTokensHandler
request_data = {
"model": model_to_use,
"messages": messages,
}
# Get the resolved model (strip prefixes like bedrock/, converse/, etc.)
resolved_model = get_bedrock_base_model(model_to_use)
try:
handler = BedrockCountTokensHandler()
result = await handler.handle_count_tokens_request(
request_data=request_data,
litellm_params=litellm_params,
resolved_model=resolved_model,
)
# Transform response to TokenCountResponse
if result is not None:
return TokenCountResponse(
total_tokens=result.get("input_tokens", 0),
request_model=request_model,
model_used=model_to_use,
tokenizer_type="bedrock_api",
original_response=result,
)
except Exception as e:
verbose_logger.warning(
f"Error calling Bedrock CountTokens API: {e}, falling back to default tokenizer"
)
return None

3
litellm/llms/bedrock/count_tokens/handler.py
View File

@ -70,6 +70,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
verbose_logger.debug(f"Making request to: {endpoint_url}")
# Use existing _sign_request method from BaseAWSLLM
# Extract api_key for bearer token auth if provided
api_key = litellm_params.get("api_key", None)
headers = {"Content-Type": "application/json"}
signed_headers, signed_body = self._sign_request(
service_name="bedrock",
@ -78,6 +80,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
request_data=bedrock_request,
api_base=endpoint_url,
model=resolved_model,
api_key=api_key,
)
async_client = get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK)

4
litellm/llms/bedrock/count_tokens/transformation.py
View File

@ -8,7 +8,7 @@ to AWS Bedrock's CountTokens API format and vice versa.
from typing import Any, Dict, List
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import BedrockModelInfo
from litellm.llms.bedrock.common_utils import get_bedrock_base_model
class BedrockCountTokensConfig(BaseAWSLLM):
@ -141,7 +141,7 @@ class BedrockCountTokensConfig(BaseAWSLLM):
Complete endpoint URL for CountTokens API
"""
# Use existing LiteLLM function to get the base model ID (removes region prefix)
model_id = BedrockModelInfo.get_base_model(model)
model_id = get_bedrock_base_model(model)
# Remove bedrock/ prefix if present
if model_id.startswith("bedrock/"):

9
litellm/utils.py
View File

@ -47,7 +47,6 @@ from tiktoken import Encoding
from tokenizers import Tokenizer
import litellm
import litellm.litellm_core_utils
# audio_utils.utils is lazy-loaded - only imported when needed for transcription calls
import litellm.litellm_core_utils.json_validation_rule
@ -291,7 +290,7 @@ if TYPE_CHECKING:
from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig
from litellm.llms.base_llm.search.transformation import BaseSearchConfig
from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig
from litellm.llms.bedrock.common_utils import BedrockModelInfo
from litellm.llms.bedrock.bedrock_model_info import BedrockModelInfo
from litellm.llms.cohere.common_utils import CohereModelInfo
from litellm.llms.mistral.ocr.transformation import MistralOCRConfig
# Type stubs for lazy-loaded functions and classes
@ -4954,7 +4953,7 @@ def _get_base_bedrock_model(model_name) -> str:
Handle model names like - "us.meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1"
AND "meta.llama3-2-11b-instruct-v1:0" -> "meta.llama3-2-11b-instruct-v1"
"""
from litellm.llms.bedrock.common_utils import BedrockModelInfo
from litellm.llms.bedrock.bedrock_model_info import BedrockModelInfo
return BedrockModelInfo.get_base_model(model_name)
@ -7779,7 +7778,7 @@ class ProviderConfigManager:
# The 'BEDROCK' provider corresponds to Amazon's implementation of Anthropic Claude v3.
# This mapping ensures that the correct configuration is returned for BEDROCK.
elif litellm.LlmProviders.BEDROCK == provider:
from litellm.llms.bedrock.common_utils import BedrockModelInfo
from litellm.llms.bedrock.bedrock_model_info import BedrockModelInfo
return BedrockModelInfo.get_bedrock_provider_config_for_messages_api(model)
elif litellm.LlmProviders.VERTEX_AI == provider:
@ -7941,6 +7940,8 @@ class ProviderConfigManager:
return litellm.LemonadeChatConfig()
elif LlmProviders.CLARIFAI == provider:
return litellm.ClarifaiConfig()
elif LlmProviders.BEDROCK == provider:
return litellm.llms.bedrock.common_utils.BedrockModelInfo()
return None
@staticmethod

182
tests/llm_translation/test_bedrock_common_utils.py Normal file
View File

@ -0,0 +1,182 @@
"""
Unit tests for litellm/llms/bedrock/common_utils.py
Tests the standalone model name utility functions and BedrockTokenCounter.
"""
import pytest
from litellm.llms.bedrock.common_utils import (
BedrockModelInfo,
extract_model_name_from_bedrock_arn,
get_bedrock_base_model,
get_bedrock_cross_region_inference_regions,
strip_bedrock_routing_prefix,
)
from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter
class TestStripBedrockRoutingPrefix:
"""Tests for strip_bedrock_routing_prefix function."""
def test_strips_bedrock_prefix(self):
assert strip_bedrock_routing_prefix("bedrock/claude-3-sonnet") == "claude-3-sonnet"
def test_strips_converse_prefix(self):
assert strip_bedrock_routing_prefix("converse/claude-3-sonnet") == "claude-3-sonnet"
def test_strips_invoke_prefix(self):
assert strip_bedrock_routing_prefix("invoke/claude-3-sonnet") == "claude-3-sonnet"
def test_strips_openai_prefix(self):
assert strip_bedrock_routing_prefix("openai/gpt-4") == "gpt-4"
def test_strips_all_known_prefixes(self):
# Function strips all known prefixes iteratively
# bedrock/converse/model -> converse/model -> model
assert strip_bedrock_routing_prefix("bedrock/converse/claude-3") == "claude-3"
def test_no_prefix_unchanged(self):
assert strip_bedrock_routing_prefix("claude-3-sonnet") == "claude-3-sonnet"
def test_model_with_dots_unchanged(self):
assert (
strip_bedrock_routing_prefix("anthropic.claude-3-sonnet-20240229-v1:0")
== "anthropic.claude-3-sonnet-20240229-v1:0"
)
class TestExtractModelNameFromBedrockArn:
"""Tests for extract_model_name_from_bedrock_arn function."""
def test_extracts_from_provisioned_model_arn(self):
arn = "arn:aws:bedrock:us-east-1:123456789012:provisioned-model/my-model-id"
assert extract_model_name_from_bedrock_arn(arn) == "my-model-id"
def test_extracts_from_foundation_model_arn(self):
arn = "arn:aws:bedrock:us-west-2:123456789012:foundation-model/anthropic.claude-v2"
assert extract_model_name_from_bedrock_arn(arn) == "anthropic.claude-v2"
def test_non_arn_unchanged(self):
model = "anthropic.claude-3-sonnet-20240229-v1:0"
assert extract_model_name_from_bedrock_arn(model) == model
def test_case_insensitive_arn_detection(self):
arn = "ARN:aws:bedrock:us-east-1:123456789012:model/my-model"
assert extract_model_name_from_bedrock_arn(arn) == "my-model"
class TestGetBedrockCrossRegionInferenceRegions:
"""Tests for get_bedrock_cross_region_inference_regions function."""
def test_returns_expected_regions(self):
regions = get_bedrock_cross_region_inference_regions()
assert "us" in regions
assert "eu" in regions
assert "global" in regions
assert "apac" in regions
def test_returns_list(self):
regions = get_bedrock_cross_region_inference_regions()
assert isinstance(regions, list)
class TestGetBedrockBaseModel:
"""Tests for get_bedrock_base_model function."""
def test_strips_bedrock_prefix(self):
assert get_bedrock_base_model("bedrock/claude-3-sonnet") == "claude-3-sonnet"
def test_strips_converse_prefix(self):
assert get_bedrock_base_model("bedrock/converse/claude-3-sonnet") == "claude-3-sonnet"
def test_strips_us_region_prefix(self):
# us.anthropic.model -> anthropic.model
assert (
get_bedrock_base_model("us.anthropic.claude-3-sonnet-20240229-v1:0")
== "anthropic.claude-3-sonnet-20240229-v1:0"
)
def test_strips_eu_region_prefix(self):
assert (
get_bedrock_base_model("eu.anthropic.claude-3-sonnet-20240229-v1:0")
== "anthropic.claude-3-sonnet-20240229-v1:0"
)
def test_extracts_from_arn(self):
arn = "arn:aws:bedrock:us-east-1:123456789012:provisioned-model/my-model"
assert get_bedrock_base_model(arn) == "my-model"
def test_model_without_prefix_unchanged(self):
model = "anthropic.claude-3-sonnet-20240229-v1:0"
assert get_bedrock_base_model(model) == model
def test_combined_bedrock_and_region_prefix(self):
# bedrock/us.anthropic.model -> anthropic.model
assert (
get_bedrock_base_model("bedrock/us.anthropic.claude-3-sonnet-20240229-v1:0")
== "anthropic.claude-3-sonnet-20240229-v1:0"
)
class TestBedrockModelInfoWrappers:
"""Tests that BedrockModelInfo methods correctly wrap standalone functions."""
def test_get_base_model_matches_standalone(self):
test_cases = [
"bedrock/claude-3-sonnet",
"us.anthropic.claude-3-sonnet-20240229-v1:0",
"arn:aws:bedrock:us-east-1:123:model/my-model",
]
for model in test_cases:
assert BedrockModelInfo.get_base_model(model) == get_bedrock_base_model(model)
def test_extract_model_name_from_arn_matches_standalone(self):
arn = "arn:aws:bedrock:us-east-1:123456789012:provisioned-model/my-model"
assert (
BedrockModelInfo.extract_model_name_from_arn(arn)
== extract_model_name_from_bedrock_arn(arn)
)
def test_get_non_litellm_routing_model_name_matches_standalone(self):
model = "bedrock/converse/claude-3"
assert (
BedrockModelInfo.get_non_litellm_routing_model_name(model)
== strip_bedrock_routing_prefix(model)
)
class TestBedrockTokenCounter:
"""Tests for BedrockTokenCounter class."""
def test_should_use_token_counting_api_for_bedrock(self):
counter = BedrockTokenCounter()
assert counter.should_use_token_counting_api("bedrock") is True
def test_should_not_use_token_counting_api_for_other_providers(self):
counter = BedrockTokenCounter()
assert counter.should_use_token_counting_api("openai") is False
assert counter.should_use_token_counting_api("anthropic") is False
assert counter.should_use_token_counting_api(None) is False
def test_get_token_counter_returns_bedrock_token_counter(self):
model_info = BedrockModelInfo()
token_counter = model_info.get_token_counter()
assert isinstance(token_counter, BedrockTokenCounter)
@pytest.mark.asyncio
async def test_count_tokens_returns_none_for_empty_messages(self):
counter = BedrockTokenCounter()
result = await counter.count_tokens(
model_to_use="anthropic.claude-3-sonnet",
messages=None,
contents=None,
)
assert result is None
result = await counter.count_tokens(
model_to_use="anthropic.claude-3-sonnet",
messages=[],
contents=None,
)
assert result is None