scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

added retry logic for gaurdrails

Browse source
This commit is contained in:
shivam 2026-02-05 16:18:36 -08:00
parent 0649720f79
commit 7c9f2c32c8
6 changed files with 263 additions and 58 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
litellm/proxy/guardrails/init_guardrails.py
View file

@ -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,
)

100
litellm/proxy/utils.py
View file

@ -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(

56
litellm/router.py
View file

@ -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(

5
litellm/types/guardrails.py
View file

@ -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)",

1
litellm/types/router.py
View file

@ -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):

144
tests/test_litellm/test_router.py
View file

@ -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