scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

feat(guardrails): add run_in_parallel opt-in for concurrent pre_call and post_call guardrails (#33770)

Browse source 
* feat(guardrails): add run_in_parallel opt-in for concurrent pre_call guardrails

Pre-call guardrails run sequentially because each may mutate the request
payload and later guardrails depend on earlier mutations. Deployments with
several slow block-only pre_call guardrails (external moderation, Bedrock,
LLM-judge) therefore pay the sum of their latencies. during_call guardrails
run concurrently but alongside the LLM call, so a violating payload has
already been sent, which is unacceptable when the request must never reach
the model.

This adds a per-guardrail run_in_parallel flag (default off). Guardrails that
opt in are pulled out of the sequential loop and run concurrently via
asyncio.gather after every sequential (payload-mutating) guardrail has run, so
they observe the mutated payload and still form a hard barrier before the LLM
call; the first to raise blocks the request. Their returned data is discarded
since they are declared block-only.

The flag is wired from LitellmParams onto the guardrail instance at the same
generic choke point in initialize_guardrail that already sets
skip_system_message_in_guardrail, so no per-provider initializer needs to
change.

* feat(guardrails): extend run_in_parallel opt-in to post_call guardrails

post_call_success_hook ran guardrails sequentially for the same reason
pre_call did: response-modifying guardrails thread the response forward. But
block-only output scanners (which read the response and reject on violation
without changing it) serialize for no benefit and add latency.

This reuses the existing run_in_parallel flag for the post_call hook. Opted-in
post_call guardrails are pulled out of the sequential loop and run concurrently
via asyncio.gather after the sequential (response-modifying) guardrails and
before the non-guardrail CustomLogger callbacks, so they inspect the final
response and still block it from reaching the client if any raises. Their
returned response is discarded since they are block-only.

The apply_guardrail path sets data["guardrail_to_apply"] immediately before
awaiting, and unified_guardrail pops it before its first suspension point, so
concurrent guardrails never race on that key under asyncio's cooperative
scheduling.

* fix(guardrails): await all parallel guardrails and prioritize blocks over reroutes

Addresses review feedback on the run_in_parallel opt-in.

asyncio.gather propagated the first exception without cancelling or awaiting
the siblings, so a block at t=0 left the other guardrails running as
unobserved background tasks (wasted external calls plus event-loop warnings),
and a fast SensitiveDataRouteException/ModifyResponseException could return a
reroute or passthrough before a slower block finished, letting crafted input
bypass the block. Both the pre_call and post_call parallel batches now gather
with return_exceptions=True so every guardrail runs to completion, then raise
any blocking exception ahead of a flow-changing one.

The registry choke point wrote bool(None)==False onto every instance when the
config omitted run_in_parallel, silently disabling a constructor-set default;
it now only writes when the config provides an explicit value.

* fix(guardrails): record lifecycle logs for every concurrently-run guardrail

The log_guardrail_information decorator skipped its auto-record when it saw
that the count of standard_logging_guardrail_information entries in the shared
request_data had grown during the wrapped call, taking that as proof the
wrapped function had recorded its own richer entry. That heuristic breaks the
moment guardrails run concurrently (parallel pre_call/post_call, during_call):
a sibling guardrail's append inflates the shared count, so a guardrail that did
not self-record wrongly concludes it already did and drops its own entry. The
result is that enabling run_in_parallel silently loses per-guardrail lifecycle
logs, so the Admin UI Request Lifecycle timeline and downstream loggers
(Datadog, Langfuse, OTEL, spend logs) show only one of the concurrent
guardrails.

Replace the shared-count heuristic with a ContextVar flag set when a guardrail
records its own entry. asyncio copies the context into each gathered task, so
the flag is isolated per concurrent guardrail while still catching the
self-record-then-skip-auto-record case within a single invocation.

* test(guardrails): declare run_in_parallel on post_call guardrail mocks

The post_call partition reads run_in_parallel on every CustomGuardrail
callback. A MagicMock(spec=CustomGuardrail) has no run_in_parallel (it is
set in __init__, not on the class) so the attribute access raised, and even
a class-level default would return a truthy child mock that wrongly routes
the double into the parallel batch. Declare the flag False on the shared
mock factories so these pre-existing hook tests exercise the sequential
path they assert on.

* fix(guardrails): harden run_in_parallel reads and address review feedback

Read run_in_parallel via getattr(..., False) in the pre_call and post_call
partitions so a third-party CustomGuardrail subclass that overrides __init__
without chaining super().__init__() no longer raises AttributeError on a path
that previously worked. Drop the redundant in-function GuardrailEventHooks
import in _run_parallel_post_call_guardrails (already imported module-level).
Remove the flaky wall-clock upper-bound assertions from the two concurrency
tests; the all-start-before-any-end overlap assertion is the timing-independent
signal that actually proves concurrency.
This commit is contained in:
Noah Nistler 2026-07-24 15:25:58 -05:00 • committed by GitHub
parent 35dc982692
commit 8177230a29
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 736 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
litellm/integrations/custom_guardrail.py
View file

@ -1,3 +1,4 @@
import contextvars
import hashlib
import os
import secrets
@ -64,6 +65,10 @@ from litellm.exceptions import (
# proxy's metadata sanitizer.
_PRE_CALL_EXECUTED_TOKEN = secrets.token_hex(16)
_guardrail_self_recorded: contextvars.ContextVar[bool] = contextvars.ContextVar(
"litellm_guardrail_self_recorded", default=False
)
def _strict_guardrail_modes_enabled() -> bool:
"""Whether guardrail-mode validation raises (default) or logs a warning.
@ -117,6 +122,7 @@ class CustomGuardrail(CustomLogger):
on_sensitive_data: Optional[str] = None,
sensitive_data_route_to_model: Optional[str] = None,
sticky_session_routing: bool = True,
run_in_parallel: bool = False,
only_scan_new_messages: bool = False,
**kwargs,
):
@ -136,6 +142,9 @@ class CustomGuardrail(CustomLogger):
on_sensitive_data: Action when sensitive data is detected. 'block' (default) or 'route'
sensitive_data_route_to_model: Model to route to when on_sensitive_data='route'
sticky_session_routing: When True, all subsequent requests in the session use the same model
run_in_parallel: When True, this pre_call or post_call guardrail runs concurrently with
other opted-in guardrails of the same hook. Only safe for block-only guardrails that
do not mutate the request or response.
"""
self.guardrail_name = guardrail_name
self.supported_event_hooks = supported_event_hooks
@ -150,6 +159,7 @@ class CustomGuardrail(CustomLogger):
self.on_sensitive_data: Optional[str] = on_sensitive_data
self.sensitive_data_route_to_model: Optional[str] = sensitive_data_route_to_model
self.sticky_session_routing: bool = sticky_session_routing
self.run_in_parallel: bool = run_in_parallel
self.only_scan_new_messages: bool = only_scan_new_messages
if supported_event_hooks:
@ -956,6 +966,8 @@ class CustomGuardrail(CustomLogger):
request_data["metadata"] = {}
_append_guardrail_info(request_data["metadata"])
_guardrail_self_recorded.set(True)
# Emit the otel guardrail span here, where every guardrail execution lands,
# rather than relying on a post-call hook that does not fire on every path
# (e.g. a pass-through request that passes its guardrails).
@ -1238,8 +1250,12 @@ def log_guardrail_information(func):
(structured detections, tracing detail) than this decorator's
"allow"/"mask"/raw-response default. To avoid double-recording in that
case (which would emit two spans, two Datadog records, two spend-log
entries, etc.), snapshot the entry count before invocation: if the
wrapped function already appended its own entry, skip the auto-record.
entries, etc.), a context-local flag records whether the wrapped function
appended its own entry; if so, the auto-record is skipped. The flag is a
``ContextVar`` rather than a count of entries in the shared ``request_data``
so it stays correct when guardrails run concurrently (asyncio copies the
context into each gathered task): counting shared entries would let one
guardrail's append hide another guardrail's missing record.
"""
import functools
import inspect
@ -1259,16 +1275,6 @@ def log_guardrail_information(func):
return GuardrailEventHooks.post_call
return None
def _count_recorded_guardrail_entries(request_data: dict) -> int:
total = 0
for container_key in ("metadata", "litellm_metadata"):
container = request_data.get(container_key)
if isinstance(container, dict):
entries = container.get("standard_logging_guardrail_information")
if isinstance(entries, list):
total += len(entries)
return total
@functools.wraps(func)
async def async_wrapper(*args, **kwargs):
start_time = datetime.now() # Move start_time inside the wrapper
@ -1282,10 +1288,10 @@ def log_guardrail_information(func):
original_inputs = kwargs.get("inputs")
logging_obj = kwargs.get("logging_obj")
entries_before = _count_recorded_guardrail_entries(request_data)
self_recorded_token = _guardrail_self_recorded.set(False)
try:
response = await func(*args, **kwargs)
if _count_recorded_guardrail_entries(request_data) > entries_before:
if _guardrail_self_recorded.get():
return response
return self._process_response(
response=response,
@ -1297,7 +1303,7 @@ def log_guardrail_information(func):
original_inputs=original_inputs,
)
except Exception as e:
if _count_recorded_guardrail_entries(request_data) > entries_before:
if _guardrail_self_recorded.get():
raise
return self._process_error(
e=e,
@ -1308,6 +1314,7 @@ def log_guardrail_information(func):
event_type=event_type,
)
finally:
_guardrail_self_recorded.reset(self_recorded_token)
_sync_guardrail_info_to_logging_obj(request_data, logging_obj)
@functools.wraps(func)
@ -1323,10 +1330,10 @@ def log_guardrail_information(func):
original_inputs = kwargs.get("inputs")
logging_obj = kwargs.get("logging_obj")
entries_before = _count_recorded_guardrail_entries(request_data)
self_recorded_token = _guardrail_self_recorded.set(False)
try:
response = func(*args, **kwargs)
if _count_recorded_guardrail_entries(request_data) > entries_before:
if _guardrail_self_recorded.get():
return response
return self._process_response(
response=response,
@ -1336,7 +1343,7 @@ def log_guardrail_information(func):
original_inputs=original_inputs,
)
except Exception as e:
if _count_recorded_guardrail_entries(request_data) > entries_before:
if _guardrail_self_recorded.get():
raise
return self._process_error(
e=e,
@ -1345,6 +1352,7 @@ def log_guardrail_information(func):
event_type=event_type,
)
finally:
_guardrail_self_recorded.reset(self_recorded_token)
_sync_guardrail_info_to_logging_obj(request_data, logging_obj)
@functools.wraps(func)

3
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -489,6 +489,9 @@ class InMemoryGuardrailHandler:
"skip_tool_message_in_guardrail",
getattr(litellm_params, "skip_tool_message_in_guardrail", None),
)
configured_run_in_parallel = getattr(litellm_params, "run_in_parallel", None)
if configured_run_in_parallel is not None:
custom_guardrail_callback.run_in_parallel = bool(configured_run_in_parallel)
parsed_guardrail = Guardrail(
guardrail_id=guardrail.get("guardrail_id"),

135
litellm/proxy/utils.py
View file

@ -1400,6 +1400,14 @@ class ProxyLogging:
self._process_guardrail_metadata(data)
return data
parallel_guardrails: tuple[CustomGuardrail, ...] = tuple(
cb
for cb in caps.resolved_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: Optional[SensitiveDataRouteException] = None
for _callback in caps.resolved_callbacks:
start_time = time.time()
@ -1409,6 +1417,9 @@ class ProxyLogging:
if _callback.guardrail_name and _callback.guardrail_name in pipeline_managed:
continue
if getattr(_callback, "run_in_parallel", False):
continue
result = await self._process_guardrail_callback(
callback=_callback,
data=data, # type: ignore
@ -1465,6 +1476,14 @@ 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)
if parallel_guardrails and data is not None:
await self._run_parallel_pre_call_guardrails(
guardrails=parallel_guardrails,
data=data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
)
if data is not None:
self._process_guardrail_metadata(data)
@ -1477,6 +1496,47 @@ class ProxyLogging:
except Exception as e:
raise e
async def _run_parallel_pre_call_guardrails(
self,
guardrails: tuple[CustomGuardrail, ...],
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: CallTypesLiteral,
) -> None:
"""
Run opted-in pre_call guardrails concurrently against one shared payload
snapshot. These guardrails are declared block-only, so any modified data
they return is discarded; they run for their blocking side effect (raising
to reject the request before it reaches the LLM). Every guardrail is
awaited to completion (``return_exceptions=True``) so a raise by one never
leaves the others running as unobserved background tasks. A guardrail that
blocks (any exception other than a reroute or passthrough) takes precedence
over one that only changes the request flow, so a fast reroute can never
let a slower block be bypassed; the request is rejected before it reaches
the LLM, preserving the pre-call barrier that ``during_call`` guardrails
cannot provide. Per-guardrail latency is recorded by
``_process_guardrail_callback``'s own metrics.
"""
results = await asyncio.gather(
*(
self._process_guardrail_callback(
callback=callback,
data=data,
user_api_key_dict=user_api_key_dict,
call_type=call_type,
event_type=GuardrailEventHooks.pre_call,
)
for callback in guardrails
),
return_exceptions=True,
)
raised = tuple(result for result in results if isinstance(result, BaseException))
blocking = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None)
if blocking is not None:
raise blocking
if raised:
raise raised[0]
async def _handle_sensitive_data_route_exception(
self,
exc: SensitiveDataRouteException,
@ -2277,9 +2337,16 @@ class ProxyLogging:
# Merge model-level guardrails before checking which guardrails to run
guardrail_data = _check_and_merge_model_level_guardrails(data=data, llm_router=llm_router)
parallel_guardrails: tuple[CustomGuardrail, ...] = tuple(
callback for callback in guardrail_callbacks if getattr(callback, "run_in_parallel", False)
)
for callback in guardrail_callbacks:
# Main - V2 Guardrails implementation
if getattr(callback, "run_in_parallel", False):
continue
if (
callback.should_run_guardrail(
data=guardrail_data,
@ -2316,6 +2383,15 @@ class ProxyLogging:
if guardrail_response is not None:
response = guardrail_response
if parallel_guardrails:
await self._run_parallel_post_call_guardrails(
guardrails=parallel_guardrails,
data=data,
guardrail_data=guardrail_data,
response=response,
user_api_key_dict=user_api_key_dict,
)
############ Handle CustomLogger ###############################
#################################################################
@ -2329,6 +2405,65 @@ class ProxyLogging:
raise e
return response
async def _run_parallel_post_call_guardrails(
self,
guardrails: tuple[CustomGuardrail, ...],
data: dict,
guardrail_data: dict,
response: LLMResponseTypes,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""
Run opted-in post_call guardrails concurrently against the response
produced by the sequential guardrails. These guardrails are declared
block-only, so any modified response they return is discarded; they run
for their blocking side effect (raising to reject the response before it
reaches the client). Every guardrail is awaited to completion
(``return_exceptions=True``) so a raise by one never leaves the others
running as unobserved background tasks. A guardrail that blocks (any
exception other than a passthrough) takes precedence over one that only
changes the response flow, so a fast passthrough can never let a slower
block be bypassed. Each per-guardrail coroutine sets ``guardrail_to_apply``
immediately before awaiting, and the unified hook pops it before its first
suspension point, so concurrent guardrails never race on that key.
"""
async def _run_one(callback: CustomGuardrail) -> None:
if callback.should_run_guardrail(data=guardrail_data, event_type=GuardrailEventHooks.post_call) is not True:
return
if "apply_guardrail" in type(callback).__dict__:
data["guardrail_to_apply"] = callback
await self._run_guardrail_with_metrics(
callback,
unified_guardrail.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict,
data=data,
response=response,
),
"post_call",
)
else:
await self._run_guardrail_with_metrics(
callback,
callback.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict,
data=data,
response=response,
),
"post_call",
)
results = await asyncio.gather(
*(_run_one(callback) for callback in guardrails),
return_exceptions=True,
)
raised = tuple(result for result in results if isinstance(result, BaseException))
blocking = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None)
if blocking is not None:
raise blocking
if raised:
raise raised[0]
async def post_call_response_headers_hook(
self,
data: dict,

11
litellm/types/guardrails.py
View file

@ -910,6 +910,17 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up
),
)
run_in_parallel: Optional[bool] = Field(
default=None,
description=(
"When True, this pre_call or post_call guardrail runs concurrently with other opted-in "
"guardrails of the same hook, after the sequential guardrails have run. Use only for "
"block-only guardrails that inspect and reject; do not enable it for guardrails that "
"modify the request or response (e.g. PII masking or sensitive-data routing), since "
"parallel runs share one snapshot and their mutations would race."
),
)
@field_validator(
"mode",
"default_action",

4
ruff-strict-budget.json
View file

@ -222,7 +222,7 @@
"limit": 38
},
"RET504": {
"limit": 721
"limit": 719
},
"RUF010": {
"limit": 874
@ -324,7 +324,7 @@
"limit": 883
},
"UP006": {
"limit": 12792
"limit": 12789
},
"UP007": {
"limit": 2570

461
tests/proxy_unit_tests/test_proxy_utils.py
View file

@ -7,10 +7,12 @@ from typing import Any, Dict, List, Optional, Union
from unittest.mock import Mock
import pytest
from fastapi import Request
from fastapi import HTTPException, Request
from starlette.datastructures import State
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy.utils import _get_docs_url, _get_openapi_url, _get_redoc_url
from litellm.types.guardrails import GuardrailEventHooks
sys.path.insert(
0, os.path.abspath("../..")
@ -2638,6 +2640,463 @@ async def test_during_call_hook_parallel_execution_with_error():
litellm.callbacks = original_callbacks
class _PreCallGuardrail(CustomGuardrail):
"""Test double for pre_call guardrails; records timing and observed payload."""
def __init__(self, name, run_in_parallel, execution_order, sleep=0.1, default_on=True):
super().__init__(
guardrail_name=name,
event_hook=GuardrailEventHooks.pre_call,
default_on=default_on,
run_in_parallel=run_in_parallel,
)
self.name = name
self.sleep = sleep
self.execution_order = execution_order
self.observed_content = None
self.was_called = False
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
self.was_called = True
self.observed_content = data["messages"][0]["content"]
self.execution_order.append(f"{self.name}_start")
await asyncio.sleep(self.sleep)
self.execution_order.append(f"{self.name}_end")
return None
@pytest.mark.asyncio
async def test_pre_call_hook_runs_opted_in_guardrails_in_parallel():
"""run_in_parallel pre_call guardrails execute concurrently (all start before any ends)."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
execution_order = []
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
try:
litellm.callbacks = [
_PreCallGuardrail(f"g{i}", run_in_parallel=True, execution_order=execution_order) for i in range(3)
]
result = await proxy_logging.pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
call_type="completion",
)
first_end_idx = next(i for i, item in enumerate(execution_order) if "end" in item)
starts_before_first_end = sum(1 for item in execution_order[:first_end_idx] if "start" in item)
assert starts_before_first_end == 3, f"expected 3 concurrent starts, got {starts_before_first_end}"
assert result["model"] == "gpt-4"
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_pre_call_hook_runs_default_guardrails_sequentially():
"""Guardrails without run_in_parallel keep the sequential, one-at-a-time behavior."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
execution_order = []
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
try:
litellm.callbacks = [
_PreCallGuardrail(f"g{i}", run_in_parallel=False, execution_order=execution_order) for i in range(2)
]
start = asyncio.get_event_loop().time()
await proxy_logging.pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
call_type="completion",
)
elapsed = asyncio.get_event_loop().time() - start
assert execution_order == ["g0_start", "g0_end", "g1_start", "g1_end"]
assert elapsed >= 0.18, f"sequential run took {elapsed}s, expected >= 0.18s"
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_pre_call_hook_sequential_mutations_precede_parallel_batch():
"""Sequential (mutating) guardrails run before the parallel batch, which sees their changes."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
execution_order = []
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
class MaskingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="masker",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=False,
)
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
data["messages"][0]["content"] = "MASKED"
return data
parallel_observer = _PreCallGuardrail("observer", run_in_parallel=True, execution_order=execution_order)
try:
litellm.callbacks = [parallel_observer, MaskingGuardrail()]
await proxy_logging.pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={"model": "gpt-4", "messages": [{"role": "user", "content": "secret"}]},
call_type="completion",
)
assert parallel_observer.observed_content == "MASKED"
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_pre_call_hook_parallel_guardrail_blocks_request():
"""A raising parallel guardrail blocks the request before it reaches the LLM."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
class BlockingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="blocker",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=True,
)
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
raise HTTPException(status_code=400, detail="blocked by guardrail")
try:
litellm.callbacks = [BlockingGuardrail()]
with pytest.raises(HTTPException) as exc_info:
await proxy_logging.pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
call_type="completion",
)
assert exc_info.value.status_code == 400
assert "blocked by guardrail" in str(exc_info.value.detail)
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_pre_call_hook_parallel_guardrail_skipped_when_should_not_run():
"""A parallel guardrail that should_run_guardrail rejects is never invoked."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
execution_order = []
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
try:
guardrail = _PreCallGuardrail(
"off_by_default", run_in_parallel=True, execution_order=execution_order, default_on=False
)
litellm.callbacks = [guardrail]
result = await proxy_logging.pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
call_type="completion",
)
assert guardrail.was_called is False
assert result["model"] == "gpt-4"
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_pre_call_hook_parallel_block_wins_over_reroute():
"""A slower block must win over a faster reroute so crafted input cannot bypass a block."""
from litellm.caching.caching import DualCache
from litellm.exceptions import SensitiveDataRouteException
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
class FastRerouteGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="rerouter",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=True,
)
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
raise SensitiveDataRouteException(route_to_model="on-prem", session_id="s1", guardrail_name="rerouter")
class SlowBlockingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="blocker",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=True,
)
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
await asyncio.sleep(0.1)
raise HTTPException(status_code=400, detail="blocked by guardrail")
try:
litellm.callbacks = [FastRerouteGuardrail(), SlowBlockingGuardrail()]
with pytest.raises(HTTPException) as exc_info:
await proxy_logging.pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
call_type="completion",
)
assert exc_info.value.status_code == 400
assert "blocked by guardrail" in str(exc_info.value.detail)
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_pre_call_hook_parallel_awaits_all_when_one_blocks():
"""A block must not orphan sibling guardrails; every parallel guardrail runs to completion."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
completed = []
class FastBlockingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="fast_blocker",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=True,
)
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
raise HTTPException(status_code=400, detail="blocked")
class SlowGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="slow",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=True,
)
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
await asyncio.sleep(0.1)
completed.append("slow")
return None
try:
litellm.callbacks = [FastBlockingGuardrail(), SlowGuardrail()]
with pytest.raises(HTTPException):
await proxy_logging.pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
call_type="completion",
)
assert completed == ["slow"], "slow guardrail was orphaned instead of awaited to completion"
finally:
litellm.callbacks = original_callbacks
class _PostCallGuardrail(CustomGuardrail):
"""Test double for post_call guardrails; records timing and invocation."""
def __init__(self, name, run_in_parallel, execution_order, sleep=0.1, default_on=True):
super().__init__(
guardrail_name=name,
event_hook=GuardrailEventHooks.post_call,
default_on=default_on,
run_in_parallel=run_in_parallel,
)
self.name = name
self.sleep = sleep
self.execution_order = execution_order
self.was_called = False
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
self.was_called = True
self.execution_order.append(f"{self.name}_start")
await asyncio.sleep(self.sleep)
self.execution_order.append(f"{self.name}_end")
return None
@pytest.mark.asyncio
async def test_post_call_hook_runs_opted_in_guardrails_in_parallel():
"""run_in_parallel post_call guardrails execute concurrently (all start before any ends)."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
execution_order = []
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
try:
litellm.callbacks = [
_PostCallGuardrail(f"g{i}", run_in_parallel=True, execution_order=execution_order) for i in range(3)
]
await proxy_logging.post_call_success_hook(
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
response=litellm.ModelResponse(),
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
)
first_end_idx = next(i for i, item in enumerate(execution_order) if "end" in item)
starts_before_first_end = sum(1 for item in execution_order[:first_end_idx] if "start" in item)
assert starts_before_first_end == 3, f"expected 3 concurrent starts, got {starts_before_first_end}"
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_post_call_hook_runs_default_guardrails_sequentially():
"""post_call guardrails without run_in_parallel keep the sequential behavior."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
execution_order = []
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
try:
litellm.callbacks = [
_PostCallGuardrail(f"g{i}", run_in_parallel=False, execution_order=execution_order) for i in range(2)
]
start = asyncio.get_event_loop().time()
await proxy_logging.post_call_success_hook(
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
response=litellm.ModelResponse(),
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
)
elapsed = asyncio.get_event_loop().time() - start
assert execution_order == ["g0_start", "g0_end", "g1_start", "g1_end"]
assert elapsed >= 0.18, f"sequential run took {elapsed}s, expected >= 0.18s"
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_post_call_hook_parallel_guardrail_blocks_response():
"""A raising parallel post_call guardrail blocks the response before it reaches the client."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
class BlockingPostCallGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="post_blocker",
event_hook=GuardrailEventHooks.post_call,
default_on=True,
run_in_parallel=True,
)
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
raise HTTPException(status_code=400, detail="blocked response by guardrail")
try:
litellm.callbacks = [BlockingPostCallGuardrail()]
with pytest.raises(HTTPException) as exc_info:
await proxy_logging.post_call_success_hook(
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
response=litellm.ModelResponse(),
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
)
assert exc_info.value.status_code == 400
assert "blocked response by guardrail" in str(exc_info.value.detail)
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_post_call_hook_parallel_awaits_all_when_one_blocks():
"""A blocking post_call guardrail must not orphan its siblings; all run to completion."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
completed = []
class FastBlockingPostCall(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="fast_post_blocker",
event_hook=GuardrailEventHooks.post_call,
default_on=True,
run_in_parallel=True,
)
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
raise HTTPException(status_code=400, detail="blocked")
class SlowPostCall(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="slow_post",
event_hook=GuardrailEventHooks.post_call,
default_on=True,
run_in_parallel=True,
)
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
await asyncio.sleep(0.1)
completed.append("slow")
return None
try:
litellm.callbacks = [FastBlockingPostCall(), SlowPostCall()]
with pytest.raises(HTTPException):
await proxy_logging.post_call_success_hook(
data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
response=litellm.ModelResponse(),
user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
)
assert completed == ["slow"], "slow post_call guardrail was orphaned instead of awaited to completion"
finally:
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_handle_logging_proxy_only_error_preserves_pass_through_call_type():
"""Ensure _handle_logging_proxy_only_error does not overwrite call_type

33
tests/test_litellm/integrations/test_custom_guardrail.py
View file

@ -1,3 +1,4 @@
import asyncio
from unittest.mock import AsyncMock
import pytest
@ -1394,6 +1395,38 @@ class TestEventTypeLogging:
assert len(logged_info) == 1
assert logged_info[0]["guardrail_status"] == "guardrail_intervened"
@pytest.mark.asyncio
async def test_log_guardrail_information_records_every_concurrent_guardrail(self):
"""Guardrails run concurrently (parallel pre_call/post_call, during_call) share one
request_data dict. Each must still record its own entry. The previous guard counted
entries in that shared dict, so a sibling's append made a guardrail think it had already
recorded and skip its own auto-record — silently dropping lifecycle logs the UI shows."""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class SleeperGuardrail(CustomGuardrail):
def __init__(self, name, sleep):
super().__init__(guardrail_name=name, event_hook=GuardrailEventHooks.pre_call)
self._sleep = sleep
@log_guardrail_information
async def async_pre_call_hook(self, data: dict, **kwargs):
await asyncio.sleep(self._sleep)
return data
request_data = {"metadata": {}}
# Different sleeps guarantee overlapping execution windows: the faster guardrail
# records while the slower one is still awaiting, which is exactly what tripped the
# old shared-count guard.
await asyncio.gather(
SleeperGuardrail("guardrail-a", 0.05).async_pre_call_hook(data=request_data),
SleeperGuardrail("guardrail-b", 0.15).async_pre_call_hook(data=request_data),
)
logged = request_data["metadata"]["standard_logging_guardrail_information"]
assert {entry["guardrail_name"] for entry in logged} == {"guardrail-a", "guardrail-b"}
assert len(logged) == 2
def test_add_standard_logging_falls_back_to_event_hook_when_event_type_is_none(
self,
):

40
tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
View file

@ -1,3 +1,5 @@
import pytest
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy.guardrails.guardrail_registry import (
get_guardrail_initializer_from_hooks,
@ -32,6 +34,44 @@ def test_noma_registry_resolution():
assert "noma_v2" in guardrail_initializer_registry
@pytest.mark.parametrize(
"configured, expected",
[(None, True), (False, False), (True, True)],
)
def test_initialize_guardrail_run_in_parallel_preserves_constructor_default(configured, expected):
"""
A guardrail whose constructor sets run_in_parallel=True must keep that default when
the config omits the key; only an explicit config value may override it. The
previous code wrote bool(None)==False on every instance, silently disabling the
opt-in for such guardrails.
"""
from litellm.proxy.guardrails import guardrail_registry as registry_module
def _initializer(litellm_params, guardrail):
return CustomGuardrail(
guardrail_name=guardrail["guardrail_name"],
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=True,
)
registry_module.guardrail_initializer_registry["parallel_default_test"] = _initializer
try:
params = {"guardrail": "parallel_default_test", "mode": "pre_call"}
if configured is not None:
params["run_in_parallel"] = configured
handler = InMemoryGuardrailHandler()
result = handler.initialize_guardrail(
guardrail={"guardrail_name": "cf-parallel-default", "litellm_params": params},
)
stored = handler.guardrail_id_to_custom_guardrail[result["guardrail_id"]]
assert stored.run_in_parallel is expected
finally:
registry_module.guardrail_initializer_registry.pop("parallel_default_test", None)
def test_update_in_memory_guardrail():
handler = InMemoryGuardrailHandler()
handler.guardrail_id_to_custom_guardrail["123"] = CustomGuardrail(

24
tests/test_litellm/proxy/guardrails/test_init_guardrails.py
View file

@ -62,3 +62,27 @@ def test_initialize_guardrail_preserves_guardrail_info():
assert result["guardrail_info"] == {"type": "PII", "description": "masks PII"}
stored = guardrail_handler.IN_MEMORY_GUARDRAILS[result["guardrail_id"]]
assert stored["guardrail_info"] == {"type": "PII", "description": "masks PII"}
@pytest.mark.parametrize(
"config_value, expected",
[(True, True), (False, False), (None, False)],
)
def test_initialize_guardrail_sets_run_in_parallel(config_value, expected):
"""run_in_parallel from litellm_params must reach the built guardrail instance."""
litellm_params = {
"guardrail": SupportedGuardrailIntegrations.PRESIDIO.value,
"mode": "pre_call",
"presidio_analyzer_api_base": "https://fakelink.com/v1/presidio/analyze",
"presidio_anonymizer_api_base": "https://fakelink.com/v1/presidio/anonymize",
}
if config_value is not None:
litellm_params["run_in_parallel"] = config_value
guardrail_handler = InMemoryGuardrailHandler()
result = guardrail_handler.initialize_guardrail(
guardrail={"guardrail_name": "test_parallel_flag", "litellm_params": litellm_params},
)
custom_guardrail = guardrail_handler.guardrail_id_to_custom_guardrail[result["guardrail_id"]]
assert custom_guardrail.run_in_parallel is expected

1
tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py
View file

@ -626,6 +626,7 @@ def _moderation_guardrail() -> MagicMock:
cb.should_run_guardrail = MagicMock(return_value=True)
cb.async_moderation_hook = AsyncMock(return_value=None)
cb.async_post_call_success_hook = AsyncMock(return_value=None)
cb.run_in_parallel = False
return cb

1
tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py
View file

@ -28,6 +28,7 @@ def _make_guardrail(name="g", should_run=True, override=None):
cb.event_hook = GuardrailEventHooks.post_call
cb.should_run_guardrail = MagicMock(return_value=should_run)
cb.async_post_call_success_hook = AsyncMock(return_value=override)
cb.run_in_parallel = False
return cb