From 7c9f2c32c835082db12ed2343c07965bca52ed67 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 5 Feb 2026 16:18:36 -0800 Subject: [PATCH] added retry logic for gaurdrails --- litellm/proxy/guardrails/init_guardrails.py | 15 +- litellm/proxy/utils.py | 100 +++++++++++++- litellm/router.py | 56 ++++++-- litellm/types/guardrails.py | 5 + litellm/types/router.py | 1 + tests/test_litellm/test_router.py | 144 +++++++++++++++----- 6 files changed, 263 insertions(+), 58 deletions(-) diff --git a/litellm/proxy/guardrails/init_guardrails.py b/litellm/proxy/guardrails/init_guardrails.py index 5db61eb9c51..3655519009f 100644 --- a/litellm/proxy/guardrails/init_guardrails.py +++ b/litellm/proxy/guardrails/init_guardrails.py @@ -78,14 +78,17 @@ def _populate_router_guardrail_list(guardrail_list: List[Guardrail]) -> None: else dict(litellm_params) ) + router_guardrail_litellm_params: Dict[str, Any] = { + "guardrail": params_dict.get("guardrail", ""), + "mode": params_dict.get("mode", ""), + "api_key": params_dict.get("api_key"), + "api_base": params_dict.get("api_base"), + } + if params_dict.get("num_retries") is not None: + router_guardrail_litellm_params["num_retries"] = params_dict["num_retries"] router_guardrail: GuardrailTypedDict = GuardrailTypedDict( guardrail_name=guardrail_name or "", - litellm_params={ - "guardrail": params_dict.get("guardrail", ""), - "mode": params_dict.get("mode", ""), - "api_key": params_dict.get("api_key"), - "api_base": params_dict.get("api_base"), - }, + litellm_params=router_guardrail_litellm_params, callback=callback, id=guardrail_id, ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 6bbf0df74de..9ddd3a2271c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -841,6 +841,61 @@ class ProxyLogging: ] return len(matching) > 1 + def _should_use_guardrail_via_router(self, guardrail_name: str) -> bool: + """ + 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. + """ + 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 + ) + + async def _execute_guardrail_via_router( + self, + guardrail_name: str, + hook_type: str, + data: dict, + user_api_key_dict: Optional[UserAPIKeyAuth], + call_type: CallTypesLiteral, + response: Optional[Any] = None, + ) -> Any: + """ + Execute a guardrail via router.aguardrail() so retries and fallbacks apply. + """ + from litellm.proxy.proxy_server import llm_router + + if llm_router is None: + raise ValueError("Router not initialized") + + async def runner(**kwargs: Any) -> Any: + selected_guardrail = kwargs.get("selected_guardrail") + if selected_guardrail is None: + raise ValueError("selected_guardrail not in kwargs") + callback = selected_guardrail.get("callback") + if callback is None: + raise ValueError( + f"No callback found for guardrail: {guardrail_name}" + ) + return await self._execute_guardrail_hook( + callback=callback, + hook_type=hook_type, + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + response=response, + ) + + return await llm_router.aguardrail( + guardrail_name=guardrail_name, + original_function=runner, + ) + async def _execute_guardrail_hook( self, callback: "CustomGuardrail", @@ -981,10 +1036,21 @@ class ProxyLogging: error_type = None try: - # Check if load balancing should be used - if guardrail_name and self._should_use_guardrail_load_balancing( + # Use router path when guardrail is in router (enables retries/fallbacks) + if guardrail_name and self._should_use_guardrail_via_router( guardrail_name ): + response = await self._execute_guardrail_via_router( + guardrail_name=guardrail_name, + hook_type="pre_call", + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + ) + elif guardrail_name and self._should_use_guardrail_load_balancing( + guardrail_name + ): + # Fallback: guardrail not in router guardrail_list response = await self._execute_guardrail_with_load_balancing( guardrail_name=guardrail_name, hook_type="pre_call", @@ -1313,8 +1379,20 @@ class ProxyLogging: ) else: user_api_key_auth_dict = user_api_key_dict - # Add task to list for parallel execution + guardrail_name = getattr(callback, "guardrail_name", None) + # Use router path when guardrail is in router (enables retries/fallbacks) if ( + guardrail_name + and self._should_use_guardrail_via_router(guardrail_name) + ): + guardrail_task = self._execute_guardrail_via_router( + guardrail_name=guardrail_name, + hook_type="during_call", + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + ) + elif ( "apply_guardrail" in type(callback).__dict__ and user_api_key_dict is not None ): @@ -1775,8 +1853,22 @@ class ProxyLogging: continue guardrail_response: Optional[Any] = None + guardrail_name = getattr(callback, "guardrail_name", None) - if "apply_guardrail" in type(callback).__dict__: + # Use router path when guardrail is in router (enables retries/fallbacks) + if ( + guardrail_name + and self._should_use_guardrail_via_router(guardrail_name) + ): + guardrail_response = await self._execute_guardrail_via_router( + guardrail_name=guardrail_name, + hook_type="post_call", + data=data, + user_api_key_dict=user_api_key_dict, + call_type=CallTypes.acompletion.value, # post_call hook doesn't use call_type + response=response, + ) + elif "apply_guardrail" in type(callback).__dict__: data["guardrail_to_apply"] = callback guardrail_response = ( await unified_guardrail.async_post_call_success_hook( diff --git a/litellm/router.py b/litellm/router.py index d01c8443dab..a0ffb2168cb 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -239,6 +239,7 @@ class Router: default_priority: Optional[int] = None, ## RELIABILITY ## num_retries: Optional[int] = None, + guardrail_num_retries: Optional[int] = None, # num retries for guardrail API calls; defaults to num_retries when None max_fallbacks: Optional[ int ] = None, # max fallbacks to try before exiting the call. Defaults to 5. @@ -490,6 +491,8 @@ class Router: else: self.num_retries = openai.DEFAULT_MAX_RETRIES + self.guardrail_num_retries = guardrail_num_retries + if max_fallbacks is not None: self.max_fallbacks = max_fallbacks elif litellm.max_fallbacks is not None: @@ -3180,6 +3183,9 @@ class Router: """ Execute a guardrail with load balancing and fallbacks. + Uses _ageneric_api_call_with_fallbacks with use_guardrail_list=True so guardrails + share the same retry/fallback path as generic API calls. + Args: guardrail_name: Name of the guardrail to execute original_function: The guardrail's execution function (e.g., async_pre_call_hook) @@ -3188,19 +3194,12 @@ class Router: Returns: Result from the guardrail execution """ - kwargs["model"] = guardrail_name # For fallback system compatibility - kwargs["original_generic_function"] = original_function - kwargs["original_function"] = self._aguardrail_helper - self._update_kwargs_before_fallbacks( + return await self._ageneric_api_call_with_fallbacks( model=guardrail_name, - kwargs=kwargs, - metadata_variable_name="litellm_metadata", + original_function=original_function, + use_guardrail_list=True, + **kwargs, ) - verbose_router_logger.debug( - f"Inside aguardrail() - guardrail_name: {guardrail_name}; kwargs: {kwargs}" - ) - response = await self.async_function_with_fallbacks(**kwargs) - return response async def _aguardrail_helper( self, @@ -3267,15 +3266,30 @@ class Router: ) async def _ageneric_api_call_with_fallbacks( - self, model: str, original_function: Callable, **kwargs + self, + model: str, + original_function: Callable, + use_guardrail_list: bool = False, + **kwargs, ): """ - Helper function to make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router + Helper function to make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router. + + When use_guardrail_list=True, treats model as guardrail_name and selects from + guardrail_list (same retry/fallback path as aguardrail). """ try: kwargs["model"] = model kwargs["original_generic_function"] = original_function kwargs["original_function"] = self._ageneric_api_call_with_fallbacks_helper + if use_guardrail_list: + kwargs["use_guardrail_list"] = True + if kwargs.get("num_retries") is None: + kwargs["num_retries"] = ( + self.guardrail_num_retries + if self.guardrail_num_retries is not None + else self.num_retries + ) self._update_kwargs_before_fallbacks( model=model, kwargs=kwargs, metadata_variable_name="litellm_metadata" ) @@ -3332,9 +3346,18 @@ class Router: self, model: str, original_generic_function: Callable, **kwargs ): """ - Helper function to make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router + Helper function to make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router. + When use_guardrail_list=True, selects from guardrail_list and calls original_generic_function with selected_guardrail in kwargs. """ + use_guardrail_list = kwargs.pop("use_guardrail_list", False) + if use_guardrail_list: + selected_guardrail = self.get_available_guardrail( + guardrail_name=model, + ) + kwargs["selected_guardrail"] = selected_guardrail + return await original_generic_function(**kwargs) + passthrough_on_no_deployment = kwargs.pop("passthrough_on_no_deployment", False) function_name = "_ageneric_api_call_with_fallbacks" try: @@ -4918,6 +4941,11 @@ class Router: except Exception as e: current_attempt = None original_exception = e + # Per-guardrail num_retries: set on exception from selected_guardrail so retry loop uses it + if kwargs.get("selected_guardrail") is not None: + self._set_deployment_num_retries_on_exception( + e, kwargs["selected_guardrail"] + ) deployment_num_retries = getattr(e, "num_retries", None) if deployment_num_retries is not None and isinstance( diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 74ccb34ca6e..c1e8a276536 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -573,6 +573,11 @@ class BaseLitellmParams( default=None, description="Base URL for the guardrail service API" ) + num_retries: Optional[int] = Field( + default=None, + description="Number of retries for this guardrail's API calls; overrides router guardrail_num_retries when set", + ) + experimental_use_latest_role_message_only: Optional[bool] = Field( default=False, description="When True, guardrails only receive the latest message for the relevant role (e.g., newest user input pre-call, newest assistant output post-call)", diff --git a/litellm/types/router.py b/litellm/types/router.py index f78789c9772..3b26074978f 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -656,6 +656,7 @@ class GuardrailLiteLLMParams(TypedDict, total=False): api_key: Optional[str] api_base: Optional[str] weight: Optional[int] # For load balancing + num_retries: Optional[int] # Per-guardrail retries; overrides router guardrail_num_retries when set class GuardrailTypedDict(TypedDict, total=False): diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 08ae804ea80..bad15ca2bac 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1803,40 +1803,6 @@ def test_get_available_guardrail_not_found(): router.get_available_guardrail(guardrail_name="non-existent") -@pytest.mark.asyncio -async def test_aguardrail_helper(): - """ - Test _aguardrail_helper selects a guardrail and executes the original function. - """ - guardrail_config = { - "guardrail_name": "content-filter", - "litellm_params": {"guardrail": "custom", "mode": "pre_call"}, - "id": "guardrail-1", - } - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo"}, - } - ], - guardrail_list=[guardrail_config], - ) - - # Mock the original function - async def mock_original_function(**kwargs): - return {"result": "success", "selected_guardrail": kwargs.get("selected_guardrail")} - - result = await router._aguardrail_helper( - model="content-filter", - original_generic_function=mock_original_function, - ) - - assert result["result"] == "success" - assert result["selected_guardrail"] == guardrail_config - - @pytest.mark.asyncio async def test_aguardrail(): """ @@ -1869,3 +1835,113 @@ async def test_aguardrail(): assert result["result"] == "success" assert result["selected_guardrail"]["id"] == "guardrail-1" + + +@pytest.mark.asyncio +async def test_aguardrail_retry_succeeds_after_retries(): + """ + Test guardrail retry logic: guardrail call fails twice then succeeds on 3rd attempt. + Per-guardrail num_retries=2 so we get 1 initial + 2 retries = 3 attempts. + """ + guardrail_config = { + "guardrail_name": "flaky-guardrail", + "litellm_params": { + "guardrail": "custom", + "mode": "pre_call", + "num_retries": 2, + }, + "id": "guardrail-retry-1", + } + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo"}, + } + ], + guardrail_list=[guardrail_config], + ) + + call_count = 0 + + async def mock_fail_twice_then_succeed(**kwargs): + nonlocal call_count + call_count += 1 + if call_count < 3: + 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, + ): + result = await router.aguardrail( + guardrail_name="flaky-guardrail", + original_function=mock_fail_twice_then_succeed, + ) + + assert call_count == 3 + assert result["result"] == "success" + assert result["attempt"] == 3 + + +@pytest.mark.asyncio +async def test_aguardrail_retry_exhausted_raises(): + """ + Test guardrail retry logic: when all attempts fail, the exception is raised. + With guardrail num_retries=2 we get at least 2 attempts (1 initial + retries) before raising. + """ + guardrail_config = { + "guardrail_name": "failing-guardrail", + "litellm_params": { + "guardrail": "custom", + "mode": "pre_call", + "num_retries": 2, + }, + "id": "guardrail-fail-1", + } + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo"}, + } + ], + guardrail_list=[guardrail_config], + ) + + call_count = 0 + + async def mock_always_fails(**kwargs): + nonlocal call_count + 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, + ): + with pytest.raises(ValueError, match="guardrail API always fails"): + await router.aguardrail( + guardrail_name="failing-guardrail", + original_function=mock_always_fails, + ) + + # At least 2 attempts (1 initial + 1 or more retries); exact count depends on router default vs per-guardrail num_retries + assert call_count >= 2