diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 67f662408c4..a76b069935f 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/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 = [] diff --git a/litellm-rust/crates/python-bridge/src/token_counter.rs b/litellm-rust/crates/python-bridge/src/token_counter.rs index 12749486679..244401e6696 100644 --- a/litellm-rust/crates/python-bridge/src/token_counter.rs +++ b/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 { - #[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 { + fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult { #[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 { + #[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", )) diff --git a/litellm-rust/crates/token-counter/README.md b/litellm-rust/crates/token-counter/README.md index effa5324822..3ad21a00697 100644 --- a/litellm-rust/crates/token-counter/README.md +++ b/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 diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 7d20fc5a02f..05a6df6d5af 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/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]]: ... diff --git a/litellm/rust_bridge/token_counter.py b/litellm/rust_bridge/token_counter.py index d232db19a4a..d36234f56c1 100644 --- a/litellm/rust_bridge/token_counter.py +++ b/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: 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 c785a7b657a..0e0025f5194 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/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() diff --git a/tests/test_litellm/rust_bridge/test_token_counter.py b/tests/test_litellm/rust_bridge/test_token_counter.py index b95e6a62e9c..71aa79cc4bb 100644 --- a/tests/test_litellm/rust_bridge/test_token_counter.py +++ b/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: