diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index dda9ed0b566..ca271e201b0 100644 --- a/litellm/rust_bridge/catalog.py +++ b/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)) diff --git a/litellm/rust_bridge/diagnostics.py b/litellm/rust_bridge/diagnostics.py index fbe1b8122d7..90a5672a500 100644 --- a/litellm/rust_bridge/diagnostics.py +++ b/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: diff --git a/litellm/rust_bridge/dispatch.py b/litellm/rust_bridge/dispatch.py index c85339eba8f..da9b9bed54f 100644 --- a/litellm/rust_bridge/dispatch.py +++ b/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: diff --git a/tests/unit/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py index 1ba95507d03..b99342388d6 100644 --- a/tests/unit/rust_bridge/test_catalog.py +++ b/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