From 6a30c58a17757a80323a95cb08a150dc895e6b3b Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 17:21:11 -0700 Subject: [PATCH] fix(rust): separate public and request token counter policies --- litellm/rust_bridge/README.md | 8 +- litellm/rust_bridge/catalog.py | 5 ++ litellm/rust_bridge/configuration.py | 1 + litellm/rust_bridge/token_counter/__init__.py | 3 + .../rust_bridge/token_counter/definition.py | 1 + litellm/rust_bridge/token_counter/value.py | 10 +-- .../rust_bridge/test_configuration_env.py | 86 +++++++++++++++++++ 7 files changed, 107 insertions(+), 7 deletions(-) diff --git a/litellm/rust_bridge/README.md b/litellm/rust_bridge/README.md index dab88fa55c3..641cad6dfc1 100644 --- a/litellm/rust_bridge/README.md +++ b/litellm/rust_bridge/README.md @@ -6,7 +6,7 @@ Every SDK API has one `NativeComponent` in the immutable `COMPONENTS` catalog. A `RustImplementationState` records whether Rust is unimplemented, experimental, or ready. `RolloutPolicy` independently selects unsupported, Python-only, Rust opt-in, Rust opt-out, or Rust-required execution. Optional Rust execution can fall back to Python. Rust-required execution cannot -OCR completed delivery is ready and default-on. Messages, chat completions, token counting, and Responses WebSocket transport are experimental and opt-in. Other completed APIs remain Python-only. Bedrock transcription requires Rust because it has no Python implementation; Python-backed transcription providers remain on Python +OCR completed delivery is ready and default-on. Messages, chat completions, raw-request input token counting, and Responses WebSocket transport are experimental and opt-in. The public `litellm.token_counter()` and other completed APIs remain Python-only. Bedrock transcription requires Rust because it has no Python implementation; Python-backed transcription providers remain on Python ```python execution = COMPONENT.resolve( @@ -28,7 +28,11 @@ Provider failures, host callback failures, cancellation, conversion failures, an `invoke` and `ainvoke` return the native result or execute the supplied fallback directly. There is no public admission, prepare, accepts, or can-handle API -Token counting follows the same component policy. Its one native counting entrypoint validates the tokenizer configuration and request body, obtains and caches the required tokenizer resource, then counts. Unsupported inputs decline, known resource loading failures report native unavailability, and unexpected counting failures propagate +Token counting has two catalog entries. `UtilityName.TOKEN_COUNTER` declares the public `litellm.token_counter()` as Python-only, with no native exports. The public function continues to execute Python directly regardless of `litellm.rust(bool)` or `LITELLM_RUST` + +`UtilityName.REQUEST_INPUT_TOKEN_COUNTER` owns the experimental raw-request optimization. Budget reservation keeps its direct `litellm.rust_bridge.token_counter.count_input_tokens` import. That adapter uses `REQUEST_COMPONENT`; `COMPONENT` describes the public API + +The request optimization follows its component policy. Its one native counting entrypoint validates the tokenizer configuration and request body, obtains and caches the required tokenizer resource, then counts. Unsupported inputs decline, known resource loading failures report native unavailability, and unexpected counting failures propagate ## Package layout diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 706d8e28b51..741bca200aa 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -136,6 +136,11 @@ COMPONENTS: Final = MappingProxyType( ), UtilityName.TOKEN_COUNTER: _component( UtilityName.TOKEN_COUNTER, + _unimplemented(), + (), + ), + UtilityName.REQUEST_INPUT_TOKEN_COUNTER: _component( + UtilityName.REQUEST_INPUT_TOKEN_COUNTER, _experimental_completed, ("count_input_tokens",), ), diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py index dafadef0bb0..d5a231842a8 100644 --- a/litellm/rust_bridge/configuration.py +++ b/litellm/rust_bridge/configuration.py @@ -25,6 +25,7 @@ class RouteName(str, Enum): class UtilityName(str, Enum): TOKEN_COUNTER = "token_counter" + REQUEST_INPUT_TOKEN_COUNTER = "request_input_token_counter" class RustImplementationState(str, Enum): diff --git a/litellm/rust_bridge/token_counter/__init__.py b/litellm/rust_bridge/token_counter/__init__.py index 1734a1799f7..acd6e86aef8 100644 --- a/litellm/rust_bridge/token_counter/__init__.py +++ b/litellm/rust_bridge/token_counter/__init__.py @@ -1,9 +1,12 @@ from typing import Final +from litellm.rust_bridge.token_counter.definition import COMPONENT, REQUEST_COMPONENT from litellm.rust_bridge.token_counter.types import InputTokenCount, RustTokenizer from litellm.rust_bridge.token_counter.value import TOKEN_COUNTER, count_input_tokens, rust_tokenizer __all__: Final = ( + "COMPONENT", + "REQUEST_COMPONENT", "TOKEN_COUNTER", "InputTokenCount", "RustTokenizer", diff --git a/litellm/rust_bridge/token_counter/definition.py b/litellm/rust_bridge/token_counter/definition.py index 8aaadc1c316..7197e6ec3bb 100644 --- a/litellm/rust_bridge/token_counter/definition.py +++ b/litellm/rust_bridge/token_counter/definition.py @@ -4,3 +4,4 @@ from litellm.rust_bridge.catalog import COMPONENTS from litellm.rust_bridge.configuration import UtilityName COMPONENT: Final = COMPONENTS[UtilityName.TOKEN_COUNTER] +REQUEST_COMPONENT: Final = COMPONENTS[UtilityName.REQUEST_INPUT_TOKEN_COUNTER] diff --git a/litellm/rust_bridge/token_counter/value.py b/litellm/rust_bridge/token_counter/value.py index 68dc9203fa0..1f2af468804 100644 --- a/litellm/rust_bridge/token_counter/value.py +++ b/litellm/rust_bridge/token_counter/value.py @@ -8,7 +8,7 @@ import litellm 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.runtime import BridgeErrorContext, ainvoke -from litellm.rust_bridge.token_counter.definition import COMPONENT +from litellm.rust_bridge.token_counter.definition import REQUEST_COMPONENT from litellm.rust_bridge.token_counter.types import InputTokenCount, RustTokenCounter, RustTokenizer from litellm.utils import claude_json_str, huggingface_tokenizer_kind @@ -19,11 +19,11 @@ def _as_counter(value: object) -> RustTokenCounter | None: return cast(RustTokenCounter, value) if callable(value) else None -TOKEN_COUNTER: Final = COMPONENT.bind("count_input_tokens", validate=_as_counter) +TOKEN_COUNTER: Final = REQUEST_COMPONENT.bind("count_input_tokens", validate=_as_counter) def rust_tokenizer(model: str) -> RustTokenizer | None: - execution: Final = COMPONENT.resolve() + execution: Final = REQUEST_COMPONENT.resolve() if execution.select(TOKEN_COUNTER) is None: return None kind: Final = None if litellm.disable_hf_tokenizer_download is True else huggingface_tokenizer_kind(model) @@ -49,7 +49,7 @@ def _tokenizer_resource(tokenizer: str) -> str: async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None: - execution: Final = COMPONENT.resolve() + execution: Final = REQUEST_COMPONENT.resolve() counter: Final = execution.select(TOKEN_COUNTER) async def python_fallback() -> None: @@ -71,5 +71,5 @@ async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputToke else None, python_fallback=python_fallback, adapt=_INPUT_TOKEN_COUNT.validate_python, - context=BridgeErrorContext(route=COMPONENT.name.value, provider="", model=""), + context=BridgeErrorContext(route=REQUEST_COMPONENT.name.value, provider="", model=""), ) diff --git a/tests/test_litellm/rust_bridge/test_configuration_env.py b/tests/test_litellm/rust_bridge/test_configuration_env.py index 8db7fb6355d..fb8f67d5f9c 100644 --- a/tests/test_litellm/rust_bridge/test_configuration_env.py +++ b/tests/test_litellm/rust_bridge/test_configuration_env.py @@ -1,10 +1,21 @@ from __future__ import annotations +import json +from collections.abc import Callable +from types import ModuleType +from typing import Final + import pytest +import litellm +from litellm.litellm_core_utils import token_counter as python_counter +from litellm.rust_bridge import bindings, configuration +from litellm.rust_bridge import token_counter as bridge from litellm.rust_bridge.configuration import ( + ExecutionDecision, _parse_env_bool, # pyright: ignore[reportPrivateUsage] # directly test env parsing contract ) +from litellm.rust_bridge.token_counter import COMPONENT @pytest.mark.parametrize(("value", "expected"), (("1", True), ("0", False), (" 1 ", True), (" 0 ", False))) @@ -20,3 +31,78 @@ def test_parse_env_bool_preserves_unset_value() -> None: def test_parse_env_bool_rejects_unknown_value(value: str) -> None: with pytest.raises(ValueError, match="must be '1' or '0'"): _parse_env_bool(value) + + +@pytest.mark.parametrize("environment", (None, "0", "1")) +@pytest.mark.parametrize("override", (None, False, True)) +def test_public_token_counter_stays_python_only( + monkeypatch: pytest.MonkeyPatch, environment: str | None, override: bool | None +) -> None: + def unexpected_native_load() -> ModuleType: + raise AssertionError("public token counting must not load Rust") + + calls: Final[list[str]] = [] + + def python_count(text: str) -> int: + calls.append(text) + return len(text) + + configuration.reset_rust_configuration() + if environment is None: + monkeypatch.delenv("LITELLM_RUST", raising=False) + else: + monkeypatch.setenv("LITELLM_RUST", environment) + if override is not None: + litellm.rust(override) + monkeypatch.setattr(bindings, "get_native_bridge", unexpected_native_load) + monkeypatch.setattr(python_counter, "_get_count_function", lambda model, custom_tokenizer: python_count) + try: + assert COMPONENT.resolve().decision is ExecutionDecision.PYTHON + assert litellm.token_counter(model="gpt-4o", text="hello") == 5 + assert calls == ["hello"] + finally: + configuration.reset_rust_configuration() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("environment", (None, "0", "1")) +@pytest.mark.parametrize("override", (None, False, True)) +async def test_budget_direct_import_follows_request_rollout( + monkeypatch: pytest.MonkeyPatch, environment: str | None, override: bool | None +) -> None: + from litellm.proxy.spend_tracking.budget_reservation import count_request_input_tokens + + configuration.reset_rust_configuration() + if environment is None: + monkeypatch.delenv("LITELLM_RUST", raising=False) + else: + monkeypatch.setenv("LITELLM_RUST", environment) + if override is not None: + litellm.rust(override) + model: Final = "gpt-4o" + messages: Final = [{"role": "user", "content": "hello"}] + body: Final = json.dumps({"model": model, "messages": messages}).encode() + native_calls: Final[list[bytes]] = [] + + async def counter( + body: bytes, + kind: str | None, + encoding: str, + disabled: bool, + legacy_accounting: bool, + resource_loader: Callable[[str], str], + ) -> object: + native_calls.append(body) + return {"model": model, "input_tokens": 42} + + bridge.TOKEN_COUNTER.override(counter) + enabled: Final = override if override is not None else environment == "1" + try: + budget: Final = await count_request_input_tokens( + request_body=json.loads(body), route="/v1/messages", llm_router=None, raw_body=body + ) + assert budget == {model: 42 if enabled else litellm.token_counter(model=model, messages=messages)} + assert native_calls == ([body] if enabled else []) + finally: + bridge.TOKEN_COUNTER.reset() + configuration.reset_rust_configuration()