diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 73c6b20e1a8..4b35e790de5 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -38,11 +38,12 @@ else: _EXCEPTION_POLICY_FIELDS: tuple[tuple[type, str], ...] = ( + # ContentPolicyViolationError subclasses BadRequestError, so it must be checked first. + (litellm.ContentPolicyViolationError, "ContentPolicyViolationErrorAllowedFails"), (litellm.BadRequestError, "BadRequestErrorAllowedFails"), (litellm.AuthenticationError, "AuthenticationErrorAllowedFails"), (litellm.Timeout, "TimeoutErrorAllowedFails"), (litellm.RateLimitError, "RateLimitErrorAllowedFails"), - (litellm.ContentPolicyViolationError, "ContentPolicyViolationErrorAllowedFails"), (litellm.InternalServerError, "InternalServerErrorAllowedFails"), (litellm.ServiceUnavailableError, "ServiceUnavailableErrorAllowedFails"), (litellm.BadGatewayError, "BadGatewayErrorAllowedFails"), @@ -103,14 +104,22 @@ def _should_cooldown_based_on_deployment_policy( dep_policy: dict[str, int] | None, dep_allowed_fails: int | None, ) -> bool: - """Resolve deployment-level allowed-fails and delegate to the shared counting logic.""" + """Resolve deployment-level allowed-fails and delegate to the shared counting logic. + + When the deployment's policy doesn't cover *original_exception*'s type and no + deployment-wide `allowed_fails` is set either, defer to router-level behavior + instead of forcing an immediate cooldown. + """ allowed_fails_from_policy = _resolve_allowed_fails_from_policy(dep_policy, original_exception) if allowed_fails_from_policy is not None: - allowed_fails_override: int = allowed_fails_from_policy - cache_key_suffix: str = type(original_exception).__name__ - else: - allowed_fails_override = dep_allowed_fails if dep_allowed_fails is not None else 0 + allowed_fails_override: int | None = allowed_fails_from_policy + cache_key_suffix: str | None = type(original_exception).__name__ + elif dep_allowed_fails is not None: + allowed_fails_override = dep_allowed_fails cache_key_suffix = "generic" + else: + allowed_fails_override = None + cache_key_suffix = None dep = litellm_router_instance.get_model_info(id=deployment) cooldown_time_override: float | None = None diff --git a/tests/test_litellm/router_utils/test_cooldown_handlers.py b/tests/test_litellm/router_utils/test_cooldown_handlers.py index 21b29f031ba..138fb508fcb 100644 --- a/tests/test_litellm/router_utils/test_cooldown_handlers.py +++ b/tests/test_litellm/router_utils/test_cooldown_handlers.py @@ -135,7 +135,11 @@ class TestShouldCooldownBasedOnDeploymentPolicy: assert call_kwargs["allowed_fails_override"] == 3 assert call_kwargs["cache_key_suffix"] == "generic" - def test_no_policy_and_no_dep_allowed_fails_defaults_to_zero(self): + def test_no_policy_and_no_dep_allowed_fails_defers_to_router_level(self): + """When neither a deployment policy nor a deployment-wide allowed_fails covers + this exception, defer to router-level behavior instead of forcing an + immediate cooldown (allowed_fails_override=0 would trip on the first failure + of any exception type the deployment's config doesn't mention).""" exc = litellm.InternalServerError("500", "openai", "gpt-4") router = self._make_router({"litellm_params": {}, "model_info": {}}) @@ -146,7 +150,26 @@ class TestShouldCooldownBasedOnDeploymentPolicy: _should_cooldown_based_on_deployment_policy(router, "dep-1", exc, None, None) call_kwargs = mock_sc.call_args[1] - assert call_kwargs["allowed_fails_override"] == 0 + assert call_kwargs["allowed_fails_override"] is None + assert call_kwargs["cache_key_suffix"] is None + + def test_partial_policy_without_dep_allowed_fails_defers_for_uncovered_exception(self): + """A deployment that only sets RateLimitErrorAllowedFails must not force a + zero-fail threshold on an unrelated TimeoutError; it should defer to + router-level behavior for exception types its policy doesn't mention.""" + policy = {"RateLimitErrorAllowedFails": 0} + exc = litellm.Timeout("timed out", "openai", "gpt-4") + router = self._make_router({"litellm_params": {}, "model_info": {}}) + + with patch( + "litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy" + ) as mock_sc: + mock_sc.return_value = False + _should_cooldown_based_on_deployment_policy(router, "dep-1", exc, policy, dep_allowed_fails=None) + + call_kwargs = mock_sc.call_args[1] + assert call_kwargs["allowed_fails_override"] is None + assert call_kwargs["cache_key_suffix"] is None def test_cooldown_time_from_litellm_params_passed_through(self): exc = litellm.RateLimitError("429", "openai", "gpt-4")