refactor(rust): keep the Python token counter on the fast backend

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm-rust/crates/python-bridge/Cargo.toml

@@ -10,7 +10,7 @@ name = "_native"
 crate-type = ["cdylib"]
 
 [features]
-default = ["abi3", "huggingface", "tiktoken"]
+default = ["abi3", "fast"]
 abi3 = ["pyo3/abi3-py310"]
 extension-module = ["pyo3/extension-module"]
 panic-test = []

litellm-rust/crates/python-bridge/src/token_counter.rs

@@ -33,28 +33,47 @@ pub(crate) struct TokenCounter {
 impl TokenCounter {
     #[new]
     fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
-        #[cfg(feature = "huggingface")]
+        #[cfg(feature = "fast")]
+        {
+            Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json))
+        }
+        #[cfg(all(not(feature = "fast"), feature = "huggingface"))]
         {
             Self::load(py, || CoreTokenCounter::from_json(tokenizer_json))
         }
-        #[cfg(not(feature = "huggingface"))]
+        #[cfg(not(any(feature = "fast", feature = "huggingface")))]
         {
             let _ = (py, tokenizer_json);
             Err(RustBridgeDeclined::new_err(
-                "tokenizer backend requires the huggingface feature",
+                "tokenizer backend requires the fast or huggingface feature",
             ))
         }
     }
 
     #[staticmethod]
-    fn from_json_fast(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
+    fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
         #[cfg(feature = "fast")]
         {
-            Self::load(py, || CoreTokenCounter::from_json_fast(tokenizer_json))
+            Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file))
         }
         #[cfg(not(feature = "fast"))]
         {
-            let _ = (py, tokenizer_json);
+            let _ = (py, rank_file);
+            Err(RustBridgeDeclined::new_err(
+                "tokenizer backend requires the fast feature",
+            ))
+        }
+    }
+
+    #[staticmethod]
+    fn from_o200k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
+        #[cfg(feature = "fast")]
+        {
+            Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file))
+        }
+        #[cfg(not(feature = "fast"))]
+        {
+            let _ = (py, rank_file);
             Err(RustBridgeDeclined::new_err(
                 "tokenizer backend requires the fast feature",
             ))

litellm-rust/crates/token-counter/README.md

@@ -8,7 +8,7 @@ The `huggingface` feature provides `huggingface::HuggingFaceTokenizer` through t
 
 The `tiktoken` feature provides `tiktoken::TiktokenTokenizer` through `tiktoken-rs`. Select an encoding with `TokenCounter::from_tiktoken`. The supported names are `cl100k_base`, `o200k_base`, `o200k_harmony`, `p50k_base`, `p50k_edit`, `r50k_base`, and `gpt2`
 
-The Hugging Face and tiktoken backends are enabled by default. The hand-written fast backend is opt-in. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend
+The Hugging Face and tiktoken backends are enabled by default. The hand-written fast backend is opt-in. The Python extension builds with `fast` only, which keeps the wheel at the size it had before the split. With `default-features = false`, callers can supply their own `Tokenizer` to `TokenCounter::new` without compiling a built-in backend
 
 Budget checks, cost calculation, and the `max_tokens` adjustment policy belong to `litellm-core-utils`. The counter does not own prices, budgets, or request limits
 

litellm/rust_bridge/_native.pyi

@@ -97,7 +97,9 @@ class ResponsesWebSocketConnection:
 class TokenCounter:
     def __new__(cls, tokenizer_json: str) -> TokenCounter: ...
     @staticmethod
-    def from_json_fast(tokenizer_json: str) -> TokenCounter: ...
+    def from_cl100k_ranks(rank_file: str) -> TokenCounter: ...
+    @staticmethod
+    def from_o200k_ranks(rank_file: str) -> TokenCounter: ...
     @staticmethod
     def from_tiktoken(encoding: str) -> TokenCounter: ...
     def acount_request(self, body: bytes) -> Future[dict[str, object]]: ...

litellm/rust_bridge/token_counter.py

@@ -11,21 +11,14 @@ from pydantic import TypeAdapter
 
 import litellm
 from litellm._logging import verbose_logger
+from litellm.litellm_core_utils.default_encoding import cl100k_base_rank_file, o200k_base_rank_file
 from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding, uses_legacy_message_accounting
 from litellm.rust_bridge.bindings import NativeBinding
 from litellm.rust_bridge.configuration import rust_enabled
 from litellm.rust_bridge.runtime import BridgeErrorContext, RustHandled, aattempt
 from litellm.utils import claude_json_str, huggingface_tokenizer_kind
 
-RustTokenizer = Literal[
-    "anthropic",
-    "cl100k_base",
-    "o200k_base",
-    "o200k_harmony",
-    "p50k_base",
-    "p50k_edit",
-    "r50k_base",
-]
+RustTokenizer = Literal["anthropic", "cl100k_base", "o200k_base"]
 
 
 class RustTokenCounter(Protocol):
@@ -37,7 +30,10 @@ class RustTokenCounterFactory(Protocol):
     def __call__(self, tokenizer_json: str) -> RustTokenCounter:
         raise NotImplementedError
 
-    def from_tiktoken(self, encoding: RustTokenizer) -> RustTokenCounter:
+    def from_cl100k_ranks(self, rank_file: str) -> RustTokenCounter:
+        raise NotImplementedError
+
+    def from_o200k_ranks(self, rank_file: str) -> RustTokenCounter:
         raise NotImplementedError
 
 
@@ -67,8 +63,9 @@ def rust_tokenizer(model: str) -> RustTokenizer | None:
     """The Rust counter for the tokenizer `litellm.token_counter` selects for `model`, `None` when Python must count.
 
     Mirrors `_select_tokenizer_helper`: the Anthropic tokenizer has a Rust port, the other HuggingFace
-    downloads do not. All tiktoken encodings used by Python are backed by tiktoken-rs. Rust prices every
-    message with the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in Python."""
+    downloads do not, and of the tiktoken encodings `cl100k_base` and `o200k_base` do (p50k/r50k do not). Rust
+    prices every message with the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in
+    Python."""
     if litellm.disable_token_counter is True:
         return None
     kind: Final = None if litellm.disable_hf_tokenizer_download is True else huggingface_tokenizer_kind(model)
@@ -76,26 +73,24 @@ def rust_tokenizer(model: str) -> RustTokenizer | None:
         return "anthropic"
     if kind is not None or uses_legacy_message_accounting(model):
         return None
-    encoding: Final = openai_tokenizer_encoding(model).name
-    if encoding in (
-        "cl100k_base",
-        "o200k_base",
-        "o200k_harmony",
-        "p50k_base",
-        "p50k_edit",
-        "r50k_base",
-    ):
-        return cast(RustTokenizer, encoding)
-    return None
+    match openai_tokenizer_encoding(model).name:
+        case "cl100k_base":
+            return "cl100k_base"
+        case "o200k_base":
+            return "o200k_base"
+        case _:
+            return None
 
 
-@lru_cache(maxsize=8)
+@lru_cache(maxsize=4)
 def _counter(factory: RustTokenCounterFactory, tokenizer: RustTokenizer) -> RustTokenCounter:
     match tokenizer:
         case "anthropic":
             return factory(claude_json_str)
-        case _:
-            return factory.from_tiktoken(tokenizer)
+        case "cl100k_base":
+            return factory.from_cl100k_ranks(cl100k_base_rank_file())
+        case "o200k_base":
+            return factory.from_o200k_ranks(o200k_base_rank_file())
 
 
 async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None:

tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py

@@ -256,7 +256,7 @@ class _RecordingCounter:
 
 
 class _RecordingFactory:
-    """Stands in for the native `TokenCounter` class."""
+    """Stands in for the native `TokenCounter` class: called with tokenizer JSON, or `from_*_ranks`."""
 
     def __init__(self) -> None:
         self.calls: list[tuple[rust_token_counter.RustTokenizer, bytes]] = []
@@ -264,8 +264,11 @@ class _RecordingFactory:
     def __call__(self, tokenizer_json: str) -> _RecordingCounter:
         return _RecordingCounter(self, "anthropic")
 
-    def from_tiktoken(self, encoding: rust_token_counter.RustTokenizer) -> _RecordingCounter:
-        return _RecordingCounter(self, encoding)
+    def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter:
+        return _RecordingCounter(self, "cl100k_base")
+
+    def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter:
+        return _RecordingCounter(self, "o200k_base")
 
 
 class _DecliningCounter:
@@ -277,7 +280,10 @@ class _DecliningFactory:
     def __call__(self, tokenizer_json: str) -> _DecliningCounter:
         return _DecliningCounter()
 
-    def from_tiktoken(self, encoding: rust_token_counter.RustTokenizer) -> _DecliningCounter:
+    def from_cl100k_ranks(self, rank_file: str) -> _DecliningCounter:
+        return _DecliningCounter()
+
+    def from_o200k_ranks(self, rank_file: str) -> _DecliningCounter:
         return _DecliningCounter()
 
 

tests/test_litellm/rust_bridge/test_token_counter.py

@@ -8,6 +8,7 @@ cases need the extension and are skipped when it is not built.
 from __future__ import annotations
 
 import json
+from types import MappingProxyType
 from typing import Final
 
 import pytest
@@ -26,6 +27,7 @@ MODEL: Final = "claude-sonnet-4-5-20250929"
 CL100K_MODEL: Final = "gpt-4"
 O200K_MODEL: Final = "gpt-4o"
 TOKENIZERS: Final[tuple[bridge.RustTokenizer, ...]] = ("anthropic", "cl100k_base", "o200k_base")
+RANK_FILE_LINES: Final = MappingProxyType({"cl100k_base": 100_256, "o200k_base": 199_998})
 BODY: Final = json.dumps({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}).encode()
 
 
@@ -53,20 +55,24 @@ class _RecordingCounter:
 
 
 class _RecordingFactory:
-    """Stands in for the native `TokenCounter` class."""
+    """Stands in for the native `TokenCounter` class: callable for tokenizer JSON, `from_*_ranks` for rank files."""
 
     def __init__(self) -> None:
         self.counters: list[_RecordingCounter] = []
-        self.encodings: list[str] = []
+        self.rank_files: list[str] = []
 
     def __call__(self, tokenizer_json: str) -> _RecordingCounter:
         counter = _RecordingCounter(tokenizer_json)
         self.counters.append(counter)
         return counter
 
-    def from_tiktoken(self, encoding: bridge.RustTokenizer) -> _RecordingCounter:
-        self.encodings.append(encoding)
-        return self(encoding)
+    def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter:
+        self.rank_files.append(rank_file)
+        return self("cl100k_base")
+
+    def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter:
+        self.rank_files.append(rank_file)
+        return self("o200k_base")
 
 
 class _RaisingCounter:
@@ -86,7 +92,10 @@ class _RaisingFactory:
     def __call__(self, tokenizer_json: str) -> _RaisingCounter:
         return _RaisingCounter(self.error)
 
-    def from_tiktoken(self, encoding: bridge.RustTokenizer) -> _RaisingCounter:
+    def from_cl100k_ranks(self, rank_file: str) -> _RaisingCounter:
+        return _RaisingCounter(self.error)
+
+    def from_o200k_ranks(self, rank_file: str) -> _RaisingCounter:
         return _RaisingCounter(self.error)
 
 
@@ -131,7 +140,7 @@ async def test_enabled_bridge_returns_typed_count_and_reuses_one_counter() -> No
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("tokenizer", ("cl100k_base", "o200k_base"))
-async def test_tiktoken_counter_is_built_from_the_encoding_once(tokenizer: bridge.RustTokenizer) -> None:
+async def test_tiktoken_counter_is_built_from_the_vendored_rank_file_once(tokenizer: bridge.RustTokenizer) -> None:
     factory: Final = _RecordingFactory()
     litellm.rust(True)
     bridge.TOKEN_COUNTER.override(factory)
@@ -140,7 +149,9 @@ async def test_tiktoken_counter_is_built_from_the_encoding_once(tokenizer: bridg
     second: Final = await bridge.count_input_tokens(BODY, tokenizer)
 
     assert first == second == bridge.InputTokenCount(model=MODEL, input_tokens=42)
-    assert factory.encodings == [tokenizer]
+    assert len(factory.rank_files) == 1
+    assert factory.rank_files[0].startswith("IQ== 0\n")
+    assert factory.rank_files[0].count("\n") == RANK_FILE_LINES[tokenizer]
     assert factory.counters[0].tokenizer_json == tokenizer
     assert factory.counters[0].bodies == [BODY, BODY]
 
@@ -224,13 +235,13 @@ def test_rust_tokenizer_mirrors_python_tokenizer_selection(model: str, expected:
     ("model", "python_encoding"),
     (("text-davinci-003", "p50k_base"), ("gpt-oss-120b", "o200k_harmony")),
 )
-def test_rust_tokenizer_uses_every_tiktoken_encoding_supported_by_rust(
-    monkeypatch: pytest.MonkeyPatch, model: str, python_encoding: bridge.RustTokenizer
+def test_rust_tokenizer_declines_tiktoken_encodings_rust_does_not_have(
+    monkeypatch: pytest.MonkeyPatch, model: str, python_encoding: str
 ) -> None:
     monkeypatch.setattr(litellm, "open_ai_chat_completion_models", litellm.open_ai_chat_completion_models | {model})
 
     assert openai_tokenizer_encoding(model).name == python_encoding
-    assert bridge.rust_tokenizer(model) == python_encoding
+    assert bridge.rust_tokenizer(model) is None
 
 
 def test_rust_tokenizer_declines_the_cohere_tokenizer_download(monkeypatch: pytest.MonkeyPatch) -> None: