diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index cfd32342d1..8c227c853c 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/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}") diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index fe9ab80ced..a37af13162 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/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 diff --git a/tests/litellm_utils_tests/test_bedrock_token_counter.py b/tests/litellm_utils_tests/test_bedrock_token_counter.py index f7c2991882..abc45b03d6 100644 --- a/tests/litellm_utils_tests/test_bedrock_token_counter.py +++ b/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")