From e66f38e447ef0291a684ade244ba9014762657b3 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 16:40:54 +0000 Subject: [PATCH 01/11] fix(guardrails): select one load-balanced callback per hook Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/guardrails/init_guardrails.py | 9 +- litellm/proxy/utils.py | 84 +++++++++++++------ .../test_guardrail_load_balancing.py | 6 +- .../test_guardrail_load_balancing.py | 61 ++++++++++---- 4 files changed, 112 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/guardrails/init_guardrails.py b/litellm/proxy/guardrails/init_guardrails.py index ecdf99e49d0..a5e522f3779 100644 --- a/litellm/proxy/guardrails/init_guardrails.py +++ b/litellm/proxy/guardrails/init_guardrails.py @@ -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)) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a2928bba287..9117f7417ac 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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 diff --git a/tests/integration/observability/test_guardrail_load_balancing.py b/tests/integration/observability/test_guardrail_load_balancing.py index 585cc7c3cc9..796e11f135c 100644 --- a/tests/integration/observability/test_guardrail_load_balancing.py +++ b/tests/integration/observability/test_guardrail_load_balancing.py @@ -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) diff --git a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py index fd38e9e35db..88f7ef3b1bb 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py +++ b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py @@ -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 From f1c79e7e4b3ed8383cf0059f91d032dce220d140 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 16:48:27 +0000 Subject: [PATCH 02/11] fix(guardrails): preserve cached streaming iterator hook kinds Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 78 +++------------- .../test_guardrail_load_balancing.py | 9 +- .../proxy_logging/test_guardrail_pipeline.py | 88 +------------------ 3 files changed, 22 insertions(+), 153 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9117f7417ac..a8a39b9d7ac 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2580,26 +2580,6 @@ class ProxyLogging: raise HTTPException(status_code=400, detail={"error": response}) return data - def _should_use_guardrail_load_balancing( - self, - guardrail_name: str, - ) -> bool: - """ - Check if load balancing should be used for this guardrail. - - Returns True if the router has multiple deployments for this guardrail name. - """ - from litellm.proxy.proxy_server import llm_router - - if llm_router is None: - return False - router_settings: Final = llm_router.settings() - if not router_settings.is_set("guardrail_list"): - return False - - matching: Final = [g for g in router_settings.guardrail_list if g.get("guardrail_name") == guardrail_name] - return len(matching) > 1 - async def _execute_guardrail_hook( self, callback: "CustomGuardrail", @@ -2657,50 +2637,6 @@ class ProxyLogging: else: raise ValueError(f"Unknown hook_type: {hook_type}") - async def _execute_guardrail_with_load_balancing( - self, - guardrail_name: str, - hook_type: str, - data: dict, - user_api_key_dict: UserAPIKeyAuth | None, - call_type: CallTypesLiteral, - response: LLMResponseTypes | None = None, - ) -> object: - """ - Execute a guardrail using the router's load balancing. - - Args: - guardrail_name: Name of the guardrail - hook_type: One of "pre_call", "during_call", "post_call" - data: Request data - user_api_key_dict: User API key auth - call_type: Type of call - response: Response object (for post_call hooks) - - Returns: - Result from the guardrail execution - """ - from litellm.proxy.proxy_server import llm_router - - if llm_router is None: - raise ValueError("Router not initialized") - - # Select guardrail using router's load balancing - selected_guardrail: Final = llm_router.settings().guardrail(guardrail_name) - - callback: Final[CustomGuardrail | None] = 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, - ) - async def _process_guardrail_callback( self, callback: CustomGuardrail, @@ -4845,7 +4781,19 @@ class ProxyLogging: 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") + ( + cb, + next( + kind + for original, kind in caps.iterator_overrides + if original is cb + or ( + isinstance(original, CustomGuardrail) + and isinstance(cb, CustomGuardrail) + and original.guardrail_name == cb.guardrail_name + ) + ), + ) for cb in selected_iterator_callbacks ) running_overrides: Final = tuple( diff --git a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py index 88f7ef3b1bb..36dbc9c3f8e 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py +++ b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py @@ -48,6 +48,10 @@ class MockGuardrail(CustomGuardrail): yield chunk +class InheritedIteratorGuardrail(MockGuardrail): + pass + + @pytest.mark.asyncio @pytest.mark.parametrize("hook", ("pre_call", "during_call", "post_call", "streaming", "iterator")) @pytest.mark.parametrize("deployment_count", (0, 1, 2)) @@ -56,7 +60,10 @@ async def test_proxy_logging_hooks_load_balance_once_per_request(hook: str, depl MockGuardrail.call_count = 0 guardrail_1 = MockGuardrail(guardrail_name="content-filter", guardrail_id="g1") - guardrail_2 = MockGuardrail(guardrail_name="content-filter", guardrail_id="g2") + second_guardrail_class: Final = ( + InheritedIteratorGuardrail if hook == "iterator" and deployment_count == 2 else MockGuardrail + ) + guardrail_2 = second_guardrail_class(guardrail_name="content-filter", guardrail_id="g2") guardrail_list = [ { diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 5fd40696224..b37b5136b14 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1,7 +1,6 @@ """Pin ProxyLogging guardrail pipeline helpers. -Covers ``_should_use_guardrail_load_balancing``, ``_execute_guardrail_hook``, -``_execute_guardrail_with_load_balancing``, ``_process_guardrail_callback``, +Covers ``_execute_guardrail_hook``, ``_process_guardrail_callback``, ``_process_prompt_template``, ``_process_guardrail_metadata``, ``_maybe_execute_pipelines``, ``_handle_pipeline_result``, ``_run_guardrail_with_metrics``, ``_emit_guardrail_metrics``. @@ -18,8 +17,6 @@ import logging from typing import Any, Callable, Dict, Final, List from unittest.mock import AsyncMock, MagicMock, patch -from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs - import pytest from fastapi import HTTPException @@ -60,39 +57,6 @@ def _clear_caps_cache(): ProxyLogging._callback_capabilities_cache.clear() -# --------------------------------------------------------------------------- -# _should_use_guardrail_load_balancing -# --------------------------------------------------------------------------- - - -def test_should_use_guardrail_load_balancing_truth_table(proxy_logging): - snapshot = {} - router = route_entry_points_to_stubs(MagicMock()) - router.guardrail_list = [{"guardrail_name": "g1"}, {"guardrail_name": "g1"}] - with patch("litellm.proxy.proxy_server.llm_router", router): - snapshot["multiple_deployments"] = proxy_logging._should_use_guardrail_load_balancing("g1") - router.guardrail_list = [{"guardrail_name": "g1"}] - with patch("litellm.proxy.proxy_server.llm_router", router): - snapshot["single_deployment"] = proxy_logging._should_use_guardrail_load_balancing("g1") - with patch("litellm.proxy.proxy_server.llm_router", None): - snapshot["no_router"] = proxy_logging._should_use_guardrail_load_balancing("g1") - router.guardrail_list = [{"guardrail_name": "other"}, {"guardrail_name": "other"}] - with patch("litellm.proxy.proxy_server.llm_router", router): - snapshot["unmatched_name"] = proxy_logging._should_use_guardrail_load_balancing("g1") - assert snapshot == { - "multiple_deployments": True, - "single_deployment": False, - "no_router": False, - "unmatched_name": False, - } - - -def test_should_use_guardrail_load_balancing_error_on_bad_guardrail_list(proxy_logging): - router = route_entry_points_to_stubs(MagicMock()) - router.guardrail_list = "not a list" - with patch("litellm.proxy.proxy_server.llm_router", router): - with pytest.raises((TypeError, AttributeError)): - proxy_logging._should_use_guardrail_load_balancing("g1") # --------------------------------------------------------------------------- @@ -165,53 +129,6 @@ async def test_execute_guardrail_hook_unknown_hook_type_raises(proxy_logging, ma ) -# --------------------------------------------------------------------------- -# _execute_guardrail_with_load_balancing -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_execute_guardrail_with_load_balancing_routes_through_router(proxy_logging, make_user_api_key_auth): - cb = _make_guardrail() - router = route_entry_points_to_stubs(MagicMock()) - router.get_available_guardrail = MagicMock(return_value={"callback": cb}) - with patch("litellm.proxy.proxy_server.llm_router", router): - out = await proxy_logging._execute_guardrail_with_load_balancing( - guardrail_name="g", - hook_type="pre_call", - data={"model": "m"}, - user_api_key_dict=make_user_api_key_auth(), - call_type="completion", - ) - assert out == {"a": 1, "b": 2, "c": 3} - - -@pytest.mark.asyncio -async def test_execute_guardrail_with_load_balancing_router_none_raises(proxy_logging, make_user_api_key_auth): - with patch("litellm.proxy.proxy_server.llm_router", None): - with pytest.raises(ValueError, match="Router not initialized"): - await proxy_logging._execute_guardrail_with_load_balancing( - guardrail_name="g", - hook_type="pre_call", - data={}, - user_api_key_dict=make_user_api_key_auth(), - call_type="completion", - ) - - -@pytest.mark.asyncio -async def test_execute_guardrail_with_load_balancing_no_callback_raises(proxy_logging, make_user_api_key_auth): - router = route_entry_points_to_stubs(MagicMock()) - router.get_available_guardrail = MagicMock(return_value={"callback": None}) - with patch("litellm.proxy.proxy_server.llm_router", router): - with pytest.raises(ValueError, match="No callback found"): - await proxy_logging._execute_guardrail_with_load_balancing( - guardrail_name="g", - hook_type="pre_call", - data={}, - user_api_key_dict=make_user_api_key_auth(), - call_type="completion", - ) # --------------------------------------------------------------------------- @@ -237,7 +154,6 @@ async def test_process_guardrail_callback_skipped_when_should_run_false(proxy_lo async def test_process_guardrail_callback_returns_data_on_success(proxy_logging, make_user_api_key_auth, monkeypatch): cb = _make_guardrail() cb.should_run_guardrail = MagicMock(return_value=True) - proxy_logging._should_use_guardrail_load_balancing = MagicMock(return_value=False) out = await proxy_logging._process_guardrail_callback( callback=cb, data={"model": "m", "messages": [{"role": "user"}], "temperature": 0.1}, @@ -257,7 +173,6 @@ async def test_process_guardrail_callback_enriches_and_reraises_http_exception( detail = {"error": "blocked"} cb.async_pre_call_hook = AsyncMock(side_effect=HTTPException(status_code=400, detail=detail)) cb.event_hook = "pre_call" - proxy_logging._should_use_guardrail_load_balancing = MagicMock(return_value=False) with pytest.raises(HTTPException): await proxy_logging._process_guardrail_callback( @@ -4255,4 +4170,3 @@ async def test_live_pipeline_disconnect_scan_that_cannot_run_records_every_step_ assert seen.get("count") is None assert _post_call_statuses(data) == ["guardrail_failed_to_respond"] assert any("cut short" in message and "RuntimeError" in message for message in _warnings(caplog)) - From c47f530b983657afa77508887f0df60c09cd2c50 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 17:17:03 +0000 Subject: [PATCH 03/11] fix: preserve load-balanced guardrail siblings and iterator kinds Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 41 +++-- .../test_guardrail_load_balancing.py | 142 +++++++++++++++++- 2 files changed, 168 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a8a39b9d7ac..6ef6d909788 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -898,6 +898,7 @@ def _load_balanced_callbacks(callbacks: Sequence[object]) -> tuple[object, ...]: def _load_balanced_callbacks(callbacks: Sequence[object]) -> tuple[object, ...]: + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import llm_router if llm_router is None or not llm_router.settings().is_set("guardrail_list"): @@ -921,7 +922,14 @@ def _load_balanced_callbacks(callbacks: Sequence[object]) -> tuple[object, ...]: target: Final = deployment.get("callback") if target is None: raise ValueError(f"No callback found for guardrail: {callback.guardrail_name}") - return (target,) + deployment_id: Final = deployment.get("id") + siblings: Final = ( + IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_sibling_callbacks.get(deployment_id, ()) + if deployment_id is not None + and IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_custom_guardrail.get(deployment_id) is target + else () + ) + return (target, *siblings) return tuple(chain.from_iterable(selected(index, callback) for index, callback in enumerate(callbacks))) @@ -1898,6 +1906,21 @@ def _overrides_moderation_hook(callback: CustomLogger) -> bool: return _overrides_hook(callback, "async_moderation_hook") +def _iterator_callback_kind(callback: CustomLogger, cached_entries: Sequence[tuple[CustomLogger, str]]) -> str | None: + cached: Final = next((kind for original, kind in cached_entries if original is callback), None) + if cached is not None: + return cached + if _overrides_hook(callback, "async_post_call_streaming_iterator_hook"): + return "override" + if ( + isinstance(callback, CustomGuardrail) + and callback.uses_apply_guardrail_interface() + and not callback.use_native_lifecycle_hooks + ): + return "apply_guardrail" + return None + + _LISTED_MODEL_NAMES: Final = TypeAdapter(tuple[str, ...]) _MCP_TOOL_DESCRIPTION: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) _MCP_TOOL_INPUT_SCHEMA: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) @@ -4780,21 +4803,11 @@ class ProxyLogging: 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, - next( - kind - for original, kind in caps.iterator_overrides - if original is cb - or ( - isinstance(original, CustomGuardrail) - and isinstance(cb, CustomGuardrail) - and original.guardrail_name == cb.guardrail_name - ) - ), - ) + (cb, kind) for cb in selected_iterator_callbacks + if (kind := _iterator_callback_kind(cb, caps.iterator_overrides)) is not None ) running_overrides: Final = tuple( (callback, kind) diff --git a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py index 36dbc9c3f8e..8fe68b7f806 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py +++ b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py @@ -4,14 +4,24 @@ Test guardrail load balancing through the Router and ProxyLogging. from unittest.mock import MagicMock, patch, AsyncMock from typing import Final +from collections.abc import AsyncIterator +from queue import SimpleQueue +import httpx import litellm import pytest from litellm import Router from litellm.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.utils import ProxyLogging +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails import guardrail_registry +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler +from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 +from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPI +from litellm.types.utils import ModelResponseStream class MockGuardrail(CustomGuardrail): @@ -109,8 +119,11 @@ async def test_proxy_logging_hooks_load_balance_once_per_request(hook: str, depl 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) + 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 @@ -130,3 +143,130 @@ async def test_proxy_logging_hooks_load_balance_once_per_request(hook: str, depl finally: litellm.callbacks = original_callbacks + + +async def _output_chunks() -> AsyncIterator[ModelResponseStream]: + yield ModelResponseStream( + model="test-model", + created=1700000000, + choices=[{"index": 0, "delta": {"role": "assistant", "content": "raw"}, "finish_reason": "stop"}], + ) + + +def _chunk_text(chunk: object) -> str: + assert isinstance(chunk, ModelResponseStream) + return chunk.choices[0].delta.content or "" + + +async def _guarded_output(logging: ProxyLogging, streaming: bool) -> str: + data: Final = {"messages": [{"role": "user", "content": "raw"}]} + auth: Final = UserAPIKeyAuth(request_route="/v1/chat/completions") + if streaming: + return "".join( + [ + _chunk_text(chunk) + async for chunk in logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, request_data=data, response=_output_chunks() + ) + ] + ) + response: Final = await logging.post_call_success_hook( + user_api_key_dict=auth, + data=data, + response=litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "raw"}}]), + ) + assert isinstance(response, litellm.ModelResponse) + choice: Final = response.choices[0] + assert isinstance(choice, litellm.Choices) + content: Final = choice.message.content + assert isinstance(content, str) + return content + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", (False, True)) +async def test_presidio_load_balancing_keeps_selected_output_sibling( + streaming: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + router: Final = Router(model_list=[]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(guardrail_registry, "IN_MEMORY_GUARDRAIL_HANDLER", InMemoryGuardrailHandler()) + monkeypatch.setattr(litellm, "callbacks", []) + init_guardrails_v2( + all_guardrails=[ + { + "guardrail_name": "presidio-balanced", + "litellm_params": { + "guardrail": "presidio", + "mode": "pre_call", + "default_on": True, + "presidio_filter_scope": "both", + "presidio_analyzer_api_base": "https://presidio.invalid/analyze", + "presidio_anonymizer_api_base": "https://presidio.invalid/anonymize", + "mock_redacted_text": {"text": masked, "items": []}, + }, + } + for masked in ("masked-east", "masked-west") + ], + llm_router=router, + ) + logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + outputs: Final[SimpleQueue[str]] = SimpleQueue() + for _ in range(50): + content: Final = await _guarded_output(logging, streaming) + assert content in ("masked-east", "masked-west") + outputs.put(content) + assert frozenset(outputs.get() for _ in range(outputs.qsize())) == {"masked-east", "masked-west"} + + +@pytest.mark.asyncio +async def test_iterator_uses_selected_generic_deployments_cached_kind(monkeypatch: pytest.MonkeyPatch) -> None: + native: Final = MockGuardrail("mixed-iterator", "native") + requests: Final[SimpleQueue[bytes]] = SimpleQueue() + + async def answer(request: httpx.Request) -> httpx.Response: + requests.put(request.content) + return httpx.Response(200, json={"action": "BLOCKED", "blocked_reason": "test policy"}) + + handler: Final = AsyncHTTPHandler(transport=httpx.MockTransport(answer)) + generic: Final = GenericGuardrailAPI( + guardrail_name="mixed-iterator", + api_base="https://guardrail.invalid", + event_hook="post_call", + default_on=True, + async_handler=handler, + streaming_sampling_rate=1, + ) + router: Final = Router( + model_list=[], + guardrail_list=[ + { + "guardrail_name": "mixed-iterator", + "callback": cb, + "id": deployment_id, + "litellm_params": {"guardrail": "custom", "mode": "post_call"}, + } + for cb, deployment_id in ((native, "native"), (generic, "generic")) + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", [native, generic]) + logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + blocked: Final[SimpleQueue[litellm.exceptions.GuardrailRaisedException]] = SimpleQueue() + try: + for _ in range(50): + try: + async for _chunk in logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"), + request_data={"messages": [{"role": "user", "content": "hi"}]}, + response=_output_chunks(), + ): + pass + except litellm.exceptions.GuardrailRaisedException as exc: + blocked.put(exc) + assert blocked.qsize() > 0 + assert requests.qsize() == blocked.qsize() + assert native.calls + blocked.qsize() == 50 + assert native.calls > 0 + finally: + await handler.close() From 1fea7d76098275894b2e941325cbdcfbb27907b2 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 18:31:29 +0000 Subject: [PATCH 04/11] fix guardrail selection to consider hook eligibility Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 169 ++++++++++++------ .../test_guardrail_load_balancing.py | 123 +++++++++++-- 2 files changed, 231 insertions(+), 61 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 6ef6d909788..10e1df2c870 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -326,6 +326,7 @@ if TYPE_CHECKING: from litellm.repositories.prisma_protocols import TableActions from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep + from litellm.types.router import GuardrailTypedDict Span = _Span | object else: @@ -889,21 +890,69 @@ 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, ...]: ... +_GUARDRAIL_ELIGIBILITY_DATA: Final = TypeAdapter(Mapping[str, object]) -@overload -def _load_balanced_callbacks(callbacks: Sequence[object]) -> tuple[object, ...]: ... +def _during_call_guardrail_eligible( + callback: CustomGuardrail, data: Mapping[str, object], call_type: CallTypesLiteral +) -> bool: + if callback.event_hook is None and hasattr(callback, "moderation_check"): + return getattr(callback, "moderation_check", None) != "pre_call" + event_type: Final = ( + GuardrailEventHooks.during_mcp_call + if call_type == CallTypes.call_mcp_tool.value + else GuardrailEventHooks.during_call + ) + return callback.should_run_guardrail(data=data, event_type=event_type) is True -def _load_balanced_callbacks(callbacks: Sequence[object]) -> tuple[object, ...]: +def _post_call_guardrail_eligible( + callback: CustomGuardrail, data: Mapping[str, object], excluded_names: frozenset[str] +) -> bool: + return ( + callback.guardrail_name not in excluded_names + and callback.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is True + ) + + +def _deployment_callbacks(deployment: "GuardrailTypedDict") -> tuple[CustomGuardrail, ...]: from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + + target: Final = deployment.get("callback") + if target is None: + raise ValueError(f"No callback found for guardrail: {deployment['guardrail_name']}") + deployment_id: Final = deployment.get("id") + siblings: Final = ( + IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_sibling_callbacks.get(deployment_id, ()) + if deployment_id is not None + and IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_custom_guardrail.get(deployment_id) is target + else () + ) + return (target, *siblings) + + +@overload +def _load_balanced_callbacks( + callbacks: Sequence[CustomLogger], eligible: Callable[[CustomGuardrail], bool] +) -> tuple[CustomLogger, ...]: ... + + +@overload +def _load_balanced_callbacks( + callbacks: Sequence[object], eligible: Callable[[CustomGuardrail], bool] +) -> tuple[object, ...]: ... + + +def _load_balanced_callbacks( + callbacks: Sequence[object], eligible: Callable[[CustomGuardrail], bool] +) -> tuple[object, ...]: from litellm.proxy.proxy_server import llm_router + from litellm.router_strategy.simple_shuffle import simple_shuffle if llm_router is None or not llm_router.settings().is_set("guardrail_list"): return tuple(callbacks) settings: Final = llm_router.settings() + registry: Final = llm_router.registry() balanced_names: Final = frozenset( entry["guardrail_name"] for entry in settings.guardrail_list @@ -918,23 +967,28 @@ def _load_balanced_callbacks(callbacks: Sequence[object]) -> tuple[object, ...]: 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}") - deployment_id: Final = deployment.get("id") - siblings: Final = ( - IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_sibling_callbacks.get(deployment_id, ()) - if deployment_id is not None - and IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_custom_guardrail.get(deployment_id) is target - else () + candidates: Final = tuple( + entry + for entry in settings.guardrail_list + if entry["guardrail_name"] == callback.guardrail_name + and any(eligible(member) for member in _deployment_callbacks(entry)) ) - return (target, *siblings) + if not candidates: + return () + deployment: Final = simple_shuffle( + resolve_model_alias=lambda name: registry.resolve(name, kind="alias"), + healthy_deployments=candidates, + model=callback.guardrail_name, + request_kwargs=None, + ) + return _deployment_callbacks(deployment) return tuple(chain.from_iterable(selected(index, callback) for index, callback in enumerate(callbacks))) -def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]: +def _partition_post_call_callbacks( + eligible: Callable[[CustomGuardrail], bool], +) -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]: resolved: Final = tuple( litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( cast( # cast-ok: the resolver returns None for unknown names, filtered below @@ -945,7 +999,9 @@ def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple else callback for callback in litellm.callbacks ) - present: Final = _load_balanced_callbacks(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), eligible + ) 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, ...]", @@ -3251,7 +3307,22 @@ class ProxyLogging: return data selected_callbacks: Final = ( - caps.resolved_callbacks if skip_guardrails else _load_balanced_callbacks(caps.resolved_callbacks) + caps.resolved_callbacks + if skip_guardrails + else _load_balanced_callbacks( + caps.resolved_callbacks, + lambda cb: ( + data is not None + and cb.guardrail_name not in pipeline_managed + and cb.should_run_guardrail( + data=data, + event_type=GuardrailEventHooks.pre_mcp_call + if call_type in MCP_GUARDRAIL_CALL_TYPES + else GuardrailEventHooks.pre_call, + ) + is True + ), + ) ) parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = ( () @@ -3737,8 +3808,12 @@ class ProxyLogging: return data # Step 1: Collect all guardrail tasks to run in parallel guardrail_tasks: Final = [] + eligibility_data: Final = _GUARDRAIL_ELIGIBILITY_DATA.validate_python(data) - for callback in _load_balanced_callbacks(tuple(litellm.callbacks)): + for callback in _load_balanced_callbacks( + tuple(litellm.callbacks), + lambda cb: _during_call_guardrail_eligible(cb, eligibility_data, call_type), + ): if ( isinstance(callback, CustomLogger) and not isinstance(callback, CustomGuardrail) @@ -3758,20 +3833,8 @@ class ProxyLogging: # Check if guardrail should be run for GuardrailEventHooks.during_call hook ################################################################ - # V1 implementation - backwards compatibility - if callback.event_hook is None and hasattr(callback, "moderation_check"): - if callback.moderation_check == "pre_call": - continue - else: - # Main - V2 Guardrails implementation - from litellm.types.guardrails import GuardrailEventHooks - - event_type = GuardrailEventHooks.during_call - if call_type == CallTypes.call_mcp_tool.value: - event_type = GuardrailEventHooks.during_mcp_call - - if callback.should_run_guardrail(data=data, event_type=event_type) is not True: - continue + if not _during_call_guardrail_eligible(callback, eligibility_data, call_type): + continue # Convert user_api_key_dict to proper format for async_moderation_hook if call_type == CallTypes.call_mcp_tool.value: user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict(user_api_key_dict) @@ -4377,10 +4440,13 @@ class ProxyLogging: response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call") - guardrail_callbacks, other_callbacks = _partition_post_call_callbacks() try: # Merge model-level guardrails before checking which guardrails to run guardrail_data: Final = check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) + eligibility_data: Final = _GUARDRAIL_ELIGIBILITY_DATA.validate_python(guardrail_data) + guardrail_callbacks, other_callbacks = _partition_post_call_callbacks( + lambda cb: _post_call_guardrail_eligible(cb, eligibility_data, pipeline_managed) + ) parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple( callback @@ -4706,28 +4772,24 @@ class ProxyLogging: if response_str is not None: # Cache model-level guardrails check per-request to avoid repeated # dict lookups + llm_router.get_deployment() per callback per chunk. - _cached_guardrail_data: dict | None = None - _guardrail_data_computed = False - pipeline_gated: Final = ( + _cached_guardrail_data: Final = _GUARDRAIL_ELIGIBILITY_DATA.validate_python( + check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) + if caps.has_guardrail + else data, + ) + pipeline_gated: Final[frozenset[str]] = ( stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset() ) - for callback in _load_balanced_callbacks(tuple(litellm.callbacks)): + for callback in _load_balanced_callbacks( + tuple(litellm.callbacks), + lambda cb: _post_call_guardrail_eligible(cb, _cached_guardrail_data, pipeline_gated), + ): try: _callback: CustomLogger | None = None if isinstance(callback, CustomGuardrail): if callback.guardrail_name in pipeline_gated: continue - # Main - V2 Guardrails implementation - from litellm.types.guardrails import GuardrailEventHooks - - ## CHECK FOR MODEL-LEVEL GUARDRAILS (cached per-request) - if not _guardrail_data_computed: - _cached_guardrail_data = check_and_merge_model_level_guardrails( - data=data, llm_router=llm_router - ) - _guardrail_data_computed = True - if ( callback.should_run_guardrail( data=_cached_guardrail_data, @@ -4802,7 +4864,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)) + eligibility_data: Final = _GUARDRAIL_ELIGIBILITY_DATA.validate_python(request_data) + selected_iterator_callbacks: Final = _load_balanced_callbacks( + tuple(cb for cb, _ in caps.iterator_overrides), + lambda cb: ( + _iterator_callback_kind(cb, caps.iterator_overrides) is not None + and _post_call_guardrail_eligible(cb, eligibility_data, pipeline_gated_names) + ), + ) selected_iterator_overrides: Final = tuple( (cb, kind) diff --git a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py index 8fe68b7f806..7b110cec59c 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py +++ b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py @@ -2,25 +2,26 @@ Test guardrail load balancing through the Router and ProxyLogging. """ -from unittest.mock import MagicMock, patch, AsyncMock -from typing import Final from collections.abc import AsyncIterator from queue import SimpleQueue - +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch import httpx -import litellm import pytest + +import litellm from litellm import Router from litellm.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.proxy.utils import ProxyLogging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails import guardrail_registry +from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPI from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 -from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPI +from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks, GuardrailStreamScope from litellm.types.utils import ModelResponseStream @@ -29,14 +30,20 @@ class MockGuardrail(CustomGuardrail): call_count = 0 - def __init__(self, guardrail_name: str, guardrail_id: str): - super().__init__(guardrail_name=guardrail_name) + def __init__( + self, + guardrail_name: str, + guardrail_id: str, + event_hook: GuardrailEventHooks | None = None, + default_on: bool = False, + stream_scope: GuardrailStreamScope = "both", + ): + super().__init__( + guardrail_name=guardrail_name, event_hook=event_hook, default_on=default_on, stream_scope=stream_scope + ) self.guardrail_id = guardrail_id self.calls = 0 - def should_run_guardrail(self, data, event_type) -> bool: - return True - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): self.calls += 1 MockGuardrail.call_count += 1 @@ -62,6 +69,10 @@ class InheritedIteratorGuardrail(MockGuardrail): pass +class NonIteratorGuardrail(CustomGuardrail): + pass + + @pytest.mark.asyncio @pytest.mark.parametrize("hook", ("pre_call", "during_call", "post_call", "streaming", "iterator")) @pytest.mark.parametrize("deployment_count", (0, 1, 2)) @@ -270,3 +281,93 @@ async def test_iterator_uses_selected_generic_deployments_cached_kind(monkeypatc assert native.calls > 0 finally: await handler.close() + + +async def _dispatch_scoped_hook(logging: ProxyLogging, hook: str) -> None: + auth: Final = UserAPIKeyAuth(request_route="/v1/chat/completions") + data: Final = {"messages": [{"role": "user", "content": "raw"}], "stream": hook in ("streaming", "iterator")} + response: Final = litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "raw"}}]) + if hook == "pre_call": + await logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="completion") + elif hook == "during_call": + await logging.during_call_hook(user_api_key_dict=auth, data=data, call_type="completion") + elif hook == "post_call": + await logging.post_call_success_hook(user_api_key_dict=auth, data=data, response=response) + elif hook == "streaming": + await logging.async_post_call_streaming_hook(user_api_key_dict=auth, data=data, response=response) + else: + async for _chunk in logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, request_data=data, response=_output_chunks() + ): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hook", ("pre_call", "during_call", "post_call", "streaming", "iterator")) +@pytest.mark.parametrize("mismatch", ("mode", "default_on", "stream_scope", "none_eligible", "weights")) +async def test_load_balancing_selects_only_hook_eligible_deployments( + hook: str, mismatch: str, monkeypatch: pytest.MonkeyPatch +) -> None: + event: Final = GuardrailEventHooks(hook if hook in ("pre_call", "during_call") else "post_call") + other_event: Final = ( + GuardrailEventHooks.post_call if event == GuardrailEventHooks.pre_call else GuardrailEventHooks.pre_call + ) + eligible: Final = MockGuardrail("scoped-lb", "eligible", event_hook=event, default_on=mismatch != "none_eligible") + ineligible: Final = MockGuardrail( + "scoped-lb", + "other", + event_hook=other_event if mismatch == "mode" else event, + default_on=mismatch not in ("default_on", "none_eligible"), + stream_scope=("non_streaming" if hook in ("streaming", "iterator") else "streaming") + if mismatch == "stream_scope" + else "both", + ) + router: Final = Router( + model_list=[], + guardrail_list=[ + { + "guardrail_name": "scoped-lb", + "id": str(index), + "callback": cb, + "litellm_params": {"guardrail": "custom", "mode": "post_call", "weight": weight}, + } + for index, (cb, weight) in enumerate( + ((eligible, 1 if mismatch == "weights" else 0), (ineligible, 0 if mismatch == "weights" else 1)) + ) + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", [ineligible, eligible]) + logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + for request_index in range(20): + await _dispatch_scoped_hook(logging, hook) + assert eligible.calls == (0 if mismatch == "none_eligible" else request_index + 1) + assert ineligible.calls == 0 + + +@pytest.mark.asyncio +async def test_streaming_selects_iterator_capable_deployment(monkeypatch: pytest.MonkeyPatch) -> None: + eligible: Final = MockGuardrail( + "iterator-lb", "iterator", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + unsupported: Final = NonIteratorGuardrail( + guardrail_name="iterator-lb", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + router: Final = Router( + model_list=[], + guardrail_list=[ + { + "guardrail_name": "iterator-lb", + "id": str(index), + "callback": cb, + "litellm_params": {"guardrail": "custom", "mode": "post_call", "weight": weight}, + } + for index, (cb, weight) in enumerate(((eligible, 0), (unsupported, 1))) + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", [unsupported, eligible]) + logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + for request_index in range(20): + await _dispatch_scoped_hook(logging, "iterator") + assert eligible.calls == request_index + 1 From 513e80dadfc418047c8304e1e23d36274cbe0ebb Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 19:06:39 +0000 Subject: [PATCH 05/11] fix: defer pre-call guardrail selection until dispatch Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 54 +++++++++----- .../test_guardrail_load_balancing.py | 70 ++++++++++++++++++- 2 files changed, 103 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 10e1df2c870..44d0cde59b9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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 chain, takewhile +from itertools import chain, takewhile, tee from operator import attrgetter from types import MappingProxyType from typing import ( @@ -946,11 +946,30 @@ def _load_balanced_callbacks( def _load_balanced_callbacks( callbacks: Sequence[object], eligible: Callable[[CustomGuardrail], bool] ) -> tuple[object, ...]: + return tuple(_iter_load_balanced_callbacks(callbacks, eligible)) + + +@overload +def _iter_load_balanced_callbacks( + callbacks: Sequence[CustomLogger], eligible: Callable[[CustomGuardrail], bool] +) -> Iterator[CustomLogger]: ... + + +@overload +def _iter_load_balanced_callbacks( + callbacks: Sequence[object], eligible: Callable[[CustomGuardrail], bool] +) -> Iterator[object]: ... + + +def _iter_load_balanced_callbacks( + callbacks: Sequence[object], eligible: Callable[[CustomGuardrail], bool] +) -> Iterator[object]: from litellm.proxy.proxy_server import llm_router from litellm.router_strategy.simple_shuffle import simple_shuffle if llm_router is None or not llm_router.settings().is_set("guardrail_list"): - return tuple(callbacks) + yield from callbacks + return settings: Final = llm_router.settings() registry: Final = llm_router.registry() balanced_names: Final = frozenset( @@ -983,7 +1002,7 @@ def _load_balanced_callbacks( ) return _deployment_callbacks(deployment) - return tuple(chain.from_iterable(selected(index, callback) for index, callback in enumerate(callbacks))) + yield from chain.from_iterable(selected(index, callback) for index, callback in enumerate(callbacks)) def _partition_post_call_callbacks( @@ -3306,10 +3325,10 @@ class ProxyLogging: self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails) return data - selected_callbacks: Final = ( + callback_iterators: Final = tee( caps.resolved_callbacks if skip_guardrails - else _load_balanced_callbacks( + else _iter_load_balanced_callbacks( caps.resolved_callbacks, lambda cb: ( data is not None @@ -3324,20 +3343,8 @@ class ProxyLogging: ), ) ) - parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = ( - () - if skip_guardrails - else tuple( - cb - 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) - ) - ) - deferred_route_exc: SensitiveDataRouteException | None = None - for _callback in selected_callbacks: + for _callback in callback_iterators[0]: start_time = time.time() try: if isinstance(_callback, CustomGuardrail) and data is not None: @@ -3415,6 +3422,17 @@ class ProxyLogging: if deferred_route_exc is not None and data is not None: data = await self._handle_sensitive_data_route_exception(deferred_route_exc, data, user_api_key_dict) + parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = ( + () + if skip_guardrails + else tuple( + cb + for cb in callback_iterators[1] + if isinstance(cb, CustomGuardrail) + and getattr(cb, "run_in_parallel", False) + and not (cb.guardrail_name and cb.guardrail_name in pipeline_managed) + ) + ) if parallel_guardrails and data is not None: await self._run_parallel_pre_call_guardrails( guardrails=parallel_guardrails, diff --git a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py index 7b110cec59c..11810569dff 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py +++ b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py @@ -2,7 +2,7 @@ Test guardrail load balancing through the Router and ProxyLogging. """ -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from queue import SimpleQueue from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -14,6 +14,7 @@ import litellm from litellm import Router from litellm.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails import guardrail_registry @@ -22,7 +23,7 @@ from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks, GuardrailStreamScope -from litellm.types.utils import ModelResponseStream +from litellm.types.utils import CallTypesLiteral, ModelResponseStream class MockGuardrail(CustomGuardrail): @@ -37,9 +38,14 @@ class MockGuardrail(CustomGuardrail): event_hook: GuardrailEventHooks | None = None, default_on: bool = False, stream_scope: GuardrailStreamScope = "both", + run_in_parallel: bool = False, ): super().__init__( - guardrail_name=guardrail_name, event_hook=event_hook, default_on=default_on, stream_scope=stream_scope + guardrail_name=guardrail_name, + event_hook=event_hook, + default_on=default_on, + stream_scope=stream_scope, + run_in_parallel=run_in_parallel, ) self.guardrail_id = guardrail_id self.calls = 0 @@ -73,6 +79,23 @@ class NonIteratorGuardrail(CustomGuardrail): pass +class GuardrailActivator(CustomLogger): + def __init__(self, action: str): + super().__init__() + self.action = action + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: Mapping[str, object], + call_type: CallTypesLiteral, + ) -> Mapping[str, object]: + if self.action == "raise": + raise ValueError("activation blocked") + return {**data, "guardrails": ["dynamic-lb"]} if self.action == "enable" else data + + @pytest.mark.asyncio @pytest.mark.parametrize("hook", ("pre_call", "during_call", "post_call", "streaming", "iterator")) @pytest.mark.parametrize("deployment_count", (0, 1, 2)) @@ -371,3 +394,44 @@ async def test_streaming_selects_iterator_capable_deployment(monkeypatch: pytest for request_index in range(20): await _dispatch_scoped_hook(logging, "iterator") assert eligible.calls == request_index + 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ("enable", "noop", "raise")) +@pytest.mark.parametrize("run_in_parallel", (False, True)) +@pytest.mark.parametrize("call_type", ("completion", "call_mcp_tool")) +async def test_pre_call_load_balancing_observes_earlier_callback_activation( + action: str, run_in_parallel: bool, call_type: CallTypesLiteral, monkeypatch: pytest.MonkeyPatch +) -> None: + event: Final = GuardrailEventHooks.pre_mcp_call if call_type == "call_mcp_tool" else GuardrailEventHooks.pre_call + replicas: Final = tuple( + MockGuardrail("dynamic-lb", region, event_hook=event, default_on=False, run_in_parallel=run_in_parallel) + for region in ("east", "west") + ) + router: Final = Router( + model_list=[], + guardrail_list=[ + { + "guardrail_name": "dynamic-lb", + "id": cb.guardrail_id, + "callback": cb, + "litellm_params": {"guardrail": "custom", "mode": event.value}, + } + for cb in replicas + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", [GuardrailActivator(action), *replicas]) + logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + for request_index in range(30): + if action == "raise": + with pytest.raises(ValueError, match="activation blocked"): + await logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), data={"model": "test", "messages": []}, call_type=call_type + ) + else: + await logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), data={"model": "test", "messages": []}, call_type=call_type + ) + assert sum(cb.calls for cb in replicas) == (request_index + 1 if action == "enable" else 0) + assert all(cb.calls > 0 for cb in replicas) is (action == "enable") From a30222c2c43a03de79232e1174e84c4fef1c8f31 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 19:39:18 +0000 Subject: [PATCH 06/11] fix: preserve raw scans and skip unavailable guardrail replicas Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 14 ++- .../test_guardrail_load_balancing.py | 91 +++++++++++++++++++ 2 files changed, 101 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 44d0cde59b9..468130c25a2 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -986,11 +986,15 @@ def _iter_load_balanced_callbacks( for previous in callbacks[:index] ): return () + deployments: Final = tuple( + entry for entry in settings.guardrail_list if entry["guardrail_name"] == callback.guardrail_name + ) + if not any(entry.get("callback") is not None for entry in deployments): + raise ValueError(f"No callback found for guardrail: {callback.guardrail_name}") candidates: Final = tuple( entry - for entry in settings.guardrail_list - if entry["guardrail_name"] == callback.guardrail_name - and any(eligible(member) for member in _deployment_callbacks(entry)) + for entry in deployments + if entry.get("callback") is not None and any(eligible(member) for member in _deployment_callbacks(entry)) ) if not candidates: return () @@ -3334,7 +3338,9 @@ class ProxyLogging: data is not None and cb.guardrail_name not in pipeline_managed and cb.should_run_guardrail( - data=data, + data=raw_request_snapshot + if cb.scan_raw_request and raw_request_snapshot is not None + else data, event_type=GuardrailEventHooks.pre_mcp_call if call_type in MCP_GUARDRAIL_CALL_TYPES else GuardrailEventHooks.pre_call, diff --git a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py index 11810569dff..23831e074da 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py +++ b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py @@ -39,6 +39,7 @@ class MockGuardrail(CustomGuardrail): default_on: bool = False, stream_scope: GuardrailStreamScope = "both", run_in_parallel: bool = False, + scan_raw_request: bool = False, ): super().__init__( guardrail_name=guardrail_name, @@ -46,6 +47,7 @@ class MockGuardrail(CustomGuardrail): default_on=default_on, stream_scope=stream_scope, run_in_parallel=run_in_parallel, + scan_raw_request=scan_raw_request, ) self.guardrail_id = guardrail_id self.calls = 0 @@ -96,6 +98,17 @@ class GuardrailActivator(CustomLogger): return {**data, "guardrails": ["dynamic-lb"]} if self.action == "enable" else data +class GuardrailRemover(CustomLogger): + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: Mapping[str, object], + call_type: CallTypesLiteral, + ) -> Mapping[str, object]: + return {key: value for key, value in data.items() if key != "guardrails"} + + @pytest.mark.asyncio @pytest.mark.parametrize("hook", ("pre_call", "during_call", "post_call", "streaming", "iterator")) @pytest.mark.parametrize("deployment_count", (0, 1, 2)) @@ -435,3 +448,81 @@ async def test_pre_call_load_balancing_observes_earlier_callback_activation( ) assert sum(cb.calls for cb in replicas) == (request_index + 1 if action == "enable" else 0) assert all(cb.calls > 0 for cb in replicas) is (action == "enable") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("requested", (False, True)) +@pytest.mark.parametrize("run_in_parallel", (False, True)) +async def test_raw_request_scans_survive_earlier_callback_removing_guardrail_names( + requested: bool, run_in_parallel: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + single: Final = MockGuardrail( + "single-raw", + "single", + event_hook=GuardrailEventHooks.pre_call, + scan_raw_request=True, + run_in_parallel=run_in_parallel, + ) + replicas: Final = tuple( + MockGuardrail( + "balanced-raw", + region, + event_hook=GuardrailEventHooks.pre_call, + scan_raw_request=True, + run_in_parallel=run_in_parallel, + ) + for region in ("east", "west") + ) + router: Final = Router( + model_list=[], + guardrail_list=[ + { + "guardrail_name": cb.guardrail_name, + "id": cb.guardrail_id, + "callback": cb, + "litellm_params": {"guardrail": "custom", "mode": "pre_call"}, + } + for cb in (single, *replicas) + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", [GuardrailRemover(), single, *replicas]) + logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + await logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + data={"model": "test", "messages": [], "guardrails": ["single-raw", "balanced-raw"] if requested else []}, + call_type="completion", + ) + assert single.calls == int(requested) + assert (sum(cb.calls for cb in replicas) > 0) is requested + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hook", ("pre_call", "during_call", "post_call", "streaming", "iterator")) +@pytest.mark.parametrize("healthy", (False, True)) +async def test_missing_callback_does_not_poison_healthy_replica( + hook: str, healthy: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + guardrail: Final = MockGuardrail("partial-lb", "healthy", default_on=True) + router: Final = Router( + model_list=[], + guardrail_list=[ + {"guardrail_name": "partial-lb", "id": "missing", "litellm_params": {"guardrail": "custom", "weight": 0}}, + { + "guardrail_name": "partial-lb", + "id": "healthy", + "callback": guardrail if healthy else None, + "litellm_params": {"guardrail": "custom", "weight": 1}, + }, + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + if healthy: + await _dispatch_scoped_hook(logging, hook) + assert guardrail.calls == 1 + else: + with pytest.raises(ValueError, match="No callback found for guardrail: partial-lb"): + await _dispatch_scoped_hook(logging, hook) + assert guardrail.calls == 0 From 4e0a6d6e43c7fde01cab3627c293567fd651c290 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 20:38:13 +0000 Subject: [PATCH 07/11] fix(guardrails): balance pipeline steps and passthrough relay Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/policy_engine/pipeline_executor.py | 105 +++++++- litellm/proxy/utils.py | 84 ++++-- .../test_guardrail_load_balancing.py | 245 +++++++++++++++++- 3 files changed, 397 insertions(+), 37 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index b9a60d9e37f..e57143a9422 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -53,6 +53,7 @@ except ImportError: HTTPException = None _LEGACY_HOOK_REPLACEMENT: Final = TypeAdapter(object) +_PIPELINE_DATA_FIELDS: Final = TypeAdapter(dict[str, object]) class UndeliverableStreamRewrite(Exception): @@ -329,6 +330,30 @@ def _prepare_hook_input( return hook_input, scans_raw_request +def _pipeline_callback_eligible( + callback: CustomGuardrail, + step: PipelineStep, + data: Mapping[str, object], + raw_request_snapshot: Mapping[str, object] | None, + mode: str, + streaming_chunks: list[object] | None, # mutable-ok: shared buffered-stream chunks, read only here +) -> bool: + source: Final = raw_request_snapshot if callback.scan_raw_request and raw_request_snapshot is not None else data + metadata: Final = source.get("metadata") + hook_input: Final[Mapping[str, object]] = { + **source, + "metadata": { + **(_PIPELINE_DATA_FIELDS.validate_python(metadata) if isinstance(metadata, Mapping) else {}), + "guardrails": [step.guardrail], + }, + **({"stream": True} if streaming_chunks is not None else {}), + } + event_type: Final = _PIPELINE_EVENT_HOOKS.get(mode) + return (event_type is None or callback.should_run_guardrail(data=hook_input, event_type=event_type) is True) and ( # pyright: ignore[reportUnknownMemberType] # legacy method accepts the str-keyed request payload + streaming_chunks is None or PipelineExecutor.supports_streaming_execution(callback) + ) + + async def _run_pre_call_hook( target: CustomLogger, callback: CustomGuardrail, @@ -434,7 +459,7 @@ class PipelineExecutor: metadata_key: Final = get_metadata_variable_name_from_kwargs( cast(Mapping[str, object], data) # cast-ok: the request payload is str-keyed ) - working_data = data.copy() + working_data = _PIPELINE_DATA_FIELDS.validate_python(data) if "metadata" in working_data: working_data["metadata"] = working_data["metadata"].copy() @@ -453,7 +478,9 @@ class PipelineExecutor: user_api_key_dict=user_api_key_dict, call_type=call_type, metadata_key=metadata_key, - raw_request_snapshot=raw_request_snapshot, + raw_request_snapshot=_PIPELINE_DATA_FIELDS.validate_python(raw_request_snapshot) + if raw_request_snapshot is not None + else None, streaming_chunks=streaming_chunks, endpoint_translation=endpoint_translation, stream_already_sent=stream_already_sent, @@ -580,6 +607,66 @@ class PipelineExecutor: @staticmethod async def _run_step( + step: PipelineStep, + mode: str, + data: dict[str, object], # mutable-ok: request payload passed to the guardrail hooks + user_api_key_dict: "UserAPIKeyAuth", + call_type: str, + metadata_key: Literal["metadata", "litellm_metadata"], + raw_request_snapshot: dict[str, object] | None = None, # mutable-ok: original request payload + streaming_chunks: list[object] | None = None, # mutable-ok: shared stream buffer rewritten by hooks + endpoint_translation: "BaseTranslation | None" = None, + stream_already_sent: bool = False, + ) -> tuple[ + Literal["pass", "fail", "error", "skip"], + dict[str, object] | None, + str | None, + Exception | None, + ]: + from litellm.proxy.utils import ( + _load_balanced_callbacks, # pyright: ignore[reportPrivateUsage] # same deployment selection as the proxy hook walks + ) + + callback: Final = PipelineExecutor.find_guardrail_callback(step.guardrail) + if callback is None: + verbose_proxy_logger.warning("Pipeline: guardrail '%s' not found in callbacks", step.guardrail) + return ("error", None, f"Guardrail '{step.guardrail}' not found", None) + callbacks: Final = _load_balanced_callbacks( + (callback,), + lambda cb: _pipeline_callback_eligible(cb, step, data, raw_request_snapshot, mode, streaming_chunks), + eligible_only=True, + ) + working_data = data + modified_data: dict[str, object] | None = None # mutable-ok: merged hook payload returned to the pipeline + outcome: Literal["pass", "skip"] = "skip" + for selected in callbacks: + if not isinstance(selected, CustomGuardrail): + continue + result = await PipelineExecutor._run_callback_step( + callback=selected, + step=step, + mode=mode, + data=working_data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + metadata_key=metadata_key, + raw_request_snapshot=raw_request_snapshot, + streaming_chunks=streaming_chunks, + endpoint_translation=endpoint_translation, + stream_already_sent=stream_already_sent, + ) + if result[0] in ("fail", "error"): + return result + if result[0] == "pass": + outcome = "pass" + if result[1] is not None: + modified_data = {**(modified_data or {}), **result[1]} + working_data = {**working_data, **result[1]} + return (outcome, modified_data, None, None) + + @staticmethod + async def _run_callback_step( + callback: CustomGuardrail, step: PipelineStep, mode: str, data: dict, @@ -592,7 +679,7 @@ class PipelineExecutor: stream_already_sent: bool = False, ) -> tuple[ Literal["pass", "fail", "error", "skip"], - dict | None, + dict[str, object] | None, str | None, Exception | None, ]: @@ -608,11 +695,6 @@ class PipelineExecutor: pipeline can re-raise it verbatim and match the direct-attachment response/trace, else None """ - callback: Final = PipelineExecutor.find_guardrail_callback(step.guardrail) - if callback is None: - verbose_proxy_logger.warning("Pipeline: guardrail '%s' not found in callbacks", step.guardrail) - return ("error", None, f"Guardrail '{step.guardrail}' not found", None) - hook_data: Final[Mapping[str, object]] = cast(Mapping[str, object], data) # cast-ok: payload if not _pipeline_stream_scope_allows(callback, hook_data, mode, streaming_chunks): return ("skip", None, None, None) @@ -680,7 +762,12 @@ class PipelineExecutor: None, None, ) - return ("pass", response if isinstance(response, dict) else None, None, None) + return ( + "pass", + _PIPELINE_DATA_FIELDS.validate_python(response) if isinstance(response, dict) else None, + None, + None, + ) except Exception as e: if CustomGuardrail._is_guardrail_intervention(e): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 468130c25a2..8e85753b114 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -701,6 +701,24 @@ def _pipeline_step_guardrail_names(pipelines: Sequence[tuple[str, "GuardrailPipe return frozenset(step.guardrail for step in _pipeline_steps(pipelines)) +def _pipeline_callback_options(guardrail_name: str) -> tuple[CustomGuardrail, ...]: + from litellm.proxy.proxy_server import llm_router + + callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) + if callback is None: + return () + deployments: Final = ( + tuple(entry for entry in llm_router.settings().guardrail_list if entry["guardrail_name"] == guardrail_name) + if llm_router is not None and llm_router.settings().is_set("guardrail_list") + else () + ) + if len(deployments) <= 1: + return (callback,) + return tuple( + chain.from_iterable(_deployment_callbacks(entry) for entry in deployments if entry.get("callback") is not None) + ) + + def _pipeline_step_streams_live(step: "PipelineStep") -> bool: """A detect-only step whose guardrail set ``streaming_buffer_until_moderated: false`` takes its verdict after the chunks reach the client. A guardrail that masks response content keeps buffering, since the client would @@ -708,10 +726,12 @@ def _pipeline_step_streams_live(step: "PipelineStep") -> bool: not-found outcome settles the stream before release""" if not pipeline_step_is_detect_only(step): return False - callback: Final = PipelineExecutor.find_guardrail_callback(step.guardrail) - if callback is None or callback.mask_response_content: - return False - return not unified_guardrail.resolve_streaming_flag(callback, "streaming_buffer_until_moderated", True) + callbacks: Final = _pipeline_callback_options(step.guardrail) + return bool(callbacks) and all( + not callback.mask_response_content + and not unified_guardrail.resolve_streaming_flag(callback, "streaming_buffer_until_moderated", True) + for callback in callbacks + ) def _drop_stream_chunk_token_ids(chunk: object) -> None: @@ -795,10 +815,13 @@ def _ran_or_runs_outside_pipelines( def _runs_only_in_pipelines( name: str, data: Mapping[str, object], pre_call: frozenset[str], post_call: frozenset[str] ) -> bool: - return not _ran_or_runs_outside_pipelines( - PipelineExecutor.find_guardrail_callback(name), - data, - frozenset(mode for mode, names in (("pre_call", pre_call), ("post_call", post_call)) if name in names), + return not any( + _ran_or_runs_outside_pipelines( + callback, + data, + frozenset(mode for mode, names in (("pre_call", pre_call), ("post_call", post_call)) if name in names), + ) + for callback in _pipeline_callback_options(name) ) @@ -933,36 +956,36 @@ def _deployment_callbacks(deployment: "GuardrailTypedDict") -> tuple[CustomGuard @overload def _load_balanced_callbacks( - callbacks: Sequence[CustomLogger], eligible: Callable[[CustomGuardrail], bool] + callbacks: Sequence[CustomLogger], eligible: Callable[[CustomGuardrail], bool], *, eligible_only: bool = False ) -> tuple[CustomLogger, ...]: ... @overload def _load_balanced_callbacks( - callbacks: Sequence[object], eligible: Callable[[CustomGuardrail], bool] + callbacks: Sequence[object], eligible: Callable[[CustomGuardrail], bool], *, eligible_only: bool = False ) -> tuple[object, ...]: ... def _load_balanced_callbacks( - callbacks: Sequence[object], eligible: Callable[[CustomGuardrail], bool] + callbacks: Sequence[object], eligible: Callable[[CustomGuardrail], bool], *, eligible_only: bool = False ) -> tuple[object, ...]: - return tuple(_iter_load_balanced_callbacks(callbacks, eligible)) + return tuple(_iter_load_balanced_callbacks(callbacks, eligible, eligible_only=eligible_only)) @overload def _iter_load_balanced_callbacks( - callbacks: Sequence[CustomLogger], eligible: Callable[[CustomGuardrail], bool] + callbacks: Sequence[CustomLogger], eligible: Callable[[CustomGuardrail], bool], *, eligible_only: bool = False ) -> Iterator[CustomLogger]: ... @overload def _iter_load_balanced_callbacks( - callbacks: Sequence[object], eligible: Callable[[CustomGuardrail], bool] + callbacks: Sequence[object], eligible: Callable[[CustomGuardrail], bool], *, eligible_only: bool = False ) -> Iterator[object]: ... def _iter_load_balanced_callbacks( - callbacks: Sequence[object], eligible: Callable[[CustomGuardrail], bool] + callbacks: Sequence[object], eligible: Callable[[CustomGuardrail], bool], *, eligible_only: bool = False ) -> Iterator[object]: from litellm.proxy.proxy_server import llm_router from litellm.router_strategy.simple_shuffle import simple_shuffle @@ -1004,7 +1027,7 @@ def _iter_load_balanced_callbacks( model=callback.guardrail_name, request_kwargs=None, ) - return _deployment_callbacks(deployment) + return tuple(member for member in _deployment_callbacks(deployment) if not eligible_only or eligible(member)) yield from chain.from_iterable(selected(index, callback) for index, callback in enumerate(callbacks)) @@ -1063,15 +1086,14 @@ def _merge_pipeline_metadata_writes(data: dict, modified_data: Mapping[str, obje def _pipeline_step_supports_streaming(guardrail_name: str, translation: "BaseTranslation | None") -> bool: - callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) - if callback is None: - return False - if PipelineExecutor.supports_unified_execution(callback): - return True - return ( - translation is not None - and type(translation).assembles_streamed_response - and PipelineExecutor.supports_streaming_execution(callback) + return any( + PipelineExecutor.supports_unified_execution(callback) + or ( + translation is not None + and type(translation).assembles_streamed_response + and PipelineExecutor.supports_streaming_execution(callback) + ) + for callback in _pipeline_callback_options(guardrail_name) ) @@ -5015,7 +5037,17 @@ class ProxyLogging: ) current_response: AsyncIterable[object] = response # rebind-ok: each guardrail wraps the previous layer guarded_layers: Final[list[AsyncGenerator[object, None]]] = [] # mutable-ok: closed on disconnect - for guardrail, holds_until_moderated in guardrails: + for guardrail in _load_balanced_callbacks( + tuple(cb for cb, _ in guardrails), + lambda cb: ( + any(cb is original for original, _ in guardrails) + and _post_call_guardrail_eligible(cb, merged_request_data, frozenset()) + ), + eligible_only=True, + ): + if not isinstance(guardrail, CustomGuardrail): + continue + holds_until_moderated = next(holds for original, holds in guardrails if original is guardrail) if guarded_layers: current_response = ProxyLogging._released_copies(current_response) guarded_layers.append(current_response) diff --git a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py index 23831e074da..1a42db0e086 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py +++ b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py @@ -3,6 +3,7 @@ Test guardrail load balancing through the Router and ProxyLogging. """ from collections.abc import AsyncIterator, Mapping +from contextlib import nullcontext from queue import SimpleQueue from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -16,12 +17,15 @@ from litellm.caching import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails import guardrail_registry from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPI from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 +from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor from litellm.proxy.utils import ProxyLogging +from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep from litellm.types.guardrails import GuardrailEventHooks, GuardrailStreamScope from litellm.types.utils import CallTypesLiteral, ModelResponseStream @@ -81,6 +85,21 @@ class NonIteratorGuardrail(CustomGuardrail): pass +class RelayGuardrail(CustomGuardrail): + def __init__(self, guardrail_id: str, fail: bool = False): + super().__init__(guardrail_name="relay-lb", event_hook="post_call", default_on=True) + self.guardrail_id = guardrail_id + self.calls = 0 + self.fail = fail + self.streaming_end_of_stream_only = True + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + self.calls += 1 + if self.fail: + raise RuntimeError("guardrail offline") + return inputs + + class GuardrailActivator(CustomLogger): def __init__(self, action: str): super().__init__() @@ -232,8 +251,9 @@ async def _guarded_output(logging: ProxyLogging, streaming: bool) -> str: @pytest.mark.asyncio @pytest.mark.parametrize("streaming", (False, True)) +@pytest.mark.parametrize("pipeline", (False, True)) async def test_presidio_load_balancing_keeps_selected_output_sibling( - streaming: bool, monkeypatch: pytest.MonkeyPatch + streaming: bool, pipeline: bool, monkeypatch: pytest.MonkeyPatch ) -> None: router: Final = Router(model_list=[]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) @@ -260,7 +280,30 @@ async def test_presidio_load_balancing_keeps_selected_output_sibling( logging: Final = ProxyLogging(user_api_key_cache=DualCache()) outputs: Final[SimpleQueue[str]] = SimpleQueue() for _ in range(50): - content: Final = await _guarded_output(logging, streaming) + if pipeline: + chunks: Final = [chunk async for chunk in _output_chunks()] if streaming else None + result: Final = await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="presidio-balanced")], + mode="post_call", + data={ + "messages": [{"role": "user", "content": "raw"}], + "response": litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "raw"}}]), + }, + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"), + call_type="completion", + policy_name="presidio-pipeline", + streaming_chunks=chunks, + endpoint_translation=OpenAIChatCompletionsHandler() if streaming else None, + ) + assert result.step_results[0].outcome == "pass" + if streaming: + assert chunks is not None + content = "".join(_chunk_text(chunk) for chunk in chunks) + else: + assert result.modified_data is not None + content = result.modified_data["response"].choices[0].message.content + else: + content = await _guarded_output(logging, streaming) assert content in ("masked-east", "masked-west") outputs.put(content) assert frozenset(outputs.get() for _ in range(outputs.qsize())) == {"masked-east", "masked-west"} @@ -526,3 +569,201 @@ async def test_missing_callback_does_not_poison_healthy_replica( with pytest.raises(ValueError, match="No callback found for guardrail: partial-lb"): await _dispatch_scoped_hook(logging, hook) assert guardrail.calls == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("pre_call", "post_call")) +async def test_pipeline_balanced_step_splits_traffic_once_per_request(mode, monkeypatch): + copies: Final = tuple( + MockGuardrail("pipeline-lb", region, event_hook=GuardrailEventHooks(mode)) + for region in ("east", "west") + ) + router: Final = Router( + model_list=[], + guardrail_list=[ + { + "guardrail_name": "pipeline-lb", + "callback": cb, + "litellm_params": {"guardrail": "custom", "mode": mode}, + } + for cb in copies + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", list(copies)) + for index in range(50): + result = await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="pipeline-lb")], + mode=mode, + data={"messages": [{"role": "user", "content": "raw"}]}, + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + policy_name="balanced-pipeline", + ) + assert result.step_results[0].outcome == "pass" + assert sum(cb.calls for cb in copies) == index + 1 + assert all(cb.calls > 0 for cb in copies) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scenario", ("mode", "stream_scope", "none", "weights", "single")) +async def test_pipeline_balancing_respects_eligibility_weights_and_single_copy(scenario, monkeypatch): + first: Final = MockGuardrail( + "pipeline-scoped", "east", + event_hook=GuardrailEventHooks.pre_call if scenario in ("mode", "single", "none") else GuardrailEventHooks.post_call, + stream_scope="streaming" if scenario == "stream_scope" else "both", + ) + second: Final = MockGuardrail( + "pipeline-scoped", "west", + event_hook=GuardrailEventHooks.pre_call if scenario == "none" else GuardrailEventHooks.post_call, + ) + copies: Final = (first, second) + router: Final = Router( + model_list=[], + guardrail_list=[ + { + "guardrail_name": "pipeline-scoped", + "callback": cb, + "litellm_params": {"guardrail": "custom", "mode": "post_call", "weight": weight}, + } + for cb, weight in ((first, 0), (second, 1))[:1 if scenario == "single" else 2] + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", list(copies)) + response: Final = object() + for _ in range(20): + result = await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="pipeline-scoped")], + mode="post_call", + data={"stream": False, "response": response}, + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + policy_name="balanced-pipeline", + ) + assert result.step_results[0].outcome == ("skip" if scenario == "none" else "pass") + assert first.calls == (20 if scenario == "single" else 0) + assert second.calls == (0 if scenario in ("single", "none") else 20) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail", (False, True)) +async def test_streamed_pipeline_selects_supported_replica_and_preserves_errors(fail, monkeypatch): + unsupported: Final = MockGuardrail("relay-lb", "native", event_hook=GuardrailEventHooks.post_call) + supported: Final = RelayGuardrail("unified", fail=fail) + supported.default_on = False + router: Final = Router( + model_list=[], + guardrail_list=[ + { + "guardrail_name": "relay-lb", + "callback": cb, + "litellm_params": {"guardrail": "custom", "mode": "post_call", "weight": weight}, + } + for cb, weight in ((unsupported, 1), (supported, 0)) + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", [unsupported, supported]) + logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + for index in range(20): + with pytest.raises(RuntimeError, match="guardrail offline") if fail else nullcontext(): + chunks = [ + chunk + async for chunk in logging.async_post_call_streaming_iterator_hook( + response=_output_chunks(), + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"), + request_data={ + "model": "test-model", + "stream": True, + "metadata": { + "_guardrail_pipelines": [("streamed-pipeline", GuardrailPipeline( + mode="post_call", steps=[PipelineStep(guardrail="relay-lb")], + ))], + }, + }, + ) + ] + assert "".join(_chunk_text(chunk) for chunk in chunks) == "raw" + assert supported.calls == index + 1 + assert unsupported.calls == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail", (False, True)) +async def test_passthrough_balanced_relay_splits_traffic_once_per_request(fail, monkeypatch): + copies: Final = tuple(RelayGuardrail(region, fail=fail) for region in ("east", "west")) + router: Final = Router( + model_list=[], + guardrail_list=[ + { + "guardrail_name": "relay-lb", + "callback": cb, + "litellm_params": {"guardrail": "custom", "mode": "post_call"}, + } + for cb in copies + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", list(copies)) + logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + for index in range(50): + with pytest.raises(RuntimeError, match="guardrail offline") if fail else nullcontext(): + chunks = [ + chunk + async for chunk in logging.async_post_call_event_stream_iterator_hook( + response=_output_chunks(), + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"), + request_data={"model": "test-model", "stream": True}, + guardrails=tuple((cb, False) for cb in copies), + ) + ] + assert "".join(_chunk_text(chunk) for chunk in chunks) == "raw" + assert sum(cb.calls for cb in copies) == index + 1 + assert all(cb.calls > 0 for cb in copies) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scenario", ("mode", "stream_scope", "default_off", "none", "weights", "single")) +async def test_passthrough_balancing_respects_eligible_layers_and_weights(scenario, monkeypatch): + first: Final = RelayGuardrail("east") + second: Final = RelayGuardrail("west") + if scenario == "mode": + first.event_hook = GuardrailEventHooks.pre_call + if scenario == "stream_scope": + first.stream_scope_default = "non_streaming" + if scenario in ("default_off", "none"): + first.default_on = False + if scenario == "none": + second.default_on = False + copies: Final = (first,) if scenario == "single" else (first, second) + router: Final = Router( + model_list=[], + guardrail_list=[ + { + "guardrail_name": "relay-lb", + "callback": cb, + "litellm_params": { + "guardrail": "custom", "mode": "post_call", + "weight": 0 if cb is first and scenario == "weights" else 1, + }, + } + for cb in copies + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", list(copies)) + logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + for _ in range(20): + chunks = [ + chunk + async for chunk in logging.async_post_call_event_stream_iterator_hook( + response=_output_chunks(), + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"), + request_data={"model": "test-model", "stream": True}, + guardrails=tuple((cb, cb is second) for cb in copies), + ) + ] + assert "".join(_chunk_text(chunk) for chunk in chunks) == "raw" + assert first.calls == (20 if scenario == "single" else 0) + assert second.calls == (0 if scenario in ("none", "single") else 20) From 471f50dc248b3b61d2b27c163bd2a59f394ec7e1 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 21:13:43 +0000 Subject: [PATCH 08/11] fix: filter pipeline replicas by streaming endpoint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/policy_engine/pipeline_executor.py | 13 ++- .../test_guardrail_load_balancing.py | 83 +++++++++++++++++++ 2 files changed, 94 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index e57143a9422..354ddeabe00 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -337,6 +337,7 @@ def _pipeline_callback_eligible( raw_request_snapshot: Mapping[str, object] | None, mode: str, streaming_chunks: list[object] | None, # mutable-ok: shared buffered-stream chunks, read only here + endpoint_translation: "BaseTranslation | None", ) -> bool: source: Final = raw_request_snapshot if callback.scan_raw_request and raw_request_snapshot is not None else data metadata: Final = source.get("metadata") @@ -350,7 +351,13 @@ def _pipeline_callback_eligible( } event_type: Final = _PIPELINE_EVENT_HOOKS.get(mode) return (event_type is None or callback.should_run_guardrail(data=hook_input, event_type=event_type) is True) and ( # pyright: ignore[reportUnknownMemberType] # legacy method accepts the str-keyed request payload - streaming_chunks is None or PipelineExecutor.supports_streaming_execution(callback) + streaming_chunks is None + or PipelineExecutor.supports_unified_execution(callback) + or ( + endpoint_translation is not None + and type(endpoint_translation).assembles_streamed_response + and PipelineExecutor.supports_streaming_execution(callback) + ) ) @@ -633,7 +640,9 @@ class PipelineExecutor: return ("error", None, f"Guardrail '{step.guardrail}' not found", None) callbacks: Final = _load_balanced_callbacks( (callback,), - lambda cb: _pipeline_callback_eligible(cb, step, data, raw_request_snapshot, mode, streaming_chunks), + lambda cb: _pipeline_callback_eligible( + cb, step, data, raw_request_snapshot, mode, streaming_chunks, endpoint_translation + ), eligible_only=True, ) working_data = data diff --git a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py index 1a42db0e086..dba740c8bc9 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py +++ b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py @@ -100,6 +100,26 @@ class RelayGuardrail(CustomGuardrail): return inputs +class LegacyPostCallGuardrail(CustomGuardrail): + def __init__(self, fail: bool): + super().__init__(guardrail_name="relay-lb", event_hook="post_call", default_on=True) + self.calls = 0 + self.response: object = None + self.fail = fail + + async def async_post_call_success_hook( + self, + data: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + response: object, + ) -> object: + self.calls += 1 + self.response = response + if response is not None and self.fail: + raise RuntimeError("guardrail offline") + return response + + class GuardrailActivator(CustomLogger): def __init__(self, action: str): super().__init__() @@ -689,6 +709,69 @@ async def test_streamed_pipeline_selects_supported_replica_and_preserves_errors( assert unsupported.calls == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("fail", (False, True)) +@pytest.mark.parametrize("empty", (False, True)) +@pytest.mark.parametrize("route", ("/v1/chat/completions", "/v1beta/models/gemini-2.5-flash:streamGenerateContent")) +async def test_streamed_pipeline_selects_replica_compatible_with_endpoint(route, empty, fail, monkeypatch): + legacy: Final = LegacyPostCallGuardrail(fail=fail) + unified: Final = RelayGuardrail("unified", fail=fail) + router: Final = Router( + model_list=[], + guardrail_list=[ + { + "guardrail_name": "relay-lb", + "callback": cb, + "litellm_params": {"guardrail": "custom", "mode": "post_call", "weight": weight}, + } + for cb, weight in ((legacy, 1), (unified, 0)) + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", [legacy, unified]) + logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + gemini: Final = route.endswith(":streamGenerateContent") + payload: Final = ( + ({"candidates": [{"content": {"parts": [{"text": "raw"}]}, "finishReason": "STOP"}]},) + if gemini + else tuple([chunk async for chunk in _output_chunks()]) + ) + source: Final = () if empty else payload + + async def response_chunks() -> AsyncIterator[object]: + for chunk in source: + yield chunk + + for index in range(3): + with pytest.raises(RuntimeError, match="guardrail offline") if fail and not empty else nullcontext(): + delivered = [ + chunk + async for chunk in logging.async_post_call_streaming_iterator_hook( + response=response_chunks(), + user_api_key_dict=UserAPIKeyAuth(request_route=route), + request_data={ + "model": "test-model", + "stream": True, + "metadata": { + "_guardrail_pipelines": [ + ( + "streamed-pipeline", + GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail="relay-lb")], + ), + ) + ], + }, + }, + ) + ] + assert delivered == list(source) + assert unified.calls == (index + 1 if gemini and not empty else 0) + assert legacy.calls == (index + 1 if not gemini and not empty else 0) + assert isinstance(legacy.response, litellm.ModelResponse) is (not gemini and not empty) + + @pytest.mark.asyncio @pytest.mark.parametrize("fail", (False, True)) async def test_passthrough_balanced_relay_splits_traffic_once_per_request(fail, monkeypatch): From ade14ec3d41a7b897b706ec551c3dc74bad21051 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 21:43:25 +0000 Subject: [PATCH 09/11] fix: share pipeline streaming admission eligibility Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/policy_engine/pipeline_executor.py | 4 +- litellm/proxy/utils.py | 28 ++++----- .../test_guardrail_load_balancing.py | 63 +++++++++++++++++++ 3 files changed, 77 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 354ddeabe00..2b8622d8bfc 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -330,7 +330,7 @@ def _prepare_hook_input( return hook_input, scans_raw_request -def _pipeline_callback_eligible( +def pipeline_callback_eligible( callback: CustomGuardrail, step: PipelineStep, data: Mapping[str, object], @@ -640,7 +640,7 @@ class PipelineExecutor: return ("error", None, f"Guardrail '{step.guardrail}' not found", None) callbacks: Final = _load_balanced_callbacks( (callback,), - lambda cb: _pipeline_callback_eligible( + lambda cb: pipeline_callback_eligible( cb, step, data, raw_request_snapshot, mode, streaming_chunks, endpoint_translation ), eligible_only=True, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8e85753b114..583eb9615b3 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -272,6 +272,7 @@ from litellm.proxy.litellm_pre_call_utils import ( from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.policy_engine.pipeline_executor import ( PipelineExecutor, + pipeline_callback_eligible, pipeline_step_is_detect_only, recorded_guardrail_information, ) @@ -1085,15 +1086,12 @@ def _merge_pipeline_metadata_writes(data: dict, modified_data: Mapping[str, obje _merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key)) -def _pipeline_step_supports_streaming(guardrail_name: str, translation: "BaseTranslation | None") -> bool: +def _pipeline_step_supports_streaming( + step: "PipelineStep", data: Mapping[str, object], translation: "BaseTranslation | None" +) -> bool: return any( - PipelineExecutor.supports_unified_execution(callback) - or ( - translation is not None - and type(translation).assembles_streamed_response - and PipelineExecutor.supports_streaming_execution(callback) - ) - for callback in _pipeline_callback_options(guardrail_name) + pipeline_callback_eligible(callback, step, data, None, "post_call", [], translation) + for callback in _pipeline_callback_options(step.guardrail) ) @@ -1223,21 +1221,19 @@ def _body_selected_deferrals( def _pipeline_unsupported_streaming_guardrails( - pipeline: "GuardrailPipeline", translation: "BaseTranslation | None" + pipeline: "GuardrailPipeline", translation: "BaseTranslation | None", data: Mapping[str, object] ) -> tuple[str, ...]: return tuple( dict.fromkeys( - step.guardrail - for step in pipeline.steps - if not _pipeline_step_supports_streaming(step.guardrail, translation) + step.guardrail for step in pipeline.steps if not _pipeline_step_supports_streaming(step, data, translation) ) ) def _pipeline_is_streamable( - policy_name: str, pipeline: "GuardrailPipeline", translation: "BaseTranslation | None" + policy_name: str, pipeline: "GuardrailPipeline", translation: "BaseTranslation | None", data: Mapping[str, object] ) -> bool: - unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline, translation) + unsupported: Final = _pipeline_unsupported_streaming_guardrails(pipeline, translation, data) if not unsupported: return True verbose_proxy_logger.warning( @@ -1266,7 +1262,7 @@ def stream_gated_guardrail_names( tuple( (policy_name, pipeline) for policy_name, pipeline in _post_call_pipelines(request_data) - if not _pipeline_unsupported_streaming_guardrails(pipeline, translation) + if not _pipeline_unsupported_streaming_guardrails(pipeline, translation, request_data) ) ) @@ -1302,7 +1298,7 @@ def _streamable_post_call_pipelines( return tuple( (policy_name, pipeline) for policy_name, pipeline in post_call_pipelines - if _pipeline_is_streamable(policy_name, pipeline, translation) + if _pipeline_is_streamable(policy_name, pipeline, translation, request_data) ) diff --git a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py index dba740c8bc9..dac9df277a3 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py +++ b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py @@ -120,6 +120,20 @@ class LegacyPostCallGuardrail(CustomGuardrail): return response +class FallbackIteratorGuardrail(CustomGuardrail): + def __init__(self, fail: bool): + super().__init__(guardrail_name="relay-lb", event_hook="post_call", default_on=True) + self.calls = 0 + self.fail = fail + + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + self.calls += 1 + if self.fail: + raise RuntimeError("guardrail offline") + async for chunk in response: + yield chunk + + class GuardrailActivator(CustomLogger): def __init__(self, action: str): super().__init__() @@ -772,6 +786,55 @@ async def test_streamed_pipeline_selects_replica_compatible_with_endpoint(route, assert isinstance(legacy.response, litellm.ModelResponse) is (not gemini and not empty) +@pytest.mark.asyncio +@pytest.mark.parametrize("scenario", ("mode", "stream_scope", "overlap", "none")) +@pytest.mark.parametrize("output", ("normal", "empty", "error")) +async def test_streaming_pipeline_falls_back_when_eligible_and_capable_replicas_do_not_overlap( + scenario, output, monkeypatch, +): + unified: Final = RelayGuardrail("unified", fail=output == "error") + unified.event_hook = GuardrailEventHooks.pre_call if scenario in ("mode", "none") else GuardrailEventHooks.post_call + unified.apply_stream_scope("non_streaming" if scenario == "stream_scope" else "both") + native: Final = FallbackIteratorGuardrail(fail=output == "error") + native.event_hook = GuardrailEventHooks.pre_call if scenario == "none" else GuardrailEventHooks.post_call + router: Final = Router( + model_list=[], + guardrail_list=[ + { + "guardrail_name": "relay-lb", "callback": cb, + "litellm_params": {"guardrail": "custom", "mode": "post_call", "weight": weight}, + } + for cb, weight in ((native, 1), (unified, 0)) + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", [native, unified]) + logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + + async def response_chunks() -> AsyncIterator[ModelResponseStream]: + if output != "empty": + async for chunk in _output_chunks(): + yield chunk + + with pytest.raises(RuntimeError, match="guardrail offline") if output == "error" and scenario != "none" else nullcontext(): + delivered: Final = [ + chunk + async for chunk in logging.async_post_call_streaming_iterator_hook( + response=response_chunks(), + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"), + request_data={ + "model": "test-model", "stream": True, + "metadata": {"_guardrail_pipelines": [("streamed-pipeline", GuardrailPipeline( + mode="post_call", steps=[PipelineStep(guardrail="relay-lb")], + ))]}, + }, + ) + ] + assert "".join(_chunk_text(chunk) for chunk in delivered) == ("" if output == "empty" else "raw") + assert native.calls == (1 if scenario in ("mode", "stream_scope") else 0) + assert unified.calls == (1 if scenario == "overlap" and output != "empty" else 0) + + @pytest.mark.asyncio @pytest.mark.parametrize("fail", (False, True)) async def test_passthrough_balanced_relay_splits_traffic_once_per_request(fail, monkeypatch): From 4a532008c667b2473cec2ef00504289741f025f9 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 21:52:52 +0000 Subject: [PATCH 10/11] fix: preserve skipped steps in streaming pipelines Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 27 +++++++++- .../test_guardrail_load_balancing.py | 52 +++++++++++++++++++ 2 files changed, 77 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 583eb9615b3..7386447b0c4 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1089,9 +1089,32 @@ def _merge_pipeline_metadata_writes(data: dict, modified_data: Mapping[str, obje def _pipeline_step_supports_streaming( step: "PipelineStep", data: Mapping[str, object], translation: "BaseTranslation | None" ) -> bool: - return any( + from litellm.proxy.proxy_server import llm_router + + callbacks: Final = _pipeline_callback_options(step.guardrail) + if not callbacks: + return False + balanced: Final = ( + llm_router is not None + and llm_router.settings().is_set("guardrail_list") + and sum(entry["guardrail_name"] == step.guardrail for entry in llm_router.settings().guardrail_list) > 1 + ) + if not balanced: + callback: Final = callbacks[0] + return PipelineExecutor.supports_unified_execution(callback) or ( + translation is not None + and type(translation).assembles_streamed_response + and PipelineExecutor.supports_streaming_execution(callback) + ) + streaming_data: Final = {**data, "stream": True} + eligible_callbacks: Final = tuple( + callback + for callback in callbacks + if pipeline_callback_eligible(callback, step, streaming_data, None, "post_call", None, translation) + ) + return not eligible_callbacks or any( pipeline_callback_eligible(callback, step, data, None, "post_call", [], translation) - for callback in _pipeline_callback_options(step.guardrail) + for callback in eligible_callbacks ) diff --git a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py index dac9df277a3..2f3fce83d11 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py +++ b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py @@ -835,6 +835,58 @@ async def test_streaming_pipeline_falls_back_when_eligible_and_capable_replicas_ assert unified.calls == (1 if scenario == "overlap" and output != "empty" else 0) +@pytest.mark.asyncio +@pytest.mark.parametrize("first_balanced", (False, True)) +@pytest.mark.parametrize("output", ("normal", "empty", "error")) +async def test_skipped_streaming_step_does_not_drop_later_policy_only_guardrail( + first_balanced, output, monkeypatch, +): + skipped: Final = tuple(RelayGuardrail(region) for region in (("east", "west") if first_balanced else ("east",))) + for callback in skipped: + callback.guardrail_name = "skipped-step" + callback.apply_stream_scope("non_streaming") + checked: Final = RelayGuardrail("checked", fail=output == "error") + checked.guardrail_name = "policy-only" + checked.default_on = False + callbacks: Final = (*skipped, checked) + router: Final = Router( + model_list=[], + guardrail_list=[ + { + "guardrail_name": cb.guardrail_name, "callback": cb, + "litellm_params": {"guardrail": "custom", "mode": "post_call"}, + } + for cb in callbacks + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", list(callbacks)) + logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + + async def response_chunks() -> AsyncIterator[ModelResponseStream]: + if output != "empty": + async for chunk in _output_chunks(): + yield chunk + + with pytest.raises(RuntimeError, match="guardrail offline") if output == "error" else nullcontext(): + delivered: Final = [ + chunk + async for chunk in logging.async_post_call_streaming_iterator_hook( + response=response_chunks(), + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"), + request_data={ + "model": "test-model", "stream": True, + "metadata": {"_guardrail_pipelines": [("streamed-pipeline", GuardrailPipeline( + mode="post_call", steps=[PipelineStep(guardrail="skipped-step"), PipelineStep(guardrail="policy-only")], + ))]}, + }, + ) + ] + assert "".join(_chunk_text(chunk) for chunk in delivered) == ("" if output == "empty" else "raw") + assert checked.calls == (0 if output == "empty" else 1) + assert all(cb.calls == 0 for cb in skipped) + + @pytest.mark.asyncio @pytest.mark.parametrize("fail", (False, True)) async def test_passthrough_balanced_relay_splits_traffic_once_per_request(fail, monkeypatch): From 505091c1f5eb243d17a4c632066b1f1efe1f3df7 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 22:39:29 +0000 Subject: [PATCH 11/11] fix: reselect parallel guardrails after sequential callbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 59 ++++++++----- .../test_guardrail_load_balancing.py | 82 ++++++++++++++++++- 2 files changed, 118 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 7386447b0c4..597b8547c83 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3370,25 +3370,25 @@ class ProxyLogging: self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails) return data + def _eligible_pre_call(callback: CustomGuardrail) -> bool: + return ( + data is not None + and callback.guardrail_name not in pipeline_managed + and callback.should_run_guardrail( + data=raw_request_snapshot + if callback.scan_raw_request and raw_request_snapshot is not None + else data, + event_type=GuardrailEventHooks.pre_mcp_call + if call_type in MCP_GUARDRAIL_CALL_TYPES + else GuardrailEventHooks.pre_call, + ) + is True + ) + callback_iterators: Final = tee( caps.resolved_callbacks if skip_guardrails - else _iter_load_balanced_callbacks( - caps.resolved_callbacks, - lambda cb: ( - data is not None - and cb.guardrail_name not in pipeline_managed - and cb.should_run_guardrail( - data=raw_request_snapshot - if cb.scan_raw_request and raw_request_snapshot is not None - else data, - event_type=GuardrailEventHooks.pre_mcp_call - if call_type in MCP_GUARDRAIL_CALL_TYPES - else GuardrailEventHooks.pre_call, - ) - is True - ), - ) + else _iter_load_balanced_callbacks(caps.resolved_callbacks, _eligible_pre_call) ) deferred_route_exc: SensitiveDataRouteException | None = None for _callback in callback_iterators[0]: @@ -3469,15 +3469,36 @@ class ProxyLogging: if deferred_route_exc is not None and data is not None: data = await self._handle_sensitive_data_route_exception(deferred_route_exc, data, user_api_key_dict) + selected_callbacks: Final = tuple(callback_iterators[1]) + sequential_names: Final = frozenset( + cb.guardrail_name + for cb in selected_callbacks + if isinstance(cb, CustomGuardrail) and not getattr(cb, "run_in_parallel", False) + ) parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = ( () if skip_guardrails else tuple( cb - for cb in callback_iterators[1] + for cb in _load_balanced_callbacks( + tuple( + cb + for cb in caps.resolved_callbacks + if isinstance(cb, CustomGuardrail) + and getattr(cb, "run_in_parallel", False) + and cb.guardrail_name not in pipeline_managed + ), + lambda cb: ( + getattr(cb, "run_in_parallel", False) + and _eligible_pre_call(cb) + and ( + cb.guardrail_name not in sequential_names + or any(cb is selected for selected in selected_callbacks) + ) + ), + eligible_only=True, + ) if isinstance(cb, CustomGuardrail) - and getattr(cb, "run_in_parallel", False) - and not (cb.guardrail_name and cb.guardrail_name in pipeline_managed) ) ) if parallel_guardrails and data is not None: diff --git a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py index 2f3fce83d11..5011b37c471 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py +++ b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py @@ -44,6 +44,7 @@ class MockGuardrail(CustomGuardrail): stream_scope: GuardrailStreamScope = "both", run_in_parallel: bool = False, scan_raw_request: bool = False, + fail: bool = False, ): super().__init__( guardrail_name=guardrail_name, @@ -55,10 +56,13 @@ class MockGuardrail(CustomGuardrail): ) self.guardrail_id = guardrail_id self.calls = 0 + self.fail = fail async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): self.calls += 1 MockGuardrail.call_count += 1 + if self.fail: + raise RuntimeError("guardrail offline") return None async def async_moderation_hook(self, data, user_api_key_dict, call_type): @@ -490,8 +494,10 @@ async def test_streaming_selects_iterator_capable_deployment(monkeypatch: pytest @pytest.mark.parametrize("action", ("enable", "noop", "raise")) @pytest.mark.parametrize("run_in_parallel", (False, True)) @pytest.mark.parametrize("call_type", ("completion", "call_mcp_tool")) +@pytest.mark.parametrize("activator_first", (False, True)) async def test_pre_call_load_balancing_observes_earlier_callback_activation( - action: str, run_in_parallel: bool, call_type: CallTypesLiteral, monkeypatch: pytest.MonkeyPatch + action: str, run_in_parallel: bool, call_type: CallTypesLiteral, activator_first: bool, + monkeypatch: pytest.MonkeyPatch, ) -> None: event: Final = GuardrailEventHooks.pre_mcp_call if call_type == "call_mcp_tool" else GuardrailEventHooks.pre_call replicas: Final = tuple( @@ -511,7 +517,8 @@ async def test_pre_call_load_balancing_observes_earlier_callback_activation( ], ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) - monkeypatch.setattr(litellm, "callbacks", [GuardrailActivator(action), *replicas]) + activator: Final = GuardrailActivator(action) + monkeypatch.setattr(litellm, "callbacks", [activator, *replicas] if activator_first else [*replicas, activator]) logging: Final = ProxyLogging(user_api_key_cache=DualCache()) for request_index in range(30): if action == "raise": @@ -523,8 +530,75 @@ async def test_pre_call_load_balancing_observes_earlier_callback_activation( await logging.pre_call_hook( user_api_key_dict=UserAPIKeyAuth(), data={"model": "test", "messages": []}, call_type=call_type ) - assert sum(cb.calls for cb in replicas) == (request_index + 1 if action == "enable" else 0) - assert all(cb.calls > 0 for cb in replicas) is (action == "enable") + assert sum(cb.calls for cb in replicas) == ( + request_index + 1 if action == "enable" and (activator_first or run_in_parallel) else 0 + ) + assert all(cb.calls > 0 for cb in replicas) is (action == "enable" and (activator_first or run_in_parallel)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scan_raw_request", (False, True)) +@pytest.mark.parametrize("requested", (False, True)) +@pytest.mark.parametrize("fail", (False, True)) +async def test_late_parallel_activation_rechecks_final_request( + scan_raw_request: bool, requested: bool, fail: bool, monkeypatch: pytest.MonkeyPatch, +) -> None: + replicas: Final = tuple( + MockGuardrail( + "dynamic-lb", region, event_hook=GuardrailEventHooks.pre_call, default_on=False, + run_in_parallel=True, scan_raw_request=scan_raw_request, fail=fail, + ) + for region in ("east", "west") + ) + router: Final = Router( + model_list=[], + guardrail_list=[ + { + "guardrail_name": "dynamic-lb", "id": cb.guardrail_id, "callback": cb, + "litellm_params": {"guardrail": "custom", "mode": "pre_call"}, + } + for cb in replicas + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", [*replicas, GuardrailActivator("enable")]) + logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + should_run: Final = requested or not scan_raw_request + with pytest.raises(RuntimeError, match="guardrail offline") if fail and should_run else nullcontext(): + await logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + data={"model": "test", "messages": [], "guardrails": ["dynamic-lb"] if requested else []}, + call_type="completion", + ) + assert any(cb.calls > 0 for cb in replicas) is should_run + + +@pytest.mark.asyncio +@pytest.mark.parametrize("parallel_weight", (0, 1)) +async def test_mixed_parallel_and_sequential_replicas_keep_one_weighted_selection( + parallel_weight: int, monkeypatch: pytest.MonkeyPatch, +) -> None: + sequential: Final = MockGuardrail("dynamic-lb", "sequential", default_on=True) + parallel: Final = MockGuardrail("dynamic-lb", "parallel", default_on=True, run_in_parallel=True) + router: Final = Router( + model_list=[], + guardrail_list=[ + { + "guardrail_name": "dynamic-lb", "callback": cb, + "litellm_params": {"guardrail": "custom", "mode": "pre_call", "weight": weight}, + } + for cb, weight in ((sequential, 1 - parallel_weight), (parallel, parallel_weight)) + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr(litellm, "callbacks", [sequential, parallel]) + logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + for request_index in range(20): + await logging.pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), data={"model": "test", "messages": []}, call_type="completion", + ) + assert sequential.calls == (request_index + 1) * (1 - parallel_weight) + assert parallel.calls == (request_index + 1) * parallel_weight @pytest.mark.asyncio