fix(proxy): run parallel guardrails before an operator hook and Custom Code guardrails for a refused model

litellm/proxy/utils.py

@@ -230,6 +230,7 @@ from litellm.proxy.db.token_auth import (
     resolve_database_token_auth,
 )
 from litellm.proxy.guardrails.exception_utils import enrich_http_exception_with_guardrail_context
+from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import CustomCodeGuardrail
 from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
     UnifiedLLMGuardrails,
     resolve_endpoint_translation,
@@ -1677,8 +1678,9 @@ _LITELLM_PACKAGES: Final = frozenset({"litellm", "litellm_enterprise", "enterpri
 
 
 def _is_custom_pre_call_hook(callback: CustomLogger) -> bool:
-    return type(callback).__module__.partition(".")[0] not in _LITELLM_PACKAGES and (
-        isinstance(callback, CustomGuardrail) or "async_pre_call_hook" in type(callback).__dict__
+    return isinstance(callback, CustomCodeGuardrail) or (
+        type(callback).__module__.partition(".")[0] not in _LITELLM_PACKAGES
+        and (isinstance(callback, CustomGuardrail) or "async_pre_call_hook" in type(callback).__dict__)
     )
 
 
@@ -1686,10 +1688,11 @@ async def _skipped_for_unroutable_model(
     guardrail: CustomGuardrail,
     data: Mapping[str, object],
     probe: "UnroutableModelProbe | None",
+    caps: _CallbackCapabilities,
 ) -> bool:
     return (
         probe is not None
-        and not _is_custom_pre_call_hook(guardrail)
+        and caps.resolved_callbacks.index(guardrail) > caps.last_custom_pre_call_hook
         and guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call)  # pyright: ignore[reportUnknownMemberType]  # takes the untyped request dict
         and await probe.refuses(data)
     )
@@ -3109,7 +3112,7 @@ class ProxyLogging:
             )
 
             deferred_route_exc: SensitiveDataRouteException | None = None
-            for position, _callback in enumerate(caps.resolved_callbacks):
+            for _callback in caps.resolved_callbacks:
                 start_time = time.time()
                 try:
                     if isinstance(_callback, CustomGuardrail) and data is not None:
@@ -3123,10 +3126,11 @@ class ProxyLogging:
                         if getattr(_callback, "run_in_parallel", False):
                             continue
 
-                        if position > caps.last_custom_pre_call_hook and await _skipped_for_unroutable_model(
+                        if await _skipped_for_unroutable_model(
                             _callback,
                             data,  # pyright: ignore[reportUnknownArgumentType]  # data is the untyped request dict
                             probe,
+                            caps,
                         ):
                             data.update(without_unscanned_content(data))  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]  # data is the untyped request dict
                             continue
@@ -3205,6 +3209,7 @@ class ProxyLogging:
                                 guardrail,
                                 data,  # pyright: ignore[reportUnknownArgumentType]  # data is the untyped request dict
                                 probe,
+                                caps,
                             )
                         ]
                     ),

tests/unit/proxy/test_proxy_utils.py

@@ -8,7 +8,7 @@ import pytest
 from fastapi import HTTPException, Request
 from starlette.datastructures import State
 
-from litellm.integrations.custom_guardrail import CustomGuardrail
+from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.proxy.utils import get_docs_url, get_openapi_url, get_redoc_url
 from litellm.types.guardrails import GuardrailEventHooks
@@ -3819,17 +3819,32 @@ async def test_pre_call_hook_without_a_probe_runs_the_shipped_guardrail_for_an_u
     assert result["messages"] == _SECRET_PROMPT
 
 
-@pytest.mark.parametrize("operator_in_parallel", [False, True], ids=["sequential", "parallel"])
+@pytest.mark.parametrize("before_in_parallel", [False, True], ids=["before-sequential", "before-parallel"])
+@pytest.mark.parametrize("operator_in_parallel", [False, True], ids=["operator-sequential", "operator-parallel"])
 async def test_pre_call_hook_runs_an_operator_guardrail_and_every_guardrail_before_it(
-    monkeypatch, operator_in_parallel
+    monkeypatch, before_in_parallel, operator_in_parallel
 ):
-    before: Final = _ShippedGuardrail(name="before")
+    before: Final = _ShippedGuardrail(name="before", run_in_parallel=before_in_parallel)
     operator: Final = _OperatorGuardrail(name="operator", run_in_parallel=operator_in_parallel)
     after: Final = _ShippedGuardrail(name="after")
 
     await _walk(monkeypatch, [before, operator, after], "no-such-model", _unroutable_model_probe())
 
-    assert (before.seen, len(operator.seen), after.seen) == ([_SECRET_PROMPT], 1, [])
+    assert (len(before.seen), len(operator.seen), after.seen) == (1, 1, [])
+
+
+async def test_pre_call_hook_runs_a_custom_code_guardrail_for_a_model_routing_refuses(monkeypatch):
+    from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import CustomCodeGuardrail
+
+    operator_code: Final = CustomCodeGuardrail(
+        custom_code='def apply_guardrail(inputs, request_data, input_type):\n    return block("operator rule")\n',
+        guardrail_name="operator-code",
+        event_hook=GuardrailEventHooks.pre_call,
+        default_on=True,
+    )
+
+    with pytest.raises(ModifyResponseException, match="operator rule"):
+        await _walk(monkeypatch, [operator_code], "no-such-model", _unroutable_model_probe())
 
 
 async def test_pre_call_hook_runs_every_guardrail_when_one_routes_on_sensitive_data(monkeypatch):