fix(rust): fall back to Python when the Anthropic Rust tokenizer fails to load

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo 2026-09-14 21:08:43 +00:00
parent c849e0b766
commit 174037f49d
3 changed files with 59 additions and 1 deletions

View file

@ -103,8 +103,13 @@ def text_counter(tokenizer: RustTokenizer) -> Callable[[str], int] | None:
factory: Final = TOKEN_COUNTER.load()
if factory is None:
return None
try:
counter: Final = _counter(factory, tokenizer)
except (RuntimeError, ValueError) as error:
verbose_logger.debug("Rust token counter (%s) failed to load, counting in Python: %s", tokenizer, error)
return None
verbose_logger.debug("Rust token counter (%s) counting text", tokenizer)
return _counter(factory, tokenizer).count_text
return counter.count_text
async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None:

View file

@ -1723,6 +1723,18 @@ def test_rust_encode_failure_falls_back_to_python_per_string(rust_bridge) -> Non
assert token_counter_new(model=ANTHROPIC_MODEL, text="hello") == python_count
def test_rust_tokenizer_load_failure_falls_back_to_python(rust_bridge) -> None:
def failing_factory(tokenizer_json: str) -> _FakeTextCounter:
raise ValueError("tokenizer json rejected")
litellm.rust(False)
python_count: Final = token_counter_new(model=ANTHROPIC_MODEL, text="hello")
litellm.rust(True)
rust_bridge.TOKEN_COUNTER.override(failing_factory)
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

View file

@ -0,0 +1,41 @@
from typing import Final
import pytest
import litellm
from litellm.litellm_core_utils.token_counter import token_counter
from litellm.rust_bridge import bindings
from litellm.rust_bridge import token_counter as bridge
pytestmark = pytest.mark.requires_rust_extension
ANTHROPIC_MODEL: Final = "claude-sonnet-4-5-20250929"
TEXTS: Final = (
"Hello, how are you today?",
"I'VE got 1234567 things; it's \"fine\"...\r\n\r\n café مرحبا 😀 <|endoftext|>",
"x " * 5_000,
"",
)
@pytest.fixture
def native_bridge() -> None:
bridge.TOKEN_COUNTER.reset()
bridge._counter.cache_clear() # pyright: ignore[reportPrivateUsage] # the counter cache is keyed on the factory object
assert bindings.get_native_bridge() is not None
@pytest.mark.parametrize("text", TEXTS)
def test_native_anthropic_text_count_matches_python(native_bridge: None, text: str) -> None:
messages: Final = [{"role": "system", "content": "You are terse."}, {"role": "user", "name": "bob", "content": text}]
litellm.rust(False)
python_text: Final = token_counter(model=ANTHROPIC_MODEL, text=text)
python_messages: Final = token_counter(model=ANTHROPIC_MODEL, messages=messages)
litellm.rust(True)
rust_count: Final = bridge.text_counter("anthropic")
assert rust_count is not None
assert rust_count(text) == python_text
assert token_counter(model=ANTHROPIC_MODEL, text=text) == python_text
assert token_counter(model=ANTHROPIC_MODEL, messages=messages) == python_messages