mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
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:
parent
c849e0b766
commit
174037f49d
3 changed files with 59 additions and 1 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
41
tests/test_litellm_rust/test_token_counter.py
Normal file
41
tests/test_litellm_rust/test_token_counter.py
Normal 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
|
||||
Loading…
Add table
Reference in a new issue