refactor(rust-bridge): remove the LoggerRule concept from the catalog (#46711)

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

litellm/rust_bridge/catalog.py

@@ -1,4 +1,4 @@
-"""Ordered rollout policy for routes and loggers.
+"""Ordered rollout policy for routes.
 
 The first matching rule wins; unmatched contexts stay on Python. Native
 admission separately decides whether the selected implementation can execute.
@@ -39,34 +39,17 @@ class RouteRule:
     providers: frozenset[str] | None = None
     models: frozenset[str] | None = None
 
-    def matches(self, context: Context) -> bool:
+    def matches(self, context: RouteContext) -> bool:
         return (
-            isinstance(context, RouteContext)
-            and context.route is self.route
+            context.route is self.route
             and (self.providers is None or context.provider in self.providers)
             and (self.models is None or context.model in self.models)
         )
 
 
-@dataclass(frozen=True, slots=True)
-class LoggerContext:
-    pass
-
-
-@dataclass(frozen=True, slots=True)
-class LoggerRule:
-    rollout: Rollout
-
-    def matches(self, context: Context) -> bool:
-        return isinstance(context, LoggerContext)
-
-
-Context: TypeAlias = RouteContext | LoggerContext
-Rule: TypeAlias = RouteRule | LoggerRule
-Rules: TypeAlias = tuple[Rule, ...]
+Rules: TypeAlias = tuple[RouteRule, ...]
 
 RULES: Final[Rules] = (
-    LoggerRule(Rollout.RUST_OPT_IN),
     RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY),
     RouteRule(Route.EMBEDDINGS, Rollout.PYTHON_ONLY),
     RouteRule(Route.OCR, Rollout.RUST_REQUIRED),
@@ -79,10 +62,10 @@ RULES: Final[Rules] = (
 )
 
 
-def rollout(context: Context, rules: Rules | None = None) -> Rollout:
+def rollout(context: RouteContext, rules: Rules | None = None) -> Rollout:
     selected_rules: Final = RULES if rules is None else rules
     return next((rule.rollout for rule in selected_rules if rule.matches(context)), Rollout.PYTHON_ONLY)
 
 
-def decision(context: Context, rules: Rules | None = None) -> Decision:
+def decision(context: RouteContext, rules: Rules | None = None) -> Decision:
     return _decision(rollout(context, rules))

litellm/rust_bridge/diagnostics.py

@@ -44,11 +44,9 @@ def _construct(factory: NativeDiagnosticFactory, minimum_custom_key_length: int)
 
 def run(native: Callable[[NativeDiagnosticProcessor], ResultT], python: Callable[[], ResultT]) -> ResultT:
     from litellm.constants import MINIMUM_CUSTOM_KEY_LENGTH
-    from litellm.rust_bridge.catalog import LoggerContext, decision
-    from litellm.rust_bridge.configuration import Decision
+    from litellm.rust_bridge.configuration import rust_enabled
 
-    selected: Final = decision(LoggerContext())
-    if selected is Decision.PYTHON:
+    if not rust_enabled():
         return python()
     factory: Final = PROCESSOR.load()
     if factory is None:

litellm/rust_bridge/dispatch.py

@@ -6,7 +6,7 @@ from typing import Final, Generic, TypeAlias, TypeVar
 
 from litellm.rust_bridge import catalog, runtime
 from litellm.rust_bridge.bindings import NativeBinding
-from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule, Rules
+from litellm.rust_bridge.catalog import Route, RouteContext, Rules
 from litellm.rust_bridge.configuration import Decision
 from litellm.rust_bridge.configuration import decision as rollout_decision
 
@@ -35,7 +35,7 @@ class PublicDispatch(Generic[RequestT]):
 
     def _requires_projection(self, rules: Rules) -> bool:
         for rule in rules:
-            if not isinstance(rule, RouteRule) or rule.route is not self.route:
+            if rule.route is not self.route:
                 continue
             if rule.providers is not None or rule.models is not None:
                 if rollout_decision(rule.rollout) is not Decision.PYTHON:

tests/unit/rust_bridge/test_catalog.py

@@ -7,9 +7,6 @@ import pytest
 
 from litellm.rust_bridge import catalog, configuration
 from litellm.rust_bridge.catalog import (
-    Context,
-    LoggerContext,
-    LoggerRule,
     Route,
     RouteContext,
     RouteRule,
@@ -65,13 +62,6 @@ def test_missing_rule_stays_on_python_even_when_rust_is_enabled(monkeypatch: pyt
     assert catalog.decision(RouteContext(route), rules=()) is Decision.PYTHON
 
 
-def test_logger_rollout_obeys_the_global_switch() -> None:
-    assert catalog.rollout(LoggerContext()) is Rollout.RUST_OPT_IN
-    assert catalog.decision(LoggerContext()) is Decision.PYTHON
-    configuration.rust(True)
-    assert catalog.decision(LoggerContext()) is Decision.RUST_WITH_FALLBACK
-
-
 @pytest.mark.parametrize(
     ("context", "expected"),
     (
@@ -117,12 +107,10 @@ def test_ocr_has_no_python_path_to_opt_out_to(
         (RouteContext(Route.OCR, provider="local"), Decision.RUST_REQUIRED),
         (RouteContext(Route.OCR, provider="other"), Decision.PYTHON),
         (RouteContext(Route.MESSAGES, provider="local"), Decision.PYTHON),
-        (LoggerContext(), Decision.RUST_REQUIRED),
     ),
 )
-def test_mixed_rules_select_only_the_matching_domain(context: Context, expected: Decision) -> None:
+def test_rules_select_only_the_matching_route(context: RouteContext, expected: Decision) -> None:
     rules: Final[Rules] = (
-        LoggerRule(Rollout.RUST_REQUIRED),
         RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"local"})),
         RouteRule(Route.OCR, Rollout.PYTHON_ONLY),
     )
@@ -130,7 +118,6 @@ def test_mixed_rules_select_only_the_matching_domain(context: Context, expected:
     assert catalog.decision(context, rules) is expected
 
 
-@pytest.mark.parametrize("context", (RouteContext(Route.OCR), LoggerContext()))
 @pytest.mark.parametrize(
     ("rollout", "process", "environment", "expected"),
     (
@@ -144,9 +131,8 @@ def test_mixed_rules_select_only_the_matching_domain(context: Context, expected:
         (Rollout.RUST_OPT_OUT, True, "0", Decision.PYTHON),
     ),
 )
-def test_all_domains_share_rollout_switches_and_first_match(
+def test_rollout_switches_and_first_match(
     monkeypatch: pytest.MonkeyPatch,
-    context: Context,
     rollout: Rollout,
     process: bool | None,
     environment: str | None,
@@ -157,19 +143,17 @@ def test_all_domains_share_rollout_switches_and_first_match(
         monkeypatch.setenv("LITELLM_RUST", environment)
     rules: Final[Rules] = (
         RouteRule(Route.OCR, rollout),
-        LoggerRule(rollout),
         RouteRule(Route.OCR, Rollout.RUST_REQUIRED),
     )
 
-    assert catalog.decision(context, rules) is expected
-    assert catalog.decision(context, ()) is Decision.PYTHON
+    assert catalog.decision(RouteContext(Route.OCR), rules) is expected
+    assert catalog.decision(RouteContext(Route.OCR), ()) is Decision.PYTHON
 
 
-@pytest.mark.parametrize("context", (RouteContext(Route.OCR),))
-def test_empty_constraints_match_nothing(context: Context) -> None:
+def test_empty_constraints_match_nothing() -> None:
     rules: Final[Rules] = (
         RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset()),
         RouteRule(Route.OCR, Rollout.RUST_REQUIRED, models=frozenset()),
     )
 
-    assert catalog.decision(context, rules) is Decision.PYTHON
+    assert catalog.decision(RouteContext(Route.OCR), rules) is Decision.PYTHON