From bc4608e718320c2f3930dcf909514f52d93cd6e9 Mon Sep 17 00:00:00 2001 From: stias Date: Fri, 20 Mar 2026 17:58:20 +0900 Subject: [PATCH 1/4] fix(bedrock): respect api_base and aws_bedrock_runtime_endpoint in count_tokens endpoint The /v1/messages/count_tokens endpoint was hardcoding the Bedrock runtime URL, ignoring api_base and aws_bedrock_runtime_endpoint settings. This aligns it with invoke/converse handlers by using the existing get_runtime_endpoint() method for consistent endpoint resolution. Signed-off-by: stias --- litellm/llms/bedrock/common_utils.py | 1 + litellm/llms/bedrock/count_tokens/handler.py | 9 ++- .../bedrock/count_tokens/transformation.py | 14 +++- .../test_bedrock_token_counter.py | 64 +++++++++++++++++++ 4 files changed, 85 insertions(+), 3 deletions(-) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 9666aa68c9..6e659f06d5 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -322,6 +322,7 @@ def init_bedrock_client( endpoint_url=endpoint_url, config=config, verify=ssl_verify, + ) elif aws_profile_name is not None: # uses auth values from AWS profile usually stored in ~/.aws/credentials 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..7239ec0bcf 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 AsyncMock, MagicMock, 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.startswith(api_base) + + 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") From 72902c39c511d33fdc1c13310633f46c37eadfc5 Mon Sep 17 00:00:00 2001 From: Seokjun Yang Date: Fri, 20 Mar 2026 18:09:52 +0900 Subject: [PATCH 2/4] Remove extra newline in common_utils.py --- litellm/llms/bedrock/common_utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 6e659f06d5..9666aa68c9 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -322,7 +322,6 @@ def init_bedrock_client( endpoint_url=endpoint_url, config=config, verify=ssl_verify, - ) elif aws_profile_name is not None: # uses auth values from AWS profile usually stored in ~/.aws/credentials From eb733702fcd5a1c3ddd57ffc2416005f5d51cd8a Mon Sep 17 00:00:00 2001 From: Seokjun Yang Date: Fri, 20 Mar 2026 22:21:15 +0900 Subject: [PATCH 3/4] Update tests/litellm_utils_tests/test_bedrock_token_counter.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/litellm_utils_tests/test_bedrock_token_counter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/litellm_utils_tests/test_bedrock_token_counter.py b/tests/litellm_utils_tests/test_bedrock_token_counter.py index 7239ec0bcf..b0e37914af 100644 --- a/tests/litellm_utils_tests/test_bedrock_token_counter.py +++ b/tests/litellm_utils_tests/test_bedrock_token_counter.py @@ -150,7 +150,7 @@ class TestBedrockCountTokensEndpoint: api_base=api_base, aws_bedrock_runtime_endpoint=runtime_endpoint, ) - assert url.startswith(api_base) + assert url == f"{api_base}/model/amazon.nova-lite-v1:0/count-tokens" def test_env_var_overrides_default(self, monkeypatch): monkeypatch.setenv( From d3afaf613dc0865ad92f374f41979cb3bee1c9e7 Mon Sep 17 00:00:00 2001 From: Seokjun Yang Date: Fri, 20 Mar 2026 22:21:22 +0900 Subject: [PATCH 4/4] Update tests/litellm_utils_tests/test_bedrock_token_counter.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/litellm_utils_tests/test_bedrock_token_counter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/litellm_utils_tests/test_bedrock_token_counter.py b/tests/litellm_utils_tests/test_bedrock_token_counter.py index b0e37914af..abc45b03d6 100644 --- a/tests/litellm_utils_tests/test_bedrock_token_counter.py +++ b/tests/litellm_utils_tests/test_bedrock_token_counter.py @@ -11,7 +11,7 @@ counting, the test will be skipped. import os import sys from typing import Any, Dict, List -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import patch import pytest