From c2c152f0607189e2ab95bdeaccb3e4a3612486e8 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Mon, 2 Nov 2026 16:40:50 -0800
Subject: [PATCH] refactor(count_tokens): share the tokenizer count call and
 keep estimate types exact

---
 litellm/proxy/proxy_server.py                 | 43 ++++++++++++-------
 ui/litellm-dashboard/src/lib/http/schema.d.ts |  7 +++
 2 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 59e9a93e5f2..439978e1b0b 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -15259,7 +15259,12 @@ def _provider_error_after_failed_estimate(
     return provider_failure.to_proxy_exception()
 
 
-async def _estimate_or_provider_error(estimate: Awaitable[int], provider_failure: ProviderCountFailure) -> int:
+_EstimateT = TypeVar("_EstimateT", int, int | None)
+
+
+async def _estimate_or_provider_error(
+    estimate: Awaitable[_EstimateT], provider_failure: ProviderCountFailure
+) -> _EstimateT:
     try:
         return await estimate
     except BaseException as e:
@@ -15413,6 +15418,22 @@ async def _tokenizer_or_provider_error(
     raise _provider_error_after_failed_estimate(provider_failure, tokenizer) from tokenizer
 
 
+def _tokenizer_count(
+    template_model: str,
+    prompt: str | None,
+    messages: Sequence[AllMessageValues] | None,
+    tools: list[ChatCompletionToolParam] | None,
+    tokenizer: SelectTokenizerResponse,
+) -> Awaitable[int]:
+    return offload_token_count(litellm.token_counter)(
+        model=template_model,
+        text=prompt,
+        messages=messages,
+        tools=tools,
+        custom_tokenizer=tokenizer,
+    )
+
+
 async def _local_token_estimate(
     model: str,
     template_model: str,
@@ -15427,13 +15448,7 @@ async def _local_token_estimate(
         if provider_failure is None
         else await _tokenizer_or_provider_error(model, custom_tokenizer, provider_failure)
     )
-    count: Final = offload_token_count(litellm.token_counter)(
-        model=template_model,
-        text=prompt,
-        messages=messages,
-        tools=tools,
-        custom_tokenizer=tokenizer,
-    )
+    count: Final = _tokenizer_count(template_model, prompt, messages, tools, tokenizer)
     total_tokens: Final = await (
         count if provider_failure is None else _estimate_or_provider_error(count, provider_failure)
     )
@@ -15492,10 +15507,9 @@ async def _kimi_k3_local_tokens(
     encoding: Final = await _KIMI_K3_ENCODINGS.encoding(TOKENIZER_LOAD_WAIT_SECONDS)
     if encoding is None:
         raise unavailable.to_proxy_exception()
-    counted: Final = await _estimate_or_provider_error(
-        count_messages(encoding) if prompt is None else offload_token_count(kimi_k3_text_tokens)(encoding, prompt),
-        unreadable,
-    )
+    if prompt is not None:
+        return await _estimate_or_provider_error(offload_token_count(kimi_k3_text_tokens)(encoding, prompt), unreadable)
+    counted: Final = await _estimate_or_provider_error(count_messages(encoding), unreadable)
     if counted is None:
         raise unreadable.to_proxy_exception()
     return counted
@@ -15521,10 +15535,7 @@ async def _local_tokenizer_tokens(
     if selected is None:
         raise unavailable.to_proxy_exception()
     total_tokens: Final = await _estimate_or_provider_error(
-        offload_token_count(litellm.token_counter)(
-            model=template_model, text=prompt, messages=messages, tools=tools, custom_tokenizer=selected
-        ),
-        unreadable,
+        _tokenizer_count(template_model, prompt, messages, tools, selected), unreadable
     )
     return _LocalTokenEstimate(total_tokens=total_tokens, tokenizer_type=selected["type"])
 
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 3aa7c888e98..ad41cc10a2a 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -47732,6 +47732,11 @@ export interface components {
              * @enum {string}
              */
             litellm_count_source?: "provider" | "token_count_model" | "local_estimate";
+            /**
+             * Litellm Count Tokenizer
+             * @enum {string}
+             */
+            litellm_count_tokenizer?: "openai" | "llama3" | "deepseek" | "cohere" | "kimi_k3";
             /** Prompttokensdetails */
             promptTokensDetails: components["schemas"]["PromptTokensDetails"][];
             /** Totaltokens */
@@ -47762,6 +47767,8 @@ export interface components {
         TokenCountResponse: {
             /** Count Source */
             count_source?: ("provider" | "token_count_model" | "local_estimate") | null;
+            /** Count Tokenizer */
+            count_tokenizer?: ("openai" | "llama3" | "deepseek" | "cohere" | "kimi_k3") | null;
             /**
              * Error
              * @default false