diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index dda9ed0b566..0b09c6ab2d7 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -73,8 +73,8 @@ RULES: Final[Rules] = ( RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic", "deepseek", "vertex_ai"})), RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY), RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), - RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY), - RouteRule(Route.TOKENIZER, Rollout.PYTHON_ONLY), + RouteRule(Route.TOKEN_COUNTER, Rollout.RUST_OPT_IN), + RouteRule(Route.TOKENIZER, Rollout.RUST_OPT_IN), RouteRule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), ) diff --git a/tests/unit/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py index 1ba95507d03..2621ad50146 100644 --- a/tests/unit/rust_bridge/test_catalog.py +++ b/tests/unit/rust_bridge/test_catalog.py @@ -43,11 +43,14 @@ def test_shipped_decisions( if environment is not None: monkeypatch.setenv("LITELLM_RUST", environment) context: Final = RouteContext(route, provider=provider, model="test-model") + opt_in_route: Final = route in (Route.TOKEN_COUNTER, Route.TOKENIZER) or ( + route is Route.MESSAGES and provider in ("anthropic", "deepseek", "vertex_ai") + ) if route is Route.OCR or (route is Route.TRANSCRIPTION and provider == "bedrock"): assert catalog.rollout(context) is Rollout.RUST_REQUIRED assert catalog.decision(context) is Decision.RUST_REQUIRED - elif route is Route.MESSAGES and provider in ("anthropic", "deepseek", "vertex_ai"): + elif opt_in_route: assert catalog.rollout(context) is Rollout.RUST_OPT_IN opted_in: Final = environment == "1" or (environment is None and process is True) assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if opted_in else Decision.PYTHON)