diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ac42528c511..f8f58f12db5 100644 --- a/litellm/proxy/utils.py +++ b/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, ) ] ), diff --git a/tests/unit/proxy/test_proxy_utils.py b/tests/unit/proxy/test_proxy_utils.py index b1df0193762..b1e011d92ee 100644 --- a/tests/unit/proxy/test_proxy_utils.py +++ b/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):