scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(guardrails): select one load-balanced callback per hook

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo 2026-11-05 16:40:54 +00:00
parent aa7917a3d2
commit e66f38e447
4 changed files with 112 additions and 48 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -47,10 +47,10 @@ def init_guardrails_v2(
# verbose_proxy_logger.debug(f"\nGuardrail List:{guardrail_list}\n")
# Populate router's guardrail_list for load balancing support
_populate_router_guardrail_list(guardrail_list=guardrail_list)
_populate_router_guardrail_list(guardrail_list=guardrail_list, router=llm_router)
def _populate_router_guardrail_list(guardrail_list: list[Guardrail]) -> None:
def _populate_router_guardrail_list(guardrail_list: list[Guardrail], router: Router | None = None) -> None:
"""
Populate the router's guardrail_list from initialized guardrails.
@ -61,7 +61,8 @@ def _populate_router_guardrail_list(guardrail_list: list[Guardrail]) -> None:
from litellm.proxy.proxy_server import llm_router
from litellm.types.router import GuardrailTypedDict
if llm_router is None:
target_router: Final = router if router is not None else llm_router
if target_router is None:
verbose_proxy_logger.debug("Router not initialized yet, skipping guardrail_list population")
return
@ -94,7 +95,7 @@ def _populate_router_guardrail_list(guardrail_list: list[Guardrail]) -> None:
router_guardrail_list.append(router_guardrail)
llm_router.apply(ChangeSet(assign={"guardrail_list": router_guardrail_list}))
target_router.apply(ChangeSet(assign={"guardrail_list": router_guardrail_list}))
verbose_proxy_logger.debug("Populated router guardrail_list with %s guardrails", len(router_guardrail_list))

84
litellm/proxy/utils.py
View file

@ -30,7 +30,7 @@ from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.utils import formatdate
from functools import partial, reduce
from itertools import takewhile
from itertools import chain, takewhile
from operator import attrgetter
from types import MappingProxyType
from typing import (
@ -889,6 +889,43 @@ def _drop_applied_guardrails_in_place(bucket: object, names: frozenset[str]) ->
metadata.pop("applied_guardrails")
@overload
def _load_balanced_callbacks(callbacks: Sequence[CustomLogger]) -> tuple[CustomLogger, ...]: ...
@overload
def _load_balanced_callbacks(callbacks: Sequence[object]) -> tuple[object, ...]: ...
def _load_balanced_callbacks(callbacks: Sequence[object]) -> tuple[object, ...]:
from litellm.proxy.proxy_server import llm_router
if llm_router is None or not llm_router.settings().is_set("guardrail_list"):
return tuple(callbacks)
settings: Final = llm_router.settings()
balanced_names: Final = frozenset(
entry["guardrail_name"]
for entry in settings.guardrail_list
if sum(other["guardrail_name"] == entry["guardrail_name"] for other in settings.guardrail_list) > 1
)
def selected(index: int, callback: object) -> tuple[object, ...]:
if not isinstance(callback, CustomGuardrail) or callback.guardrail_name not in balanced_names:
return (callback,)
if any(
isinstance(previous, CustomGuardrail) and previous.guardrail_name == callback.guardrail_name
for previous in callbacks[:index]
):
return ()
deployment: Final = settings.guardrail(callback.guardrail_name)
target: Final = deployment.get("callback")
if target is None:
raise ValueError(f"No callback found for guardrail: {callback.guardrail_name}")
return (target,)
return tuple(chain.from_iterable(selected(index, callback) for index, callback in enumerate(callbacks)))
def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]:
resolved: Final = tuple(
litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class(
@ -900,7 +937,7 @@ def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple
else callback
for callback in litellm.callbacks
)
present: Final = tuple(callback for callback in resolved if callback is not None)
present: Final = _load_balanced_callbacks(tuple(callback for callback in resolved if callback is not None))
guardrails: Final = tuple(callback for callback in present if isinstance(callback, CustomGuardrail))
others: Final = cast( # cast-ok: mirrors the legacy loop, which treated every non-guardrail entry as a CustomLogger
"tuple[CustomLogger, ...]",
@ -2705,24 +2742,13 @@ class ProxyLogging:
error_type = None
try:
# Check if load balancing should be used
if guardrail_name and self._should_use_guardrail_load_balancing(guardrail_name):
response = await self._execute_guardrail_with_load_balancing(
guardrail_name=guardrail_name,
hook_type="pre_call",
data=data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
)
else:
# Single guardrail - execute directly
response = await self._execute_guardrail_hook(
callback=callback,
hook_type="pre_call",
data=data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
)
response = await self._execute_guardrail_hook(
callback=callback,
hook_type="pre_call",
data=data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
)
# Process the response if one was returned
if response is not None:
@ -3265,12 +3291,15 @@ class ProxyLogging:
self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails)
return data
selected_callbacks: Final = (
caps.resolved_callbacks if skip_guardrails else _load_balanced_callbacks(caps.resolved_callbacks)
)
parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = (
()
if skip_guardrails
else tuple(
cb
for cb in caps.resolved_callbacks
for cb in selected_callbacks
if isinstance(cb, CustomGuardrail)
and getattr(cb, "run_in_parallel", False)
and not (cb.guardrail_name and cb.guardrail_name in pipeline_managed)
@ -3278,7 +3307,7 @@ class ProxyLogging:
)
deferred_route_exc: SensitiveDataRouteException | None = None
for _callback in caps.resolved_callbacks:
for _callback in selected_callbacks:
start_time = time.time()
try:
if isinstance(_callback, CustomGuardrail) and data is not None:
@ -3750,7 +3779,7 @@ class ProxyLogging:
# Step 1: Collect all guardrail tasks to run in parallel
guardrail_tasks: Final = []
for callback in litellm.callbacks:
for callback in _load_balanced_callbacks(tuple(litellm.callbacks)):
if (
isinstance(callback, CustomLogger)
and not isinstance(callback, CustomGuardrail)
@ -4724,7 +4753,7 @@ class ProxyLogging:
stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset()
)
for callback in litellm.callbacks:
for callback in _load_balanced_callbacks(tuple(litellm.callbacks)):
try:
_callback: CustomLogger | None = None
if isinstance(callback, CustomGuardrail):
@ -4814,9 +4843,14 @@ class ProxyLogging:
stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict)
pipeline_gated_names: Final = _pipeline_step_guardrail_names(post_call_pipelines)
selected_iterator_callbacks: Final = _load_balanced_callbacks(tuple(cb for cb, _ in caps.iterator_overrides))
selected_iterator_overrides: Final = tuple(
(cb, "override" if "async_post_call_streaming_iterator_hook" in type(cb).__dict__ else "apply_guardrail")
for cb in selected_iterator_callbacks
)
running_overrides: Final = tuple(
(callback, kind)
for callback, kind in caps.iterator_overrides
for callback, kind in selected_iterator_overrides
if not isinstance(callback, CustomGuardrail)
or (
callback.guardrail_name not in pipeline_gated_names

6
tests/integration/observability/test_guardrail_load_balancing.py
View file

@ -42,10 +42,6 @@ def _deployment(name: str, api_base: str) -> dict[str, object]:
def test_requests_naming_a_shared_guardrail_are_spread_over_its_deployments(gateway: Gateway, tmp_path: Path) -> None:
pytest.skip(
"BUG: init_guardrails_v2 populates llm_router.guardrail_list from the proxy_server global, which is still "
"None during load_config, so every deployment sharing a guardrail_name checks every request"
)
name: Final = f"lb-guardrail-{uuid4().hex}"
prompts: Final = tuple(f"Hello request {index}" for index in range(_REQUESTS))
with (
@ -75,5 +71,7 @@ def test_requests_naming_a_shared_guardrail_are_spread_over_its_deployments(gate
(200, name)
] * _REQUESTS
seen: Final = tuple(frozenset(json.loads(call.body)["texts"][0] for call in calls) for calls in checked)
assert sum(len(calls) for calls in checked) == _REQUESTS
assert seen[0].isdisjoint(seen[1])
assert all(0 < len(prompts_seen) < _REQUESTS for prompts_seen in seen), [len(prompts_seen) for prompts_seen in seen]
assert seen[0] | seen[1] == frozenset(prompts)

61
tests/unit/proxy/guardrails/test_guardrail_load_balancing.py
View file

@ -3,6 +3,7 @@ Test guardrail load balancing through the Router and ProxyLogging.
"""
from unittest.mock import MagicMock, patch, AsyncMock
from typing import Final
import litellm
@ -31,9 +32,26 @@ class MockGuardrail(CustomGuardrail):
MockGuardrail.call_count += 1
return None
async def async_moderation_hook(self, data, user_api_key_dict, call_type):
self.calls += 1
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
self.calls += 1
return response
async def async_post_call_streaming_hook(self, user_api_key_dict, response):
self.calls += 1
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
self.calls += 1
async for chunk in response:
yield chunk
@pytest.mark.asyncio
async def test_proxy_logging_pre_call_hook_load_balancing():
@pytest.mark.parametrize("hook", ("pre_call", "during_call", "post_call", "streaming", "iterator"))
@pytest.mark.parametrize("deployment_count", (0, 1, 2))
async def test_proxy_logging_hooks_load_balance_once_per_request(hook: str, deployment_count: int):
"""Test that async_pre_call_hook load balances across multiple guardrails."""
MockGuardrail.call_count = 0
@ -62,33 +80,46 @@ async def test_proxy_logging_pre_call_hook_load_balancing():
"litellm_params": {"model": "gpt-5.5", "api_key": "fake-key"},
}
],
guardrail_list=guardrail_list,
guardrail_list=guardrail_list[:deployment_count],
)
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [guardrail_1]
litellm.callbacks = [guardrail_1, guardrail_2]
try:
with patch("litellm.proxy.proxy_server.llm_router", router):
with patch("litellm.proxy.proxy_server.llm_router", router if deployment_count else None):
for _ in range(50):
await proxy_logging.pre_call_hook(
user_api_key_dict=MagicMock(),
data={"messages": [{"role": "user", "content": "test"}]},
call_type="completion",
)
previous_calls: Final = guardrail_1.calls + guardrail_2.calls
data: Final = {"messages": [{"role": "user", "content": "test"}]}
auth: Final = MagicMock()
response: Final = litellm.ModelResponse(choices=[{"message": {"content": "test", "role": "assistant"}}])
if hook == "pre_call":
await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="completion")
elif hook == "during_call":
await proxy_logging.during_call_hook(user_api_key_dict=auth, data=data, call_type="completion")
elif hook == "post_call":
await proxy_logging.post_call_success_hook(user_api_key_dict=auth, data=data, response=response)
elif hook == "streaming":
await proxy_logging.async_post_call_streaming_hook(user_api_key_dict=auth, data=data, response=response)
else:
async def chunks(chunk=response):
yield chunk
async for _chunk in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=auth, request_data=data, response=chunks()
):
pass
expected_calls: Final = 1 if deployment_count == 2 else 2
assert guardrail_1.calls + guardrail_2.calls == previous_calls + expected_calls
assert guardrail_1.calls > 0, "Guardrail 1 should have been called"
assert guardrail_2.calls > 0, "Guardrail 2 should have been called"
total = guardrail_1.calls + guardrail_2.calls
assert total == 50, f"Expected 50 total calls, got {total}"
min_calls = min(guardrail_1.calls, guardrail_2.calls)
assert (
min_calls >= 10
), f"Expected at least 10 calls to each guardrail, got min={min_calls}"
expected_total: Final = 50 if deployment_count == 2 else 100
assert total == expected_total, f"Expected {expected_total} total calls, got {total}"
finally:
litellm.callbacks = original_callbacks