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

fix(count_tokens): keep transient provider status and catch nested documents

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo 2026-10-21 01:12:21 +00:00
parent 55bb1ae3f3
commit 026df91689
3 changed files with 184 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -107,7 +107,14 @@ def _is_object_sequence(value: object) -> TypeGuard[Sequence[object]]:
def _is_document_content_block(block: object) -> bool:
block_mapping: Final = as_str_mapping(block)
return block_mapping is not None and block_mapping.get("type") in {"document", "file"}
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:

8
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,
@ -14118,6 +14119,13 @@ async def count_request_tokens(
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,

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

@ -2678,6 +2678,50 @@ def test_unreliable_local_estimate_preserves_local_count(
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
@pytest.mark.parametrize(
("resolved_model", "messages", "tools", "call_endpoint", "reject"),
[
@ -2754,6 +2798,130 @@ async def test_count_request_tokens_rejects_only_known_unreliable_estimates(
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:
router: Final = Router(
model_list=[
{
"model_name": "count-target",
"litellm_params": {
"model": "bedrock/meta.llama4-maverick",
"api_key": "test-key",
"api_base": "https://count.test",
},
}
]
)
monkeypatch.setattr(proxy_server, "llm_router", router)
failure: Final = ProviderCountFailure(
status_code=status_code,
message=f"Provider count failed with {status_code}",
)
with patch.object(
proxy_server,
"_provider_token_count",
new=AsyncMock(return_value=failure),
):
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:
router: Final = Router(
model_list=[
{
"model_name": "count-target",
"litellm_params": {
"model": "bedrock/meta.llama4-maverick",
"api_key": "test-key",
"api_base": "https://count.test",
},
}
]
)
monkeypatch.setattr(proxy_server, "llm_router", router)
failure: Final = ProviderCountFailure(status_code=400, message="Runtime count rejected the document")
with patch.object(
proxy_server,
"_provider_token_count",
new=AsyncMock(return_value=failure),
):
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(