Merge pull request #24199 from stias/fix/bedrock-count-tokens-custom-endpoint

fix(bedrock): respect api_base and aws_bedrock_runtime_endpoint in count_tokens endpoint

litellm/llms/bedrock/count_tokens/handler.py

@@ -64,8 +64,15 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
             verbose_logger.debug(f"Transformed request: {bedrock_request}")
 
             # Get endpoint URL using simplified function
+            api_base = litellm_params.get("api_base", None)
+            aws_bedrock_runtime_endpoint = litellm_params.get(
+                "aws_bedrock_runtime_endpoint", None
+            )
             endpoint_url = self.get_bedrock_count_tokens_endpoint(
-                resolved_model, aws_region_name
+                model=resolved_model,
+                aws_region_name=aws_region_name,
+                api_base=api_base,
+                aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
             )
 
             verbose_logger.debug(f"Making request to: {endpoint_url}")

litellm/llms/bedrock/count_tokens/transformation.py

@@ -177,7 +177,11 @@ class BedrockCountTokensConfig(BaseAWSLLM):
         return {"input": {"invokeModel": {"body": json.dumps(body_data)}}}
 
     def get_bedrock_count_tokens_endpoint(
-        self, model: str, aws_region_name: str
+        self,
+        model: str,
+        aws_region_name: str,
+        api_base: Optional[str] = None,
+        aws_bedrock_runtime_endpoint: Optional[str] = None,
     ) -> str:
         """
         Construct the AWS Bedrock CountTokens API endpoint using existing LiteLLM functions.
@@ -185,6 +189,8 @@ class BedrockCountTokensConfig(BaseAWSLLM):
         Args:
             model: The resolved model ID from router lookup
             aws_region_name: AWS region (e.g., "eu-west-1")
+            api_base: Optional custom API base URL (takes highest priority)
+            aws_bedrock_runtime_endpoint: Optional custom Bedrock runtime endpoint
 
         Returns:
             Complete endpoint URL for CountTokens API
@@ -196,7 +202,11 @@ class BedrockCountTokensConfig(BaseAWSLLM):
         if model_id.startswith("bedrock/"):
             model_id = model_id[8:]  # Remove "bedrock/" prefix
 
-        base_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com"
+        base_url, _ = self.get_runtime_endpoint(
+            api_base=api_base,
+            aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
+            aws_region_name=aws_region_name,
+        )
         endpoint = f"{base_url}/model/{model_id}/count-tokens"
 
         return endpoint

tests/litellm_utils_tests/test_bedrock_token_counter.py

@@ -11,6 +11,7 @@ counting, the test will be skipped.
 import os
 import sys
 from typing import Any, Dict, List
+from unittest.mock import patch
 
 import pytest
 
@@ -99,3 +100,66 @@ class TestBedrockTokenCounter(BaseTokenCounterTest):
         assert result.total_tokens > 0, f"Token count should be > 0, got {result.total_tokens}"
         assert result.tokenizer_type is not None, "tokenizer_type should be set"
         assert result.error is not True, f"Token counting should not error: {result.error_message}"
+
+
+class TestBedrockCountTokensEndpoint:
+    """Unit tests for custom endpoint URL resolution in BedrockCountTokensConfig."""
+
+    def _make_handler(self):
+        from litellm.llms.bedrock.count_tokens.transformation import (
+            BedrockCountTokensConfig,
+        )
+
+        return BedrockCountTokensConfig()
+
+    def test_default_endpoint(self):
+        handler = self._make_handler()
+        url = handler.get_bedrock_count_tokens_endpoint(
+            model="amazon.nova-lite-v1:0",
+            aws_region_name="us-east-1",
+        )
+        assert url == "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.nova-lite-v1:0/count-tokens"
+
+    def test_api_base_overrides_default(self):
+        handler = self._make_handler()
+        custom_base = "https://vpce-xxx.bedrock-runtime.us-east-1.vpce.amazonaws.com"
+        url = handler.get_bedrock_count_tokens_endpoint(
+            model="amazon.nova-lite-v1:0",
+            aws_region_name="us-east-1",
+            api_base=custom_base,
+        )
+        assert url == f"{custom_base}/model/amazon.nova-lite-v1:0/count-tokens"
+
+    def test_aws_bedrock_runtime_endpoint_overrides_default(self):
+        handler = self._make_handler()
+        custom_endpoint = "https://vpce-yyy.bedrock-runtime.eu-west-1.vpce.amazonaws.com"
+        url = handler.get_bedrock_count_tokens_endpoint(
+            model="amazon.nova-lite-v1:0",
+            aws_region_name="eu-west-1",
+            aws_bedrock_runtime_endpoint=custom_endpoint,
+        )
+        assert url == f"{custom_endpoint}/model/amazon.nova-lite-v1:0/count-tokens"
+
+    def test_api_base_takes_priority_over_aws_bedrock_runtime_endpoint(self):
+        handler = self._make_handler()
+        api_base = "https://api-base.example.com"
+        runtime_endpoint = "https://runtime-endpoint.example.com"
+        url = handler.get_bedrock_count_tokens_endpoint(
+            model="amazon.nova-lite-v1:0",
+            aws_region_name="us-east-1",
+            api_base=api_base,
+            aws_bedrock_runtime_endpoint=runtime_endpoint,
+        )
+        assert url == f"{api_base}/model/amazon.nova-lite-v1:0/count-tokens"
+
+    def test_env_var_overrides_default(self, monkeypatch):
+        monkeypatch.setenv(
+            "AWS_BEDROCK_RUNTIME_ENDPOINT",
+            "https://env-endpoint.bedrock-runtime.us-west-2.amazonaws.com",
+        )
+        handler = self._make_handler()
+        url = handler.get_bedrock_count_tokens_endpoint(
+            model="amazon.nova-lite-v1:0",
+            aws_region_name="us-west-2",
+        )
+        assert url.startswith("https://env-endpoint.bedrock-runtime.us-west-2.amazonaws.com")