diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index ca271e201b0..15a2d432761 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -56,8 +56,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 b99342388d6..d3b857c6a13 100644 --- a/tests/unit/rust_bridge/test_catalog.py +++ b/tests/unit/rust_bridge/test_catalog.py @@ -40,11 +40,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)