diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 17e3d1256d0..9bdcb45a789 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -1694,16 +1694,23 @@ class ComplexityRouter(CustomLogger): } } - response: Final[ModelResponse] = await self.litellm_router_instance.acompletion( - model=llm_config.model, - messages=messages_for_call, - response_format=response_format, - timeout=llm_config.timeout_ms / 1000, - metadata=metadata, - proxy_server_request=proxy_server_request, - turn_off_message_logging=turn_off_message_logging, - **classifier_call_params, - **_parent_session_kwargs(request_kwargs), + classifier_timeout_s: Final[float] = llm_config.timeout_ms / 1000 + response: Final[ModelResponse] = await asyncio.wait_for( + self.litellm_router_instance.acompletion( + model=llm_config.model, + messages=messages_for_call, + stream=False, + response_format=response_format, + timeout=classifier_timeout_s, + num_retries=0, + disable_fallbacks=True, + metadata=metadata, + proxy_server_request=proxy_server_request, + turn_off_message_logging=turn_off_message_logging, + **classifier_call_params, + **_parent_session_kwargs(request_kwargs), + ), + timeout=classifier_timeout_s, ) content: Final = response.choices[0].message.content if not content: diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index da3791da39a..17594f7d444 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -1994,6 +1994,77 @@ class TestLLMClassifier: assert outcome.cause == "llm_classifier" assert outcome.classifier_cost == pytest.approx(1.35e-05) + @pytest.mark.asyncio + async def test_aclassify_timeout_does_not_inherit_router_retries_or_fallbacks( + self, llm_classifier_config + ): + real_router = Router( + model_list=[ + { + "model_name": "haiku-classifier", + "litellm_params": { + "model": "openai/mock-classifier", + "api_key": "mock-key", + "mock_timeout": True, + }, + }, + { + "model_name": "backup-classifier", + "litellm_params": { + "model": "openai/mock-backup-classifier", + "api_key": "mock-key", + "mock_response": '{"tier": "COMPLEX"}', + }, + }, + ], + num_retries=2, + fallbacks=[{"haiku-classifier": ["backup-classifier"]}], + ) + config = { + **llm_classifier_config, + "classifier_llm_config": {"model": "haiku-classifier", "timeout_ms": 10}, + } + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=real_router, + complexity_router_config=config, + ) + + outcome = await router.aclassify("hi") + + assert outcome.cause == "heuristic_scorer" + assert real_router.total_calls["openai/mock-classifier"] == 1 + assert real_router.total_calls["openai/mock-backup-classifier"] == 0 + + @pytest.mark.asyncio + async def test_aclassify_enforces_total_classifier_deadline( + self, mock_router_instance, llm_classifier_config + ): + cancelled = asyncio.Event() + + async def slow_classifier(**_kwargs: object) -> None: + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + cancelled.set() + raise + + mock_router_instance.acompletion = AsyncMock(side_effect=slow_classifier) + config = { + **llm_classifier_config, + "classifier_llm_config": {"model": "haiku-classifier", "timeout_ms": 10}, + } + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=config, + ) + + outcome = await router.aclassify("hi") + + assert outcome.cause == "heuristic_scorer" + assert cancelled.is_set() + @pytest.mark.asyncio async def test_aclassify_classifier_cost_is_none_when_call_is_unpriced( self, llm_complexity_router, mock_router_instance