From b656d80d7b9bbc454bc80fb4b56d3c77a113bd71 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 5 Feb 2026 16:41:33 -0800 Subject: [PATCH] fixed greptile issues --- litellm/proxy/utils.py | 22 ++++++++++++++++++---- litellm/router.py | 22 ++++++++++++++-------- tests/test_litellm/test_router.py | 26 ++++---------------------- 3 files changed, 36 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9ddd3a2271c..b47279646ae 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -17,6 +17,7 @@ from typing import ( List, Literal, Optional, + Set, Union, cast, overload, @@ -329,6 +330,10 @@ class ProxyLogging: self.daily_report_started: bool = False self.hanging_requests_check_started: bool = False + # Cache guardrail names for router path (avoids repeated linear scans per request) + self._guardrail_names_via_router_cache: Optional[Set[str]] = None + self._guardrail_names_via_router_cache_router_id: Optional[int] = None + def startup_event( self, llm_router: Optional[Router], @@ -846,15 +851,24 @@ class ProxyLogging: Check if this guardrail should be executed via the router (for retries/fallbacks). Returns True when the router exists and has this guardrail in its guardrail_list. + Uses a cache keyed by router id to avoid repeated linear scans of guardrail_list per request. """ from litellm.proxy.proxy_server import llm_router if llm_router is None or not hasattr(llm_router, "guardrail_list"): return False - return any( - g.get("guardrail_name") == guardrail_name - for g in llm_router.guardrail_list - ) + router_id = id(llm_router) + if ( + self._guardrail_names_via_router_cache_router_id != router_id + or self._guardrail_names_via_router_cache is None + ): + self._guardrail_names_via_router_cache = { + g.get("guardrail_name") + for g in llm_router.guardrail_list + if g.get("guardrail_name") + } + self._guardrail_names_via_router_cache_router_id = router_id + return guardrail_name in self._guardrail_names_via_router_cache async def _execute_guardrail_via_router( self, diff --git a/litellm/router.py b/litellm/router.py index fbf451539a6..ddf2d5f4c68 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4954,13 +4954,12 @@ class Router: Retry Logic """ # For guardrail calls, use guardrail_list so should_retry_this_error allows retries - if kwargs.get("selected_guardrail") is not None: - _guardrail_name = kwargs.get("model") or "" - _healthy_deployments = [ - g - for g in self.guardrail_list - if g.get("guardrail_name") == _guardrail_name - ] + _model = kwargs.get("model") or "" + _guardrail_deployments = [ + g for g in self.guardrail_list if g.get("guardrail_name") == _model + ] + if kwargs.get("selected_guardrail") is not None or _guardrail_deployments: + _healthy_deployments = _guardrail_deployments _all_deployments = _healthy_deployments else: ( @@ -5044,7 +5043,14 @@ class Router: kwargs = self.log_retry(kwargs=kwargs, e=e) remaining_retries = num_retries - current_attempt - 1 _model = kwargs.get("model") # type: ignore - if kwargs.get("selected_guardrail") is not None and _model: + _is_guardrail = ( + kwargs.get("selected_guardrail") is not None + or any( + g.get("guardrail_name") == _model + for g in self.guardrail_list + ) + ) + if _is_guardrail and _model: _healthy_deployments = [ g for g in self.guardrail_list diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index bad15ca2bac..3ac04deacbd 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1872,18 +1872,8 @@ async def test_aguardrail_retry_succeeds_after_retries(): raise ValueError("guardrail API temporarily failed") return {"result": "success", "attempt": call_count} - # Allow retries: router's should_retry_this_error requires healthy_deployments > 0 - # for non-RateLimit errors; guardrail "model" has no LLM deployments so we patch. - with patch.object( - router, - "_async_get_healthy_deployments", - new_callable=AsyncMock, - return_value=([{"litellm_params": {}}], []), - ), patch.object( - router, - "_time_to_sleep_before_retry", - return_value=0, - ): + # Skip sleep so test stays fast; guardrail retries use guardrail_list for healthy_deployments (no patch needed) + with patch.object(router, "_time_to_sleep_before_retry", return_value=0): result = await router.aguardrail( guardrail_name="flaky-guardrail", original_function=mock_fail_twice_then_succeed, @@ -1927,16 +1917,8 @@ async def test_aguardrail_retry_exhausted_raises(): call_count += 1 raise ValueError("guardrail API always fails") - with patch.object( - router, - "_async_get_healthy_deployments", - new_callable=AsyncMock, - return_value=([{"litellm_params": {}}], []), - ), patch.object( - router, - "_time_to_sleep_before_retry", - return_value=0, - ): + # Skip sleep so test stays fast; guardrail retries use guardrail_list for healthy_deployments (no patch needed) + with patch.object(router, "_time_to_sleep_before_retry", return_value=0): with pytest.raises(ValueError, match="guardrail API always fails"): await router.aguardrail( guardrail_name="failing-guardrail",