fix(rust): count the proxy's preselected Anthropic tokenizer in Rust

The proxy token counting route selects the tokenizer itself and passes it in as custom_tokenizer, so
only a tokenizer other than the one token_counter selects for the model stays in Python

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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:

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
 
 

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)