From 026df9168951ce9cf150cae58bf7733d859a6f7d Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 21 Oct 2026 01:12:21 +0000 Subject: [PATCH] fix(count_tokens): keep transient provider status and catch nested documents Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../common_utils/token_count_settings.py | 9 +- litellm/proxy/proxy_server.py | 8 + tests/unit/proxy/test_proxy_token_counter.py | 168 ++++++++++++++++++ 3 files changed, 184 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py index be3984bf803..e8048d404d6 100644 --- a/litellm/proxy/common_utils/token_count_settings.py +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -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: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 66deb8c69e3..58b72dc18ee 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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, diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 72f51256ac7..c5de167f8d2 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -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(