diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 4fcf555035d..acf1445ac36 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -4597,7 +4597,8 @@ class ComplexityRouter(CustomLogger): pinned_pin.tier if pinned_pin.tier is not None else self._tier_for_model(pinned_pin.model), ) ) - if selected_pin is not None: + pin_plan_sentinel: Final = self._matched_plan_mode_signal(request_kwargs, resolved_messages) + if selected_pin is not None and (selected_pin.tier is not None or pin_plan_sentinel is None): escalated: Final = selected_pin.model != pinned_pin.model or ( pin_escalation_keyword is not None and pinned_pin.tier is not None @@ -4624,7 +4625,6 @@ class ComplexityRouter(CustomLogger): # the floor, and the stored pin deliberately keeps the session's own model so # the first turn after plan mode exits auto-routes exactly as it would have. # Escalation is the opposite on purpose -- an explicit ask to re-pin higher. - pin_plan_sentinel: Final = self._matched_plan_mode_signal(request_kwargs, resolved_messages) pinned_tier: Final = resolved_pin_tier if pin_plan_sentinel is not None else None plan_floored: Final = ( pinned_tier is not None and self._apply_plan_mode_floor(pinned_tier) != pinned_tier @@ -4772,10 +4772,16 @@ class ComplexityRouter(CustomLogger): # classified at or above the floor keeps its ordinary cause, yet on an adaptive router # the hard floor constrained its pick, so pinning it would carry a plan-mode-shaped # choice past plan mode's exit. No sentinel turn writes the pin, whatever its cause. + holds_fallback_for_this_ask: Final = ( + not self.config.session_affinity + and response is not None + and response.routing_decision is not None + and response.routing_decision.get("cause") == "default_model_fallback" + ) pinnable: Final = ( cache_key is not None and response is not None - and _decision_is_pinnable(response.routing_decision) + and (_decision_is_pinnable(response.routing_decision) or holds_fallback_for_this_ask) and self._matched_plan_mode_signal(request_kwargs, resolved_messages) is None ) if pinnable and cache_key is not None and response is not None: diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index f06cb0d282d..59f637bb294 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -8284,6 +8284,130 @@ class TestClassificationMode: assert suppressed._uses_deployment_pin is False +class TestUserTurnClassifierFailureHoldsTheDefaultModel: + """In user_turn mode a classifier failure holds the default model for that ask, so the ask's + continuations never replay the previous ask's decision.""" + + HELLO: Final[dict[str, object]] = {"role": "user", "content": "Hello"} + HELLO_REPLY: Final[dict[str, object]] = {"role": "assistant", "content": "Hello! How can I help?"} + FIX: Final[dict[str, object]] = {"role": "user", "content": "Find and fix the concurrency bug in the queue worker."} + TOOL_CALL: Final[dict[str, object]] = { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "t1", "type": "function", "function": {"name": "read_file", "arguments": "{}"}}], + } + TOOL_RESULT: Final[dict[str, object]] = {"role": "tool", "tool_call_id": "t1", "content": "File contents returned."} + + @staticmethod + def _router( + *classifier_replies: object, default_model: str = "gpt-4o", plan_mode_min_tier: str | None = None + ) -> tuple[ComplexityRouter, MagicMock]: + router_instance: Final = MagicMock() + router_instance.health.return_value.store = DualCache() + router_instance.acompletion = AsyncMock(side_effect=list(classifier_replies)) + router: Final = ComplexityRouter( + model_name="test-router", + ctx=RoutingContext(router_instance), + complexity_router_config={ + "tiers": { + "SIMPLE": "gpt-4o-mini", + "MEDIUM": "gpt-4o", + "COMPLEX": "claude-sonnet-4-20250514", + "REASONING": "o1-preview", + }, + "classifier_type": "llm", + "classifier_llm_config": { + "model": "haiku-classifier", + "timeout_ms": 400, + "circuit_breaker_enabled": False, + }, + "classifier_fallback": "default_model", + "default_model": default_model, + "classification_mode": "user_turn", + **({"plan_mode_min_tier": plan_mode_min_tier} if plan_mode_min_tier is not None else {}), + }, + ) + return router, router_instance + + @staticmethod + async def _route( + router: ComplexityRouter, messages: list[dict[str, object]] | None, **request: object + ) -> tuple[str, str]: + result: Final = await router.async_pre_routing_hook( + model="test-model", + request_kwargs={"metadata": {"session_id": "fallback-holds-ask"}, **request}, + messages=messages, + input=request.get("input"), + ) + assert result is not None and result.routing_decision is not None + return result.model, result.routing_decision["cause"] + + @pytest.mark.asyncio + async def test_a_failed_ask_keeps_the_default_model_for_its_continuations(self): + router, router_instance = self._router( + _llm_response('{"tier": "SIMPLE"}'), TimeoutError("classifier timed out"), _llm_response('{"tier": "SIMPLE"}') + ) + fix_ask: Final = [self.HELLO, self.HELLO_REPLY, self.FIX] + fix_loop: Final = [*fix_ask, self.TOOL_CALL, self.TOOL_RESULT] + thanks: Final[dict[str, object]] = {"role": "user", "content": "Thanks"} + + routed: Final = [ + await self._route(router, turn) + for turn in ( + [self.HELLO], + fix_ask, + fix_loop, + [*fix_loop, self.TOOL_CALL, self.TOOL_RESULT], + [*fix_loop, {"role": "assistant", "content": "Fixed it."}, thanks], + ) + ] + + assert routed == [ + ("gpt-4o-mini", "llm_classifier"), + ("gpt-4o", "default_model_fallback"), + ("gpt-4o", "user_turn_continuation"), + ("gpt-4o", "user_turn_continuation"), + ("gpt-4o-mini", "llm_classifier"), + ] + assert router_instance.acompletion.await_count == 3 + + @pytest.mark.asyncio + async def test_a_server_side_history_continuation_keeps_the_default_model(self): + router, _ = self._router(_llm_response('{"tier": "SIMPLE"}'), TimeoutError("classifier timed out")) + tool_output: Final = [{"type": "function_call_output", "call_id": "t1", "output": "File contents returned."}] + + assert await self._route(router, None, input="Hello") == ("gpt-4o-mini", "llm_classifier") + assert await self._route(router, None, input=self.FIX["content"], previous_response_id="resp_1") == ( + "gpt-4o", + "default_model_fallback", + ) + continued: Final = await self._route(router, None, input=tool_output, previous_response_id="resp_2") + + assert continued == ("gpt-4o", "user_turn_continuation") + + @pytest.mark.asyncio + async def test_a_plan_mode_continuation_still_gets_the_plan_floor_over_a_held_fallback(self): + router, _ = self._router( + _llm_response('{"tier": "SIMPLE"}'), + TimeoutError("classifier timed out"), + _llm_response('{"tier": "SIMPLE"}'), + default_model="untiered-default", + plan_mode_min_tier="COMPLEX", + ) + fix_ask: Final = [self.HELLO, self.HELLO_REPLY, self.FIX] + plan_mode_body: Final = TestPlanModeTierFloor.PLAN_BODY + + assert await self._route(router, [self.HELLO]) == ("gpt-4o-mini", "llm_classifier") + assert await self._route(router, fix_ask) == ("untiered-default", "default_model_fallback") + continued: Final = await self._route( + router, + [*fix_ask, self.TOOL_CALL, self.TOOL_RESULT], + proxy_server_request={"body": plan_mode_body}, + ) + + assert continued == ("claude-sonnet-4-20250514", "plan_mode") + + class TestRoutingPlugins: """Test the `complexity_router_config.plugins` field: narrows the classified tier's candidate pool before a model is picked. Discussion: