From ddcd22099bb7d670d1caddb7e95d790df68aa40c Mon Sep 17 00:00:00 2001 From: mateo Date: Mon, 14 Sep 2026 20:47:10 +0000 Subject: [PATCH] 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> --- .../crates/python-bridge/src/token_counter.rs | 13 +- litellm/litellm_core_utils/token_counter.py | 26 +++- litellm/rust_bridge/token_counter.py | 15 +- .../litellm_core_utils/test_token_counter.py | 146 ++++++++++++++++++ .../spend_tracking/test_budget_reservation.py | 3 + 5 files changed, 197 insertions(+), 6 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/token_counter.rs b/litellm-rust/crates/python-bridge/src/token_counter.rs index b4de50c5f1a..e18ac1f7d3e 100644 --- a/litellm-rust/crates/python-bridge/src/token_counter.rs +++ b/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, @@ -43,6 +44,10 @@ impl TokenCounter { Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file)) } + fn count_text(&self, py: Python<'_>, text: &str) -> PyResult { + 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> { let counter = Arc::clone(&self.inner); let encode_slots = Arc::clone(&self.encode_slots); diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 3b128899f45..ac9c6097254 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/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 diff --git a/litellm/rust_bridge/token_counter.py b/litellm/rust_bridge/token_counter.py index d36234f56c1..24d48e4508e 100644 --- a/litellm/rust_bridge/token_counter.py +++ b/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 diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 60f25c48443..c3f572eb0b0 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/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 diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 3e0acf917aa..3dea40f6bf0 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/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: