feat: add parallel execution handling in during_call_hook (#16279)

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