From 911e802969c3cac9a2ffd8a6f203d8a9efded387 Mon Sep 17 00:00:00 2001 From: Petre Alexandru Date: Thu, 6 Nov 2025 04:35:25 +0200 Subject: [PATCH] feat: add parallel execution handling in during_call_hook (#16279) --- litellm/proxy/utils.py | 78 ++++++++++-------- tests/proxy_unit_tests/test_proxy_utils.py | 96 ++++++++++++++++++++++ 2 files changed, 141 insertions(+), 33 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 7ad39bff6b..62ab69c32e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1053,51 +1053,63 @@ class ProxyLogging: ], ): """ - Runs the CustomGuardrail's async_moderation_hook() + Runs the CustomGuardrail's async_moderation_hook() in parallel """ + # Step 1: Collect all guardrail tasks to run in parallel + guardrail_tasks = [] + for callback in litellm.callbacks: - try: - if isinstance(callback, CustomGuardrail): - ################################################################ - # Check if guardrail should be run for GuardrailEventHooks.during_call hook - ################################################################ + if isinstance(callback, CustomGuardrail): + ################################################################ + # 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": # type: ignore - return - else: - # Main - V2 Guardrails implementation - from litellm.types.guardrails import GuardrailEventHooks + # V1 implementation - backwards compatibility + if callback.event_hook is None and hasattr( + callback, "moderation_check" + ): + if callback.moderation_check == "pre_call": # type: ignore + return + else: + # Main - V2 Guardrails implementation + from litellm.types.guardrails import GuardrailEventHooks - event_type = GuardrailEventHooks.during_call - if call_type == "mcp_call": - event_type = GuardrailEventHooks.during_mcp_call - - if ( - callback.should_run_guardrail( - data=data, event_type=event_type - ) - is not True - ): - continue - # Convert user_api_key_dict to proper format for async_moderation_hook + event_type = GuardrailEventHooks.during_call if call_type == "mcp_call": - user_api_key_auth_dict = ( - self._convert_user_api_key_auth_to_dict(user_api_key_dict) - ) - else: - user_api_key_auth_dict = user_api_key_dict + event_type = GuardrailEventHooks.during_mcp_call - await callback.async_moderation_hook( + if ( + callback.should_run_guardrail( + data=data, event_type=event_type + ) + is not True + ): + continue + # Convert user_api_key_dict to proper format for async_moderation_hook + if call_type == "mcp_call": + user_api_key_auth_dict = ( + self._convert_user_api_key_auth_to_dict(user_api_key_dict) + ) + else: + user_api_key_auth_dict = user_api_key_dict + + # Add task to list for parallel execution + guardrail_tasks.append( + callback.async_moderation_hook( data=data, user_api_key_dict=user_api_key_auth_dict, # type: ignore call_type=call_type, ) + ) + + # Step 2: Run all guardrail tasks in parallel + if guardrail_tasks: + try: + await asyncio.gather(*guardrail_tasks) except Exception as e: + # If any guardrail raises an exception, it will propagate here raise e + return data async def failed_tracking_alert( diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 38c0bf934d..66b748e548 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -2154,3 +2154,99 @@ async def test_post_call_failure_hook_auth_error_llm_api_route(): # Assert that _handle_logging_proxy_only_error WAS called # because /v1/chat/completions is an LLM API route mock_handle_logging.assert_called_once() + + +@pytest.mark.asyncio +async def test_during_call_hook_parallel_execution(): + """ + Test that multiple guardrails in during_call_hook are executed in parallel. + Verifies parallel execution by checking timing and execution order. + """ + from litellm.caching.caching import DualCache + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.utils import ProxyLogging + from litellm.types.guardrails import GuardrailEventHooks + + cache = DualCache() + proxy_logging = ProxyLogging(user_api_key_cache=cache) + execution_order = [] + + class TestGuardrail(CustomGuardrail): + def __init__(self, name): + super().__init__( + guardrail_name=name, + event_hook=GuardrailEventHooks.during_call, + default_on=True + ) + self.name = name + + async def async_moderation_hook(self, data, user_api_key_dict, call_type): + execution_order.append(f"{self.name}_start") + await asyncio.sleep(0.1) + execution_order.append(f"{self.name}_end") + return data + + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + litellm.callbacks = [TestGuardrail(f"g{i}") for i in range(3)] + + start_time = asyncio.get_event_loop().time() + result = await proxy_logging.during_call_hook( + data={"model": "gpt-4", "messages": [{"role": "user", "content": "test"}]}, + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + call_type="completion", + ) + execution_time = asyncio.get_event_loop().time() - start_time + + # Verify parallel execution: all start before any end + first_end_idx = next(i for i, item in enumerate(execution_order) if "end" in item) + starts_before_end = sum(1 for item in execution_order[:first_end_idx] if "start" in item) + assert starts_before_end == 3, f"Expected 3 starts before first end, got {starts_before_end}" + + # Verify timing: parallel ~0.1s vs sequential ~0.3s + assert execution_time < 0.2, f"Parallel execution took {execution_time}s, expected < 0.2s" + assert result["model"] == "gpt-4" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_during_call_hook_parallel_execution_with_error(): + """ + Test that exceptions from guardrails are properly raised in parallel execution. + """ + from litellm.caching.caching import DualCache + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.utils import ProxyLogging + from litellm.types.guardrails import GuardrailEventHooks + + cache = DualCache() + proxy_logging = ProxyLogging(user_api_key_cache=cache) + + class FailingGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="failing_guardrail", + event_hook=GuardrailEventHooks.during_call, + default_on=True + ) + + async def async_moderation_hook(self, data, user_api_key_dict, call_type): + raise ValueError("Guardrail violation detected!") + + original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] + + try: + litellm.callbacks = [FailingGuardrail()] + + with pytest.raises(ValueError) as exc_info: + await proxy_logging.during_call_hook( + data={"model": "gpt-4", "messages": [{"role": "user", "content": "test"}]}, + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), + call_type="completion", + ) + + assert "Guardrail violation detected!" in str(exc_info.value) + finally: + litellm.callbacks = original_callbacks \ No newline at end of file