mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
fix(guardrails): run pre_call hook once for model-level guardrails (#30543)
* fix(guardrails): run pre_call hook once for model-level guardrails A CustomGuardrail attached to a deployment via litellm_params.guardrails gets its async_pre_call_hook invoked twice per request: once by the proxy pre-call loop and again by async_pre_call_deployment_hook after the router spreads the model-level guardrails into the top-level request kwargs. Record in request metadata that the proxy pre-call loop already ran a given guardrail, and have the deployment hook skip it when the marker is present. Direct-SDK usage never runs the proxy loop, so the deployment hook stays the sole invocation there and still fires exactly once. The marker key is stripped from untrusted caller metadata so a request body cannot suppress a model-only guardrail by pre-seeding it. * fix(guardrails): mark pre_call dedup on the post-hook request data Record the exactly-once marker after async_pre_call_hook runs, on the data object that flows downstream, rather than before it. A guardrail whose hook returns a brand-new request dict (instead of mutating or spreading the one it received) would otherwise discard the marker, letting the deployment hook re-run the guardrail a second time.
This commit is contained in:
parent
bed6ce820c
commit
4faeabc254
8 changed files with 325 additions and 1 deletions
|
|
@ -190,6 +190,10 @@ DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE = int(
|
|||
# Override with LITELLM_MAX_CALLBACKS env var for large deployments (e.g., many teams with guardrails)
|
||||
MAX_CALLBACKS = get_env_int("LITELLM_MAX_CALLBACKS", 100)
|
||||
|
||||
# Metadata key recording which pre_call guardrails the proxy loop already ran,
|
||||
# so the deployment-level hook does not re-run them for the same request
|
||||
PRE_CALL_EXECUTED_GUARDRAILS_KEY = "_pre_call_executed_guardrails"
|
||||
|
||||
# Generic fallback for unknown models
|
||||
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET = int(
|
||||
os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128)
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import secrets
|
||||
from datetime import datetime
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
|
|
@ -43,6 +44,7 @@ if TYPE_CHECKING:
|
|||
dc = DualCache()
|
||||
|
||||
|
||||
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
|
||||
from litellm.exceptions import (
|
||||
BlockedPiiEntityError,
|
||||
GuardrailRaisedException,
|
||||
|
|
@ -50,6 +52,12 @@ from litellm.exceptions import (
|
|||
SensitiveDataRouteException,
|
||||
)
|
||||
|
||||
# Per-process secret tagging each recorded marker. The deployment hook only
|
||||
# honors markers carrying this token, so a caller cannot forge the metadata
|
||||
# field to suppress a guardrail on the direct-SDK path that never reaches the
|
||||
# proxy's metadata sanitizer.
|
||||
_PRE_CALL_EXECUTED_TOKEN = secrets.token_hex(16)
|
||||
|
||||
|
||||
def get_session_id_from_request_data(request_data: Dict[str, Any]) -> Optional[str]:
|
||||
"""Extract session_id from request data (litellm_session_id or metadata)."""
|
||||
|
|
@ -458,6 +466,49 @@ class CustomGuardrail(CustomLogger):
|
|||
|
||||
return False
|
||||
|
||||
def _pre_call_marker(self) -> Optional[str]:
|
||||
name = self.guardrail_name
|
||||
if not name:
|
||||
return None
|
||||
return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}"
|
||||
|
||||
def mark_pre_call_hook_ran(self, data: Dict[str, Any]) -> None:
|
||||
"""
|
||||
Record that this guardrail's ``async_pre_call_hook`` already ran for this
|
||||
request, so the deployment-level hook does not run it a second time.
|
||||
|
||||
The proxy runs pre-call guardrails in ``ProxyLogging.pre_call_hook``. The
|
||||
router later spreads a deployment's model-level ``guardrails`` into the
|
||||
top-level request kwargs, which would otherwise re-trigger the same hook
|
||||
from ``async_pre_call_deployment_hook``.
|
||||
"""
|
||||
marker = self._pre_call_marker()
|
||||
if marker is None:
|
||||
return
|
||||
for meta_key in ("metadata", "litellm_metadata"):
|
||||
meta = data.get(meta_key)
|
||||
if isinstance(meta, dict):
|
||||
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
|
||||
if isinstance(executed, list):
|
||||
if marker not in executed:
|
||||
executed.append(marker)
|
||||
else:
|
||||
meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker]
|
||||
return
|
||||
data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
|
||||
|
||||
def _pre_call_hook_already_ran(self, data: Dict[str, Any]) -> bool:
|
||||
marker = self._pre_call_marker()
|
||||
if marker is None:
|
||||
return False
|
||||
for meta_key in ("metadata", "litellm_metadata"):
|
||||
meta = data.get(meta_key)
|
||||
if isinstance(meta, dict):
|
||||
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
|
||||
if isinstance(executed, list) and marker in executed:
|
||||
return True
|
||||
return False
|
||||
|
||||
async def async_pre_call_deployment_hook(
|
||||
self, kwargs: Dict[str, Any], call_type: Optional[CallTypes]
|
||||
) -> Optional[dict]:
|
||||
|
|
@ -468,6 +519,9 @@ class CustomGuardrail(CustomLogger):
|
|||
if litellm_guardrails is None or not isinstance(litellm_guardrails, list):
|
||||
return kwargs
|
||||
|
||||
if self._pre_call_hook_already_ran(kwargs):
|
||||
return kwargs
|
||||
|
||||
if (
|
||||
self.should_run_guardrail(
|
||||
data=kwargs, event_type=GuardrailEventHooks.pre_call
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Literal,
|
|||
import litellm
|
||||
from litellm import get_secret
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
||||
from litellm.proxy._types import CommonProxyErrors, LiteLLMPromptInjectionParams
|
||||
|
|
@ -497,6 +498,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS = frozenset(
|
|||
"guardrail_config",
|
||||
"_guardrail_pipelines",
|
||||
"_pipeline_managed_guardrails",
|
||||
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
|
||||
"disable_global_guardrails",
|
||||
"disable_global_guardrail",
|
||||
"opted_out_global_guardrails",
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from starlette.datastructures import Headers
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger, verbose_proxy_logger
|
||||
from litellm._service_logger import ServiceLogging
|
||||
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.url_utils import is_url_destination_allowed_by_host
|
||||
|
|
@ -161,6 +162,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS = (
|
|||
"secret_fields",
|
||||
"_guardrail_pipelines",
|
||||
"_pipeline_managed_guardrails",
|
||||
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
|
||||
)
|
||||
|
||||
_UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS = frozenset(
|
||||
|
|
|
|||
|
|
@ -171,6 +171,10 @@ class PipelineExecutor:
|
|||
data=data,
|
||||
call_type=call_type, # type: ignore
|
||||
)
|
||||
if isinstance(callback, CustomGuardrail):
|
||||
callback.mark_pre_call_hook_ran(data)
|
||||
if isinstance(response, dict):
|
||||
callback.mark_pre_call_hook_ran(response)
|
||||
elif mode == "post_call":
|
||||
response = await target.async_post_call_success_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
|
|||
|
|
@ -1171,6 +1171,8 @@ class ProxyLogging:
|
|||
response=response, data=data, call_type=call_type
|
||||
)
|
||||
|
||||
callback.mark_pre_call_hook_ran(data)
|
||||
|
||||
except SensitiveDataRouteException:
|
||||
status = "intervened"
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -84,6 +84,112 @@ class TestCustomGuardrailDeploymentHook:
|
|||
assert result["messages"] == mock_result["messages"]
|
||||
assert result["messages"] != original_messages
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deployment_hook_skips_when_pre_call_already_ran(self):
|
||||
"""The deployment hook must not re-run async_pre_call_hook once the proxy
|
||||
pre-call loop has already run it for this request."""
|
||||
|
||||
class CountingGuardrail(CustomGuardrail):
|
||||
def __init__(self):
|
||||
super().__init__(guardrail_name="g1", default_on=True)
|
||||
self.pre_call_count = 0
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self, user_api_key_dict, cache, data, call_type
|
||||
):
|
||||
self.pre_call_count += 1
|
||||
return data
|
||||
|
||||
guardrail = CountingGuardrail()
|
||||
kwargs = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"model": "gpt-3.5-turbo",
|
||||
"guardrails": ["g1"],
|
||||
"metadata": {},
|
||||
}
|
||||
|
||||
guardrail.mark_pre_call_hook_ran(kwargs)
|
||||
await guardrail.async_pre_call_deployment_hook(
|
||||
kwargs=kwargs, call_type=CallTypes.completion
|
||||
)
|
||||
|
||||
assert guardrail.pre_call_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deployment_hook_runs_when_not_marked(self):
|
||||
"""Without the proxy marker (direct-SDK usage) the deployment hook is the
|
||||
only execution path and must still run the guardrail exactly once."""
|
||||
|
||||
class CountingGuardrail(CustomGuardrail):
|
||||
def __init__(self):
|
||||
super().__init__(guardrail_name="g1", default_on=True)
|
||||
self.pre_call_count = 0
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self, user_api_key_dict, cache, data, call_type
|
||||
):
|
||||
self.pre_call_count += 1
|
||||
return data
|
||||
|
||||
guardrail = CountingGuardrail()
|
||||
kwargs = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"model": "gpt-3.5-turbo",
|
||||
"guardrails": ["g1"],
|
||||
"metadata": {},
|
||||
}
|
||||
|
||||
await guardrail.async_pre_call_deployment_hook(
|
||||
kwargs=kwargs, call_type=CallTypes.completion
|
||||
)
|
||||
|
||||
assert guardrail.pre_call_count == 1
|
||||
|
||||
def test_mark_pre_call_hook_ran_uses_litellm_metadata(self):
|
||||
"""The marker is recorded in litellm_metadata when that is the metadata
|
||||
bucket in use, and is then visible to the skip check."""
|
||||
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
|
||||
|
||||
guardrail = CustomGuardrail(guardrail_name="g1")
|
||||
kwargs = {"litellm_metadata": {}}
|
||||
|
||||
guardrail.mark_pre_call_hook_ran(kwargs)
|
||||
|
||||
assert kwargs["litellm_metadata"][PRE_CALL_EXECUTED_GUARDRAILS_KEY]
|
||||
assert guardrail._pre_call_hook_already_ran(kwargs) is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deployment_hook_ignores_forged_caller_marker(self):
|
||||
"""A direct-SDK caller controls request metadata but cannot know the
|
||||
per-process token, so a hand-crafted marker must not suppress a
|
||||
requested guardrail in async_pre_call_deployment_hook."""
|
||||
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
|
||||
|
||||
class CountingGuardrail(CustomGuardrail):
|
||||
def __init__(self):
|
||||
super().__init__(guardrail_name="g1", default_on=True)
|
||||
self.pre_call_count = 0
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self, user_api_key_dict, cache, data, call_type
|
||||
):
|
||||
self.pre_call_count += 1
|
||||
return data
|
||||
|
||||
guardrail = CountingGuardrail()
|
||||
kwargs = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"model": "gpt-3.5-turbo",
|
||||
"guardrails": ["g1"],
|
||||
"metadata": {PRE_CALL_EXECUTED_GUARDRAILS_KEY: ["g1"]},
|
||||
}
|
||||
|
||||
await guardrail.async_pre_call_deployment_hook(
|
||||
kwargs=kwargs, call_type=CallTypes.completion
|
||||
)
|
||||
|
||||
assert guardrail.pre_call_count == 1
|
||||
|
||||
|
||||
class TestCustomGuardrailShouldRunGuardrail:
|
||||
|
||||
|
|
|
|||
|
|
@ -19,7 +19,6 @@ from litellm.proxy.utils import (
|
|||
_merge_guardrails_with_existing,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests for _check_and_merge_model_level_guardrails
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -159,6 +158,157 @@ class TestCheckAndMergeModelLevelGuardrails:
|
|||
assert "existing" in result["metadata"]["guardrails"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Regression test: pre_call hook must run exactly once with model-level guardrails
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_runs_once_with_model_level_guardrails():
|
||||
"""
|
||||
A guardrail attached at the model level (litellm_params.guardrails) is
|
||||
spread into the top-level request kwargs by the router. The proxy pre-call
|
||||
loop (async_pre_call_hook) and the deployment-level hook
|
||||
(async_pre_call_deployment_hook) must together invoke async_pre_call_hook
|
||||
exactly once, not twice.
|
||||
"""
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
class CountingGuardrail(CustomGuardrail):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
guardrail_name="counting-guardrail",
|
||||
event_hook=GuardrailEventHooks.pre_call,
|
||||
default_on=True,
|
||||
)
|
||||
self.pre_call_count = 0
|
||||
|
||||
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
||||
self.pre_call_count += 1
|
||||
return data
|
||||
|
||||
guardrail = CountingGuardrail()
|
||||
|
||||
with patch("litellm.callbacks", [guardrail]):
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
||||
|
||||
data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"metadata": {},
|
||||
}
|
||||
|
||||
# Path A: proxy pre-call loop runs the guardrail and records that it ran
|
||||
data = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
call_type="acompletion",
|
||||
)
|
||||
|
||||
# Path B: the router spreads the deployment's model-level guardrails into
|
||||
# the top-level kwargs, then litellm.acompletion fires the deployment hook
|
||||
data["guardrails"] = ["counting-guardrail"]
|
||||
await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion)
|
||||
|
||||
assert guardrail.pre_call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_runs_once_when_hook_returns_fresh_dict():
|
||||
"""
|
||||
async_pre_call_hook may return a brand-new request dict instead of mutating
|
||||
or spreading the one it received. The exactly-once marker must live on the
|
||||
data that flows downstream, so the deployment hook still skips the guardrail
|
||||
even when the proxy loop swapped in a fresh dict that never carried it.
|
||||
"""
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
class FreshDictGuardrail(CustomGuardrail):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
guardrail_name="counting-guardrail",
|
||||
event_hook=GuardrailEventHooks.pre_call,
|
||||
default_on=True,
|
||||
)
|
||||
self.pre_call_count = 0
|
||||
|
||||
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
||||
self.pre_call_count += 1
|
||||
return {"model": data["model"], "messages": data["messages"]}
|
||||
|
||||
guardrail = FreshDictGuardrail()
|
||||
|
||||
with patch("litellm.callbacks", [guardrail]):
|
||||
ProxyLogging._callback_capabilities_cache.clear()
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
||||
|
||||
data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"metadata": {},
|
||||
}
|
||||
|
||||
data = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=data,
|
||||
call_type="acompletion",
|
||||
)
|
||||
|
||||
data["guardrails"] = ["counting-guardrail"]
|
||||
await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion)
|
||||
|
||||
assert guardrail.pre_call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deployment_hook_runs_pre_call_without_proxy_loop():
|
||||
"""
|
||||
Direct-SDK usage (litellm.acompletion(..., guardrails=[...]) without the
|
||||
proxy) never runs the proxy pre-call loop, so the deployment hook is the
|
||||
only place the guardrail executes and it must still run exactly once.
|
||||
"""
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.proxy._types import CallTypes
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
class CountingGuardrail(CustomGuardrail):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
guardrail_name="counting-guardrail",
|
||||
event_hook=GuardrailEventHooks.pre_call,
|
||||
default_on=True,
|
||||
)
|
||||
self.pre_call_count = 0
|
||||
|
||||
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
||||
self.pre_call_count += 1
|
||||
return data
|
||||
|
||||
guardrail = CountingGuardrail()
|
||||
|
||||
data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"guardrails": ["counting-guardrail"],
|
||||
"metadata": {},
|
||||
}
|
||||
|
||||
await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion)
|
||||
|
||||
assert guardrail.pre_call_count == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Integration test: post_call_success_hook with model-level guardrails
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue