scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-28 06:23:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 94

Merge pull request #47705 from BerriAI/litellm_count_tokens_hard400_stable_1_104_x

Browse source 
feat(count_tokens)!: reject known-wrong local estimates with HTTP 400 on stable/1.104.x
This commit is contained in:
yuneng-jiang 2026-10-21 00:53:35 -07:00 • committed by GitHub
commit c1f73e4bf5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 591 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

72
litellm/proxy/common_utils/token_count_settings.py
View file

@ -1,7 +1,7 @@
from collections.abc import Callable, Mapping
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass, replace
from types import MappingProxyType
from typing import Final, Literal
from typing import Final, Literal, TypeGuard
from pydantic import BaseModel, ConfigDict, ValidationError
@ -77,6 +77,74 @@ class ProviderCountFailure:
)
@dataclass(frozen=True, slots=True)
class _UnreliableLocalEstimateRule:
provider: str
model_substring: str
trigger: Literal["tools", "no_tools"]
_UNRELIABLE_LOCAL_ESTIMATE_RULES: Final[tuple[_UnreliableLocalEstimateRule, ...]] = (
_UnreliableLocalEstimateRule("bedrock", "nova-2-lite", "tools"),
_UnreliableLocalEstimateRule("azure_ai", "deepseek-v4", "tools"),
_UnreliableLocalEstimateRule("azure_ai", "grok-4", "no_tools"),
)
def _content_blocks(message: object) -> tuple[object, ...]:
message_mapping: Final = as_str_mapping(message)
if message_mapping is None:
return ()
content: Final = message_mapping.get("content")
if not _is_object_sequence(content):
return ()
return tuple(content)
def _is_object_sequence(value: object) -> TypeGuard[Sequence[object]]:
return isinstance(value, Sequence) and not isinstance(value, (str, bytes))
def _is_document_content_block(block: object) -> bool:
block_mapping: Final = as_str_mapping(block)
if block_mapping is None:
return False
if block_mapping.get("type") in {"document", "file"}:
return True
nested_content: Final = block_mapping.get("content")
if not _is_object_sequence(nested_content):
return False
return any(_is_document_content_block(nested_block) for nested_block in nested_content)
def _has_document_content_block(messages: Sequence[object] | None) -> bool:
return any(
any(_is_document_content_block(block) for block in _content_blocks(message)) for message in messages or ()
)
def unreliable_local_estimate(
custom_llm_provider: str | None,
model: str,
messages: Sequence[object] | None,
tools: Sequence[object] | None,
) -> str | None:
if _has_document_content_block(messages):
return f"{model} returned no server-side token count and LiteLLM cannot count document content locally"
trigger: Final = "tools" if tools else "no_tools"
matched_rule: Final = next(
(
rule
for rule in _UNRELIABLE_LOCAL_ESTIMATE_RULES
if custom_llm_provider == rule.provider
and rule.model_substring in model.casefold()
and rule.trigger == trigger
),
None,
)
return None if matched_rule is None else f"Local token estimate for {model} is known to be off by more than 5%"
def cost_map_entry(model: str, custom_llm_provider: str | None = None) -> Mapping[str, object]:
try:
return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)

46
litellm/proxy/proxy_server.py
View file

@ -143,6 +143,7 @@ from litellm.proxy.common_utils.callback_utils import (
)
from litellm.proxy.common_utils.realtime_utils import _realtime_request_body
from litellm.proxy.common_utils.token_count_settings import (
CANNOT_COUNT_MODEL_STATUSES,
TOKENIZER_LOAD_MAX_RUNNING,
TOKENIZER_LOAD_REUSE_SECONDS,
TOKENIZER_LOAD_WAIT_SECONDS,
@ -150,6 +151,7 @@ from litellm.proxy.common_utils.token_count_settings import (
TokenCountSettings,
cost_map_entry,
resolve_token_count_settings,
unreliable_local_estimate,
)
from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog
from litellm.router_utils.access_windows import access_windows_config_error
@ -13997,6 +13999,23 @@ def _token_count_settings_model(deployment: _TokenCountDeployment, request_model
return model, _non_empty_str(params.get("custom_llm_provider"))
def _unreliable_estimate_provider_model(
deployment: _TokenCountDeployment, request_model: str
) -> tuple[str, str | None]:
settings_model: Final = _token_count_settings_model(deployment, request_model)
if settings_model[1] is not None:
return settings_model
params: Final = deployment.get("litellm_params", {})
deployment_model: Final = _non_empty_str(params.get("model"))
if deployment_model is None:
return settings_model[0], None
try:
resolved: Final = litellm.get_llm_provider(model=deployment_model)
except litellm.exceptions.BadRequestError:
return settings_model[0], None
return resolved[0], resolved[1]
_PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset(
{litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value}
)
@ -14116,6 +14135,33 @@ async def count_request_tokens(
if isinstance(ladder, TokenCountResponse):
return ladder
if call_endpoint is True:
estimate_provider_model: Final = (
_unreliable_estimate_provider_model(deployment_view, request.model)
if custom_llm_provider is None
else (model_to_use, custom_llm_provider)
)
unreliable_reason: Final = unreliable_local_estimate(
custom_llm_provider=estimate_provider_model[1],
model=estimate_provider_model[0],
messages=messages,
tools=tools,
)
if (
unreliable_reason is not None
and isinstance(ladder, ProviderCountFailure)
and not ladder.rejected_before_provider_call
and ladder.status_code not in CANNOT_COUNT_MODEL_STATUSES
):
raise ladder.to_proxy_exception()
if unreliable_reason is not None:
raise ProxyException(
message=unreliable_reason,
type="invalid_request_error",
param="model",
code=400,
)
# Check if token counter is disabled before fallback
if litellm.disable_token_counter is True:
raise ProxyException(

475
tests/unit/proxy/test_proxy_token_counter.py
View file

@ -58,6 +58,7 @@ from litellm.llms.bedrock.count_tokens.kimi_k3 import (
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler
from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import VertexAIPartnerModelsTokenCounter
from litellm.proxy import proxy_server
from litellm.proxy._types import ProxyException, TokenCountRequest
from litellm.proxy.anthropic_endpoints.endpoints import (
count_tokens as anthropic_count_tokens,
@ -67,6 +68,7 @@ from litellm.proxy.common_utils.token_count_settings import (
TokenCountModel,
TokenCountSettings,
cost_map_entry,
unreliable_local_estimate,
)
from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request
from litellm.proxy.proxy_server import (
@ -2637,3 +2639,476 @@ async def test_call_endpoint_prompt_keeps_the_local_count_on_a_counter_that_need
local.total_tokens,
local.tokenizer_type,
)
@pytest.mark.parametrize(
("provider", "model", "messages", "tools", "expected"),
[
("bedrock", "meta.llama4-maverick", [{"content": [{"type": "document"}]}], None, "document content"),
("bedrock", "us.amazon.nova-2-lite", [], [{"name": "search"}], "more than 5%"),
("azure_ai", "DeepSeek-V4-Pro", [], [{"name": "search"}], "more than 5%"),
("azure_ai", "grok-4-1-fast-reasoning", [], None, "more than 5%"),
],
)
def test_unreliable_local_estimate_returns_reason(
provider: str,
model: str,
messages: list[object],
tools: list[object] | None,
expected: str,
) -> None:
reason: Final = unreliable_local_estimate(provider, model, messages, tools)
assert reason is not None and expected in reason
@pytest.mark.parametrize(
("provider", "model", "messages", "tools"),
[
("bedrock", "meta.llama4-maverick", [{"content": "text"}], None),
("bedrock", "us.amazon.nova-2-lite", [], None),
("azure_ai", "deepseek-v4-pro", [], None),
("azure_ai", "grok-4-1-fast-reasoning", [], [{"name": "search"}]),
("azure", "gpt-4.1", [{"content": "text"}], None),
],
)
def test_unreliable_local_estimate_preserves_local_count(
provider: str, model: str, messages: list[object], tools: list[object] | None
) -> None:
assert unreliable_local_estimate(provider, model, messages, tools) is None
@pytest.mark.parametrize(
("nested_content", "expected"),
[
(
[
{
"type": "document",
"source": {
"type": "base64",
"media_type": "application/pdf",
"data": "JVBERi0xLjQK",
},
}
],
"meta.llama4-maverick returned no server-side token count and LiteLLM cannot count document content locally",
),
("text", None),
([{"type": "text", "text": "hello"}], None),
],
)
def test_unreliable_local_estimate_detects_nested_document_content(
nested_content: object, expected: str | None
) -> None:
reason: Final = unreliable_local_estimate(
custom_llm_provider="bedrock",
model="meta.llama4-maverick",
messages=[
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "t1",
"content": nested_content,
}
],
}
],
tools=None,
)
assert reason == expected
_BEDROCK_COUNT_TOKENS_URL: Final = r"^https://bedrock-runtime\.us-east-1\.amazonaws\.com/model/.+/count-tokens$"
@pytest.mark.parametrize(
("resolved_model", "messages", "tools", "call_endpoint", "reject"),
[
(
"bedrock/us.meta.llama4-maverick-17b-instruct-v1:0",
[{"role": "user", "content": [{"type": "document"}]}],
None,
True,
True,
),
(
"bedrock/us.meta.llama4-maverick-17b-instruct-v1:0",
[{"role": "user", "content": "text"}],
None,
True,
False,
),
(
"bedrock/us.amazon.nova-2-lite-v1:0",
[{"role": "user", "content": "text"}],
[{"name": "search"}],
True,
True,
),
(
"bedrock/us.amazon.nova-2-lite-v1:0",
[{"role": "user", "content": "text"}],
None,
True,
False,
),
("azure_ai/deepseek-v4-pro", [{"role": "user", "content": "text"}], [{"name": "search"}], True, True),
("azure_ai/deepseek-v4-pro", [{"role": "user", "content": "text"}], None, True, False),
("azure_ai/grok-4-1-fast-reasoning", [{"role": "user", "content": "text"}], None, True, True),
("azure_ai/grok-4-1-fast-reasoning", [{"role": "user", "content": "text"}], [{"name": "search"}], True, False),
(
"bedrock/us.meta.llama4-maverick-17b-instruct-v1:0",
[
{
"role": "user",
"content": [
{
"type": "document",
"source": {
"type": "base64",
"media_type": "application/pdf",
"data": "JVBERi0xLjQK",
},
}
],
}
],
None,
False,
False,
),
("azure/gpt-4.1", [{"role": "user", "content": "text"}], None, True, False),
],
)
@pytest.mark.asyncio
async def test_count_request_tokens_rejects_only_known_unreliable_estimates(
monkeypatch: pytest.MonkeyPatch,
resolved_model: str,
messages: list[dict[str, object]],
tools: list[dict[str, object]] | None,
call_endpoint: bool,
reject: bool,
) -> None:
is_bedrock: Final = resolved_model.startswith("bedrock/")
litellm_params: Final = (
{
"model": resolved_model,
"aws_access_key_id": "AKIATESTACCESSKEY",
"aws_secret_access_key": "test-secret",
"aws_region_name": "us-east-1",
}
if is_bedrock
else {"model": resolved_model, "api_key": "test-key", "api_base": "https://count.test"}
)
router: Final = Router(model_list=[{"model_name": "count-target", "litellm_params": litellm_params}])
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(litellm, "disable_token_counter", False)
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
should_count_with_provider: Final = is_bedrock and call_endpoint
with respx.mock(assert_all_called=should_count_with_provider) as upstream:
upstream.post(url__regex=_BEDROCK_COUNT_TOKENS_URL).mock(
return_value=httpx.Response(400, json={"message": "CountTokens rejected request"})
)
if reject:
with pytest.raises(ProxyException) as raised:
await count_request_tokens(
request=TokenCountRequest(model="count-target", messages=messages, tools=tools),
call_endpoint=call_endpoint,
messages_format="anthropic",
)
assert raised.value.code == "400"
else:
response: Final = await count_request_tokens(
request=TokenCountRequest(model="count-target", messages=messages, tools=tools),
call_endpoint=call_endpoint,
messages_format="anthropic",
)
assert response.count_source == "local_estimate"
@pytest.mark.parametrize(
("status_code", "expected_type", "expected_message"),
[
(429, "token_counting_error", "Provider count failed with 429"),
(500, "token_counting_error", "Provider count failed with 500"),
(400, "invalid_request_error", "cannot count document content locally"),
],
)
@pytest.mark.asyncio
async def test_unreliable_document_estimate_preserves_transient_provider_errors(
monkeypatch: pytest.MonkeyPatch,
status_code: int,
expected_type: str,
expected_message: str,
) -> None:
model: Final = "bedrock/us.meta.llama4-maverick-17b-instruct-v1:0"
router: Final = Router(
model_list=[
{
"model_name": "count-target",
"litellm_params": {
"model": model,
"aws_access_key_id": "AKIATESTACCESSKEY",
"aws_secret_access_key": "test-secret",
"aws_region_name": "us-east-1",
},
}
]
)
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(litellm, "disable_token_counter", False)
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
with respx.mock(assert_all_called=True) as upstream:
upstream.post(url__regex=_BEDROCK_COUNT_TOKENS_URL).mock(
return_value=httpx.Response(
status_code,
json={"message": f"Provider count failed with {status_code}"},
)
)
with pytest.raises(ProxyException) as raised:
await count_request_tokens(
request=TokenCountRequest(
model="count-target",
messages=[
{
"role": "user",
"content": [
{
"type": "document",
"source": {
"type": "base64",
"media_type": "application/pdf",
"data": "JVBERi0xLjQK",
},
}
],
}
],
),
call_endpoint=True,
messages_format="anthropic",
)
assert raised.value.code == str(status_code)
assert raised.value.type == expected_type
assert expected_message in raised.value.message
@pytest.mark.asyncio
async def test_count_request_tokens_rejects_nested_document_after_provider_count_failure(
monkeypatch: pytest.MonkeyPatch,
) -> None:
model: Final = "bedrock/us.meta.llama4-maverick-17b-instruct-v1:0"
router: Final = Router(
model_list=[
{
"model_name": "count-target",
"litellm_params": {
"model": model,
"aws_access_key_id": "AKIATESTACCESSKEY",
"aws_secret_access_key": "test-secret",
"aws_region_name": "us-east-1",
},
}
]
)
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(litellm, "disable_token_counter", False)
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
with respx.mock(assert_all_called=True) as upstream:
upstream.post(url__regex=_BEDROCK_COUNT_TOKENS_URL).mock(
return_value=httpx.Response(400, json={"message": "Runtime count rejected the document"})
)
with pytest.raises(ProxyException) as raised:
await count_request_tokens(
request=TokenCountRequest(
model="count-target",
messages=[
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "t1",
"content": [
{
"type": "document",
"source": {
"type": "base64",
"media_type": "application/pdf",
"data": "JVBERi0xLjQK",
},
}
],
}
],
}
],
),
call_endpoint=True,
messages_format="anthropic",
)
assert raised.value.code == "400"
assert raised.value.type == "invalid_request_error"
assert "cannot count document content locally" in raised.value.message
@pytest.mark.asyncio
async def test_azure_ai_estimate_guard_runs_without_a_provider_counter(monkeypatch: pytest.MonkeyPatch) -> None:
router: Final = Router(
model_list=[
{
"model_name": "azure-ai-count-target",
"litellm_params": {
"model": "azure_ai/deepseek-v4-pro",
"api_key": "test-key",
"api_base": "https://count.test",
},
}
]
)
monkeypatch.setattr(proxy_server, "llm_router", router)
with pytest.raises(ProxyException) as raised:
await count_request_tokens(
request=TokenCountRequest(
model="azure-ai-count-target",
messages=[{"role": "user", "content": "Summarize."}],
tools=[{"name": "search"}],
),
call_endpoint=True,
)
assert raised.value.code == "400"
assert raised.value.type == "invalid_request_error"
assert raised.value.param == "model"
assert "deepseek-v4-pro" in raised.value.message
@pytest.mark.asyncio
async def test_document_block_keeps_successful_provider_count(monkeypatch: pytest.MonkeyPatch) -> None:
model: Final = "bedrock/us.anthropic.claude-sonnet-4-6"
router: Final = Router(
model_list=[
{
"model_name": "count-target",
"litellm_params": {
"model": model,
"aws_access_key_id": "AKIATESTACCESSKEY",
"aws_secret_access_key": "test-secret",
"aws_region_name": "us-east-1",
},
}
]
)
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(litellm, "disable_token_counter", False)
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
with respx.mock(assert_all_called=True) as upstream:
upstream.post(url__regex=_BEDROCK_COUNT_TOKENS_URL).mock(
return_value=httpx.Response(200, json={"inputTokens": 17})
)
result: Final = await count_request_tokens(
request=TokenCountRequest(
model="count-target",
messages=[
{
"role": "user",
"content": [
{
"type": "document",
"source": {
"type": "base64",
"media_type": "application/pdf",
"data": "JVBERi0xLjQK",
},
}
],
}
],
),
call_endpoint=True,
messages_format="anthropic",
)
assert result.total_tokens == 17
assert result.count_source == "provider"
@pytest.mark.asyncio
async def test_unreliable_estimate_uses_deployment_model_during_cooldown(
monkeypatch: pytest.MonkeyPatch,
) -> None:
router: Final = Router(
model_list=[
{
"model_name": "nova-count-target",
"litellm_params": {
"model": "bedrock/us.amazon.nova-2-lite-v1:0",
"aws_access_key_id": "AKIATESTACCESSKEY",
"aws_secret_access_key": "test-secret",
"aws_region_name": "us-east-1",
},
"model_info": {"id": "nova-cooldown-deployment"},
}
]
)
router.cooldown_cache.add_deployment_to_cooldown(
model_id="nova-cooldown-deployment",
original_exception=Exception("rate limited"),
exception_status=429,
cooldown_time=60,
)
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(litellm, "disable_token_counter", False)
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
with respx.mock(assert_all_called=False) as upstream:
upstream.post(url__regex=_BEDROCK_COUNT_TOKENS_URL).mock(
return_value=httpx.Response(200, json={"inputTokens": 17})
)
with pytest.raises(ProxyException) as raised:
await count_request_tokens(
request=TokenCountRequest(
model="nova-count-target",
messages=[{"role": "user", "content": "Find a record."}],
tools=[{"name": "search"}],
),
call_endpoint=True,
messages_format="anthropic",
)
local: Final = await count_request_tokens(
request=TokenCountRequest(
model="nova-count-target",
messages=[{"role": "user", "content": "Find a record."}],
),
call_endpoint=True,
messages_format="anthropic",
)
assert raised.value.code == "400"
assert local.count_source == "local_estimate"