From 3e0c4f6e5dcffcc9255f5df8048c8d651bb8adcf Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 2 Nov 2026 17:08:50 -0800 Subject: [PATCH] refactor(count_tokens): return local tokenizer failures as values and raise them at one boundary --- litellm/proxy/proxy_server.py | 194 +++++++++++-------- tests/unit/proxy/test_proxy_token_counter.py | 50 +++++ 2 files changed, 162 insertions(+), 82 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 821ae5958dc..d3eed7ea26f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15485,15 +15485,13 @@ async def _kimi_k3_encoded_count( return await offload_token_count(kimi_k3_tokens)(encoding, model, messages, tools) -def _known_local_tokenizer(settings: TokenCountSettings, request_model: str) -> LocalTokenizer | None: - named: Final = settings.local_tokenizer - if isinstance(named, UnknownLocalTokenizer): - raise _rejected_view( - 400, - f"{request_model} sets token_count_tokenizer to {named.value!r}, " - f"which is not one of {', '.join(get_args(LocalTokenizer))}", - ) - return named +def _unknown_tokenizer_failure(named: UnknownLocalTokenizer, request_model: str) -> ProviderCountFailure: + return ProviderCountFailure( + status_code=400, + message=f"{request_model} sets token_count_tokenizer to {named.value!r}, " + f"which is not one of {', '.join(get_args(LocalTokenizer))}", + rejected_before_provider_call=True, + ) def _local_tokenizer_failure(tokenizer: str, model: str, reason: str) -> ProviderCountFailure: @@ -15504,30 +15502,34 @@ def _local_tokenizer_failure(tokenizer: str, model: str, reason: str) -> Provide ) +async def _count_or_failure( + count: Awaitable[int | None], unreadable: ProviderCountFailure +) -> int | ProviderCountFailure: + try: + counted: Final = await count + except BaseException as e: + if not _answers_provider_error(e): + raise + verbose_proxy_logger.warning("Local token count failed, answering 400 instead: %s", type(e).__name__) + return unreadable + return unreadable if counted is None else counted + + async def _kimi_k3_local_tokens( prompt: str | None, count_messages: Callable[[tiktoken.Encoding], Awaitable[int | None]], unreadable: ProviderCountFailure, unavailable: ProviderCountFailure, -) -> int: +) -> int | ProviderCountFailure: encoding: Final = await _KIMI_K3_ENCODINGS.encoding(TOKENIZER_LOAD_WAIT_SECONDS) if encoding is None: - raise unavailable.to_proxy_exception() - return await _estimate_or_provider_error( - offload_token_count(kimi_k3_text_tokens)(encoding, prompt) - if prompt is not None - else _counted_or_rejected(count_messages(encoding), unreadable), + return unavailable + return await _count_or_failure( + offload_token_count(kimi_k3_text_tokens)(encoding, prompt) if prompt is not None else count_messages(encoding), unreadable, ) -async def _counted_or_rejected(count: Awaitable[int | None], rejection: ProviderCountFailure) -> int: - counted: Final = await count - if counted is None: - raise rejection.to_proxy_exception() - return counted - - async def _local_tokenizer_tokens( tokenizer: LocalTokenizer, request_model: str, @@ -15536,57 +15538,64 @@ async def _local_tokenizer_tokens( messages: Sequence[AllMessageValues] | None, tools: Sequence[ChatCompletionToolParam] | None, kimi_k3_count: Callable[[tiktoken.Encoding], Awaitable[int | None]], -) -> _LocalTokenEstimate: +) -> _LocalTokenEstimate | ProviderCountFailure: unreadable: Final = _local_tokenizer_failure(tokenizer, request_model, "could not count this request") unavailable: Final = _local_tokenizer_failure(tokenizer, request_model, "could not be loaded") if tokenizer == "kimi_k3": - return _LocalTokenEstimate( - total_tokens=await _kimi_k3_local_tokens(prompt, kimi_k3_count, unreadable, unavailable), - tokenizer_type="huggingface_tokenizer", + kimi_k3: Final = await _kimi_k3_local_tokens(prompt, kimi_k3_count, unreadable, unavailable) + return ( + kimi_k3 + if isinstance(kimi_k3, ProviderCountFailure) + else _LocalTokenEstimate(total_tokens=kimi_k3, tokenizer_type="huggingface_tokenizer") ) selected: Final = await asyncify(litellm.utils.select_named_tokenizer)(tokenizer) if selected is None: - raise unavailable.to_proxy_exception() - total_tokens: Final = await _estimate_or_provider_error( + return unavailable + total_tokens: Final = await _count_or_failure( _tokenizer_count(template_model, prompt, messages, tools, selected), unreadable ) - return _LocalTokenEstimate(total_tokens=total_tokens, tokenizer_type=selected["type"]) + return ( + total_tokens + if isinstance(total_tokens, ProviderCountFailure) + else _LocalTokenEstimate(total_tokens=total_tokens, tokenizer_type=selected["type"]) + ) -async def _count_with_local_tokenizer( - tokenizer: LocalTokenizer, - request: TokenCountRequest, - messages_format: TokenCountMessagesFormat, - model: str, - template_model: str, - prompt: str | None, - messages: Sequence[AllMessageValues] | None, - tools: Sequence[ChatCompletionToolParam] | None, - bridged: ChatCountRequest | None, - settings: TokenCountSettings, - responses_api: bool, - contents_estimable: bool, - stored_turns_unseen: bool, -) -> TokenCountResponse: - """Count with the tokenizer the deployment names in token_count_tokenizer, never asking the provider. +@dataclass(frozen=True, slots=True) +class _LocalTokenizerRequest: + request: TokenCountRequest + messages_format: TokenCountMessagesFormat + model: str + template_model: str + prompt: str | None + messages: Sequence[AllMessageValues] | None + tools: Sequence[ChatCompletionToolParam] | None + bridged: ChatCountRequest | None + responses_api: bool + contents_estimable: bool + stored_turns_unseen: bool - A Responses request bridged to chat is counted as the chat body inference sends, previous_response_id turns - included.""" - if stored_turns_unseen: - raise _local_tokenizer_failure( + +async def _local_tokenizer_count( + tokenizer: LocalTokenizer, counted: _LocalTokenizerRequest, settings: TokenCountSettings +) -> TokenCountResponse | ProviderCountFailure: + request: Final = counted.request + if counted.stored_turns_unseen: + return _local_tokenizer_failure( tokenizer, request.model, "cannot see the earlier turns previous_response_id points to on the provider" - ).to_proxy_exception() - if not contents_estimable: - raise _local_tokenizer_failure(tokenizer, request.model, "cannot read these contents").to_proxy_exception() - counted_messages: Final = ( - messages + ) + if not counted.contents_estimable: + return _local_tokenizer_failure(tokenizer, request.model, "cannot read these contents") + bridged: Final = counted.bridged + messages: Final = ( + counted.messages if bridged is None else cast( # cast-ok: the chat body Responses inference sends, chat-shaped dicts token_counter normalizes Sequence[AllMessageValues], [dict(message) for message in bridged.messages] ) ) - counted_tools: Final = ( - tools + tools: Final = ( + counted.tools if bridged is None else cast( # cast-ok: the chat tools Responses inference sends Sequence[ChatCompletionToolParam] | None, [dict(tool) for tool in bridged.tools or ()] or None @@ -15595,36 +15604,55 @@ async def _count_with_local_tokenizer( estimate: Final = await _local_tokenizer_tokens( tokenizer, request.model, - template_model, - prompt, - counted_messages, - counted_tools, + counted.template_model, + counted.prompt, + messages, + tools, kimi_k3_count=partial( _kimi_k3_encoded_count, request=request, - messages_format=messages_format if bridged is None else "openai", - model=model, - messages=counted_messages or (), - tools=counted_tools, + messages_format=counted.messages_format if bridged is None else "openai", + model=counted.model, + messages=messages or (), + tools=tools, ), ) + if isinstance(estimate, ProviderCountFailure): + return estimate framing: Final = ( 0 - if counted_messages is None + if messages is None else settings.hidden_prompt_tokens( - carries_tools=sends_function_tools(counted_tools), responses_api=responses_api + carries_tools=sends_function_tools(tools), responses_api=counted.responses_api ) ) return TokenCountResponse( total_tokens=estimate.total_tokens + framing, request_model=request.model, - model_used=model, + model_used=counted.model, tokenizer_type=estimate.tokenizer_type, count_source="local_estimate", count_tokenizer=tokenizer, ) +async def _count_with_local_tokenizer( + named: LocalTokenizer | UnknownLocalTokenizer, counted: _LocalTokenizerRequest, settings: TokenCountSettings +) -> TokenCountResponse: + """Count with the tokenizer the deployment names in token_count_tokenizer, never asking the provider. + + A Responses request bridged to chat is counted as the chat body inference sends, previous_response_id turns + included.""" + outcome: Final = ( + _unknown_tokenizer_failure(named, counted.request.model) + if isinstance(named, UnknownLocalTokenizer) + else await _local_tokenizer_count(named, counted, settings) + ) + if isinstance(outcome, ProviderCountFailure): + raise outcome.to_proxy_exception() + return outcome + + class _TokenCountDeployment(TypedDict, total=False): litellm_params: ReadOnly[Mapping[str, object]] model_info: ReadOnly[Mapping[str, object]] @@ -16029,7 +16057,7 @@ async def _count_request_tokens( default=await asyncify(cost_map_entry)(*_token_count_settings_model(deployment_view, request.model)), deployment_region=partial(_deployment_aws_region, deployment_view), ) - local_tokenizer: Final = _known_local_tokenizer(settings, request.model) + local_tokenizer: Final = settings.local_tokenizer # Try provider-specific token counting first - only for non-direct requests (from provider endpoints) provider_counter: BaseTokenCounter | None = None @@ -16228,20 +16256,22 @@ async def _count_request_tokens( if local_tokenizer is not None: return await _count_with_local_tokenizer( local_tokenizer, - request - if counted_view is None - else request.model_copy(update={"messages": local_messages, "system": local_system}), - messages_format, - model_to_use, - template_model, - prompt=prompt if counted_messages is None else None, - messages=counted_messages, - tools=counted_tools, - bridged=chat_count, - settings=settings, - responses_api=native_responses, - contents_estimable=contents_estimable, - stored_turns_unseen=chat_count is None and bool(_responses_routing_kwargs(routed_request)), + _LocalTokenizerRequest( + request=request + if counted_view is None + else request.model_copy(update={"messages": local_messages, "system": local_system}), + messages_format=messages_format, + model=model_to_use, + template_model=template_model, + prompt=prompt if counted_messages is None else None, + messages=counted_messages, + tools=counted_tools, + bridged=chat_count, + responses_api=native_responses, + contents_estimable=contents_estimable, + stored_turns_unseen=chat_count is None and bool(_responses_routing_kwargs(routed_request)), + ), + settings, ) kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None exact: Final = ( diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 8c4f9b26302..66f10a5ee34 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -5488,6 +5488,56 @@ async def test_the_deployments_named_tokenizer_beats_the_tokenizer_the_cost_map_ } +@pytest.mark.asyncio +async def test_a_named_tokenizer_count_adds_the_framing_the_deployment_sets_for_requests_with_tools( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model_info: Final = {"token_count_tokenizer": "openai", "token_count_framing": 5, "token_count_tool_framing": 40} + + _, utils_reply, sent = await _bedrock_counts( + monkeypatch, _NOVA_PRO, model_info, {"messages": _MIXED_SCRIPT_TURNS, "tools": [_KIMI_WEATHER]} + ) + + assert sent == () + assert isinstance(utils_reply, TokenCountResponse), utils_reply + assert utils_reply.total_tokens == 40 + litellm.token_counter( + model=_NOVA_PRO, messages=_MIXED_SCRIPT_TURNS, tools=[_KIMI_WEATHER] + ) + + +@pytest.mark.asyncio +async def test_a_named_kimi_k3_tokenizer_counts_text_with_kimi_k3_and_answers_400_for_a_body_it_cannot_render( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + encoding: Final = _available_kimi_k3(monkeypatch, tmp_path) + image_turn: Final[dict[str, JsonValue]] = { + "role": "user", + "content": [{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0K"}}], + } + + counted, _, counted_sent = await _bedrock_counts( + monkeypatch, _NOVA_PRO, {"token_count_tokenizer": "kimi_k3"}, {"messages": _MIXED_SCRIPT_TURNS} + ) + rejected, _, rejected_sent = await _bedrock_counts( + monkeypatch, _NOVA_PRO, {"token_count_tokenizer": "kimi_k3"}, {"messages": [image_turn]} + ) + + assert counted == { + "input_tokens": kimi_k3_tokens(encoding, _NOVA_PRO, _MIXED_SCRIPT_TURNS, None), + "litellm_count_source": "local_estimate", + "litellm_count_tokenizer": "kimi_k3", + } + status, message = _rejection(rejected) + assert ( + status, + "kimi_k3 tokenizer set in token_count_tokenizer, which could not count this request" in message, + ) == ( + 400, + True, + ), message + assert counted_sent == rejected_sent == () + + @pytest.mark.asyncio @pytest.mark.parametrize( ("model_info", "reason"),