add support for bedrock in token counting api

litellm/llms/bedrock/common_utils.py

@@ -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(

litellm/llms/bedrock/count_tokens/bedrock_token_counter.py

@@ -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

litellm/llms/bedrock/count_tokens/handler.py

@@ -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)

litellm/llms/bedrock/count_tokens/transformation.py

@@ -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/"):

litellm/utils.py

@@ -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

tests/llm_translation/test_bedrock_common_utils.py

@@ -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