diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index ac9c6097254..45b7d251735 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -615,7 +615,8 @@ def _get_exact_count_function( ) -> TokenCounterFunction: """ Get the function to count tokens based on the model and custom tokenizer.""" - from litellm.utils import _select_tokenizer + from litellm.rust_bridge.token_counter import text_counter + from litellm.utils import _select_tokenizer, huggingface_tokenizer_kind if model is not None or custom_tokenizer is not None: tokenizer_json: Final = custom_tokenizer or _select_tokenizer(model) @@ -625,9 +626,11 @@ def _get_exact_count_function( def count_tokens(text: str) -> int: return len(tokenizer.encode_batch_fast([text])[0]) - rust_count: Final = ( - None if custom_tokenizer is not None or model is None else _rust_anthropic_count_function(model) - ) + if model is None or huggingface_tokenizer_kind(model) != "anthropic": + return count_tokens + if tokenizer is not _select_tokenizer(model)["tokenizer"]: + return count_tokens + rust_count: Final = text_counter("anthropic") 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) @@ -646,16 +649,6 @@ 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: diff --git a/litellm/rust_bridge/token_counter.py b/litellm/rust_bridge/token_counter.py index 24d48e4508e..a6ed403a5d6 100644 --- a/litellm/rust_bridge/token_counter.py +++ b/litellm/rust_bridge/token_counter.py @@ -103,6 +103,7 @@ def text_counter(tokenizer: RustTokenizer) -> Callable[[str], int] | None: factory: Final = TOKEN_COUNTER.load() if factory is None: return None + verbose_logger.debug("Rust token counter (%s) counting text", tokenizer) return _counter(factory, tokenizer).count_text 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 c3f572eb0b0..4b4da3d0e54 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -1677,6 +1677,22 @@ def test_custom_huggingface_tokenizer_stays_in_python_for_anthropic_models(rust_ assert factory.counters == [] +def test_preselected_anthropic_tokenizer_is_counted_by_rust(rust_bridge) -> None: + """The proxy's token counting route selects the tokenizer itself and passes it in as `custom_tokenizer`.""" + from litellm.utils import _select_tokenizer + + factory: Final = _FakeTextCounterFactory(1_000) + litellm.rust(True) + rust_bridge.TOKEN_COUNTER.override(factory) + + count: Final = token_counter_new( + model=ANTHROPIC_MODEL, custom_tokenizer=_select_tokenizer(ANTHROPIC_MODEL), text="hello" + ) + + assert count == 1_000 + assert [counter.texts for counter in factory.counters] == [["hello"]] + + def test_disabled_bridge_counts_anthropic_text_in_python(rust_bridge) -> None: factory: Final = _FakeTextCounterFactory(1_000) litellm.rust(False)