From f35b1698f0ee34c1d87382b1899e245e2d63549e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 2 Jul 2026 10:46:37 -0700 Subject: [PATCH] feat(proxy): add per-model use_local_token_counter to skip provider CountTokens API Token counting for a deployment could not be forced to stay local. When a Bedrock (or other provider) model is counted via /v1/messages/count_tokens or /utils/token_counter?call_endpoint=true, the proxy always dispatched to the provider's remote CountTokens API, and the only global switch (disable_token_counter) turns counting off entirely rather than falling back to the local tokenizer. Add a per-model model_info.use_local_token_counter flag. When set, the proxy resolves no provider token counter for that deployment and falls through to the local tokenizer (tiktoken/HF/anthropic-local depending on the model), so no remote CountTokens call is made. model_list: - model_name: haiku-4.5 litellm_params: model: bedrock/anthropic.claude-haiku-4-5-... model_info: use_local_token_counter: true --- litellm/proxy/proxy_server.py | 4 + litellm/types/router.py | 2 + .../test_proxy_token_counter.py | 83 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 4 files changed, 91 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a29acf3b06f..eb8833da593 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10205,6 +10205,10 @@ def _get_provider_token_counter( if deployment is None: return None + model_info = deployment.get("model_info") or {} + if model_info.get("use_local_token_counter"): + return None, None, None + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider full_model = deployment.get("litellm_params", {}).get("model", "") diff --git a/litellm/types/router.py b/litellm/types/router.py index 0c3485deae7..ef1c3a3d0ef 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -133,6 +133,8 @@ class ModelInfo(BaseModel): base_model: Optional[str] = None # specify if the base model is azure/gpt-3.5-turbo etc for accurate cost tracking tier: Optional[Literal["free", "paid"]] = None + use_local_token_counter: Optional[bool] = None + """ Team Model Specific Fields """ diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index 1079a5228a1..f3d6df7f0d6 100644 --- a/tests/proxy_unit_tests/test_proxy_token_counter.py +++ b/tests/proxy_unit_tests/test_proxy_token_counter.py @@ -1364,3 +1364,86 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): finally: anthropic_endpoints._read_request_body = original_read_request_body proxy_server.token_counter = original_token_counter + + +@pytest.mark.asyncio +async def test_get_provider_token_counter_respects_use_local_token_counter(): + """ + Regression: a deployment whose model_info sets use_local_token_counter=True must NOT + resolve a provider (remote) token counter, so callers fall back to the local tokenizer. + Without the flag, the Bedrock provider counter is still returned. + """ + from litellm.proxy.proxy_server import _get_provider_token_counter + + bedrock_params = {"model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0"} + resolved_model = "anthropic.claude-3-sonnet-20240229-v1:0" + + with_flag = { + "litellm_params": bedrock_params, + "model_info": {"use_local_token_counter": True}, + } + without_flag = { + "litellm_params": bedrock_params, + "model_info": {}, + } + + assert _get_provider_token_counter(with_flag, resolved_model) == (None, None, None) + + counter, _model, provider = _get_provider_token_counter(without_flag, resolved_model) + assert isinstance(counter, BedrockTokenCounter) + assert provider == "bedrock" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_local", [True, False]) +async def test_token_counter_use_local_flag_bypasses_bedrock_api(use_local): + """ + Regression for the per-model model_info.use_local_token_counter flag. + + call_endpoint=True normally routes Bedrock models to the remote CountTokens API. With + use_local_token_counter=True the proxy must skip the remote API entirely and count locally, + never instantiating/calling the Bedrock handler. Without it, the remote API is still used. + """ + import litellm.proxy.proxy_server as proxy_server + + model_info = {"use_local_token_counter": True} if use_local else {} + llm_router = Router( + model_list=[ + { + "model_name": "claude-bedrock", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet-20240229-v1:0" + }, + "model_info": model_info, + } + ] + ) + setattr(proxy_server, "llm_router", llm_router) + + with patch( + "litellm.llms.bedrock.count_tokens.bedrock_token_counter.BedrockCountTokensHandler" + ) as MockHandler: + mock_handler_instance = MockHandler.return_value + mock_handler_instance.handle_count_tokens_request = AsyncMock( + return_value={"input_tokens": 999} + ) + + response = await token_counter( + request=TokenCountRequest( + model="claude-bedrock", + messages=[ + {"role": "user", "content": "Hello Claude, how are you today?"} + ], + ), + call_endpoint=True, + ) + + if use_local: + mock_handler_instance.handle_count_tokens_request.assert_not_called() + assert response.tokenizer_type != "bedrock_api" + assert response.total_tokens > 0 + assert response.total_tokens != 999 + else: + mock_handler_instance.handle_count_tokens_request.assert_called_once() + assert response.tokenizer_type == "bedrock_api" + assert response.total_tokens == 999 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index ddf2040cd04..6fbaad211dc 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -33010,6 +33010,8 @@ export interface components { updated_at?: string | null; /** Updated By */ updated_by?: string | null; + /** Use Local Token Counter */ + use_local_token_counter?: boolean | null; } & { [key: string]: unknown; };