feat(rust): count Anthropic tokens in Rust from litellm.token_counter

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

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

@@ -16,10 +16,11 @@ use crate::errors::RustBridgeDeclined;
 use crate::execution::run_async;
 
 /// Counts the input tokens of a raw request body off the Python event loop with
-/// the GIL released. Python owns which requests get here and what to do with
-/// the count. At most one encode per core runs at a time; the rest wait in the
-/// async task, where a cancelled Python awaiter drops them before any blocking
-/// work is scheduled.
+/// the GIL released, or of one string synchronously on the calling thread.
+/// Python owns which requests get here and what to do with the count. At most
+/// one request encode per core runs at a time; the rest wait in the async task,
+/// where a cancelled Python awaiter drops them before any blocking work is
+/// scheduled.
 #[pyclass(frozen)]
 struct TokenCounter {
     inner: Arc<CoreTokenCounter>,
@@ -43,6 +44,10 @@ impl TokenCounter {
         Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file))
     }
 
+    fn count_text(&self, py: Python<'_>, text: &str) -> PyResult<usize> {
+        release_gil(py, || self.inner.count_text(text)).map_err(token_count_error_to_pyerr)
+    }
+
     fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult<Bound<'py, PyAny>> {
         let counter = Arc::clone(&self.inner);
         let encode_slots = Arc::clone(&self.encode_slots);

litellm/litellm_core_utils/token_counter.py

@@ -625,7 +625,10 @@ def _get_exact_count_function(
             def count_tokens(text: str) -> int:
                 return len(tokenizer.encode_batch_fast([text])[0])
 
-            return count_tokens
+            rust_count: Final = (
+                None if custom_tokenizer is not None or model is None else _rust_anthropic_count_function(model)
+            )
+            return count_tokens if rust_count is None else _with_python_fallback(rust_count, count_tokens)
         elif tokenizer_json["type"] == "openai_tokenizer":
             encoding: Final = openai_tokenizer_encoding(model)
 
@@ -643,6 +646,27 @@ def _get_exact_count_function(
         return _get_tiktoken_count_function(encode_length)
 
 
+def _rust_anthropic_count_function(model: str) -> TokenCounterFunction | None:
+    """The Rust port of the Anthropic tokenizer when the bridge is enabled; the other HuggingFace tokenizers stay in Python."""
+    from litellm.rust_bridge.token_counter import text_counter
+    from litellm.utils import huggingface_tokenizer_kind
+
+    if huggingface_tokenizer_kind(model) != "anthropic":
+        return None
+    return text_counter("anthropic")
+
+
+def _with_python_fallback(rust_count: TokenCounterFunction, python_count: TokenCounterFunction) -> TokenCounterFunction:
+    def count_tokens(text: str) -> int:
+        try:
+            return rust_count(text)
+        except RuntimeError as error:
+            verbose_logger.debug("Rust token counter failed, counting in Python: %s", error)
+            return python_count(text)
+
+    return count_tokens
+
+
 def openai_tokenizer_encoding(model: str) -> tiktoken.Encoding:
     """The tiktoken encoding `token_counter` uses for a model on the `openai_tokenizer` path."""
     from litellm.utils import print_verbose

litellm/rust_bridge/token_counter.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Awaitable
+from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from functools import lru_cache
 from typing import Final, Literal, Protocol, cast  # noqa: TID251  # native extension exposes untyped callables
@@ -25,6 +25,9 @@ class RustTokenCounter(Protocol):
     def acount_request(self, body: bytes) -> Awaitable[object]:
         raise NotImplementedError
 
+    def count_text(self, text: str) -> int:
+        raise NotImplementedError
+
 
 class RustTokenCounterFactory(Protocol):
     def __call__(self, tokenizer_json: str) -> RustTokenCounter:
@@ -93,6 +96,16 @@ def _counter(factory: RustTokenCounterFactory, tokenizer: RustTokenizer) -> Rust
             return factory.from_o200k_ranks(o200k_base_rank_file())
 
 
+def text_counter(tokenizer: RustTokenizer) -> Callable[[str], int] | None:
+    """The Rust per-string counter for `tokenizer`, `None` when the bridge is off or the extension is missing."""
+    if not rust_enabled():
+        return None
+    factory: Final = TOKEN_COUNTER.load()
+    if factory is None:
+        return None
+    return _counter(factory, tokenizer).count_text
+
+
 async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None:
     if not rust_enabled():
         return None

tests/test_litellm/litellm_core_utils/test_token_counter.py

@@ -1576,3 +1576,149 @@ def test_high_detail_image_token_upper_bound_covers_every_image_size(width: int,
 def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_image() -> None:
     assert calculate_img_tokens(_png_data_url(2000, 768), mode="high") == high_detail_image_token_upper_bound()
     assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound()
+
+
+ANTHROPIC_MODEL: Final = "claude-sonnet-4-5-20250929"
+RUST_TEXTS: Final = (
+    "Hello, how are you today?",
+    "I'VE got 1234567 things; it's \"fine\"...\r\n\r\n  caf\u00e9 \u0645\u0631\u062d\u0628\u0627 \U0001f600 <|endoftext|>",
+    "x " * 5_000,
+    "",
+)
+
+
+class _FakeTextCounter:
+    def __init__(self, count: int | Exception) -> None:
+        self.count = count
+        self.texts: list[str] = []
+
+    def count_text(self, text: str) -> int:
+        self.texts.append(text)
+        if isinstance(self.count, Exception):
+            raise self.count
+        return self.count
+
+
+class _FakeTextCounterFactory:
+    def __init__(self, count: int | Exception) -> None:
+        self.count = count
+        self.counters: list[_FakeTextCounter] = []
+
+    def __call__(self, tokenizer_json: str) -> _FakeTextCounter:
+        counter = _FakeTextCounter(self.count)
+        self.counters.append(counter)
+        return counter
+
+    def from_cl100k_ranks(self, rank_file: str) -> _FakeTextCounter:
+        return self(rank_file)
+
+    def from_o200k_ranks(self, rank_file: str) -> _FakeTextCounter:
+        return self(rank_file)
+
+
+@pytest.fixture
+def rust_bridge(monkeypatch: pytest.MonkeyPatch):
+    from litellm.rust_bridge import bindings, configuration
+    from litellm.rust_bridge import token_counter as bridge
+
+    bridge.TOKEN_COUNTER.reset()
+    bridge._counter.cache_clear()
+    configuration.reset_rust_configuration()
+    monkeypatch.setattr(bindings, "get_native_bridge", lambda: object())
+    yield bridge
+    bridge.TOKEN_COUNTER.reset()
+    bridge._counter.cache_clear()
+    configuration.reset_rust_configuration()
+
+
+def test_anthropic_text_is_counted_by_rust_when_the_bridge_is_enabled(rust_bridge) -> None:
+    factory: Final = _FakeTextCounterFactory(1_000)
+    litellm.rust(True)
+    rust_bridge.TOKEN_COUNTER.override(factory)
+
+    assert token_counter_new(model=ANTHROPIC_MODEL, text="hello") == 1_000
+    assert token_counter_new(model=ANTHROPIC_MODEL, text="hello again") == 1_000
+    assert len(factory.counters) == 1
+    assert factory.counters[0].texts == ["hello", "hello again"]
+
+
+def test_anthropic_messages_are_counted_by_rust_when_the_bridge_is_enabled(rust_bridge) -> None:
+    factory: Final = _FakeTextCounterFactory(1_000)
+    litellm.rust(True)
+    rust_bridge.TOKEN_COUNTER.override(factory)
+
+    count: Final = token_counter_new(model=ANTHROPIC_MODEL, messages=[{"role": "user", "content": "hello"}])
+
+    assert count >= 2_000
+    assert "hello" in factory.counters[0].texts
+
+
+@pytest.mark.parametrize("model", ("gpt-4", "gpt-4o", "replicate/meta/llama-2-70b-chat"))
+def test_only_the_anthropic_tokenizer_is_routed_to_rust(rust_bridge, model: str) -> None:
+    factory: Final = _FakeTextCounterFactory(1_000)
+    litellm.rust(True)
+    rust_bridge.TOKEN_COUNTER.override(factory)
+
+    assert token_counter_new(model=model, text="hello") < 1_000
+    assert factory.counters == []
+
+
+def test_custom_huggingface_tokenizer_stays_in_python_for_anthropic_models(rust_bridge) -> None:
+    from tokenizers import Tokenizer
+
+    from litellm.utils import claude_json_str
+
+    factory: Final = _FakeTextCounterFactory(1_000)
+    litellm.rust(True)
+    rust_bridge.TOKEN_COUNTER.override(factory)
+    custom: Final = {"type": "huggingface_tokenizer", "tokenizer": Tokenizer.from_str(claude_json_str)}
+
+    assert token_counter_new(model=ANTHROPIC_MODEL, custom_tokenizer=custom, text="hello") < 1_000
+    assert factory.counters == []
+
+
+def test_disabled_bridge_counts_anthropic_text_in_python(rust_bridge) -> None:
+    factory: Final = _FakeTextCounterFactory(1_000)
+    litellm.rust(False)
+    rust_bridge.TOKEN_COUNTER.override(factory)
+
+    assert token_counter_new(model=ANTHROPIC_MODEL, text="hello") < 1_000
+    assert factory.counters == []
+
+
+def test_missing_native_module_counts_anthropic_text_in_python(rust_bridge, monkeypatch: pytest.MonkeyPatch) -> None:
+    from litellm.rust_bridge import bindings
+
+    litellm.rust(True)
+    monkeypatch.setattr(bindings, "get_native_bridge", lambda: None)
+    litellm.rust(False)
+    python_count: Final = token_counter_new(model=ANTHROPIC_MODEL, text="hello")
+    litellm.rust(True)
+
+    assert token_counter_new(model=ANTHROPIC_MODEL, text="hello") == python_count
+
+
+def test_rust_encode_failure_falls_back_to_python_per_string(rust_bridge) -> None:
+    litellm.rust(False)
+    python_count: Final = token_counter_new(model=ANTHROPIC_MODEL, text="hello")
+    litellm.rust(True)
+    rust_bridge.TOKEN_COUNTER.override(_FakeTextCounterFactory(RuntimeError("encode failed")))
+
+    assert token_counter_new(model=ANTHROPIC_MODEL, text="hello") == python_count
+
+
+@pytest.mark.parametrize("text", RUST_TEXTS)
+def test_native_anthropic_text_count_matches_python(rust_bridge, monkeypatch: pytest.MonkeyPatch, text: str) -> None:
+    from litellm.rust_bridge import bindings
+
+    native: Final = pytest.importorskip("litellm.rust_bridge._native")
+    monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
+    messages: Final = [{"role": "system", "content": "You are terse."}, {"role": "user", "name": "bob", "content": text}]
+
+    litellm.rust(False)
+    python_text: Final = token_counter_new(model=ANTHROPIC_MODEL, text=text)
+    python_messages: Final = token_counter_new(model=ANTHROPIC_MODEL, messages=messages)
+    litellm.rust(True)
+
+    assert token_counter_new(model=ANTHROPIC_MODEL, text=text) == python_text
+    assert token_counter_new(model=ANTHROPIC_MODEL, messages=messages) == python_messages

tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py

@@ -264,6 +264,9 @@ class _DecliningCounter:
     async def acount_request(self, body: bytes) -> object:
         raise _FakeDeclined("unsupported content block")
 
+    def count_text(self, text: str) -> int:
+        return len(text)
+
 
 class _DecliningFactory:
     def __call__(self, tokenizer_json: str) -> _DecliningCounter: