scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(guardrails): keep guardrail telemetry when a policy pipeline blocks or modifies the response (#40211)

Browse source 
* fix(guardrails): keep guardrail telemetry when a policy pipeline blocks or modifies the response

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(guardrails): count every raw-snapshot guardrail evaluation and type the telemetry carry helpers

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(policy_engine): type the recording guardrail hooks and telemetry test parameters

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-08 11:45:02 -07:00 • committed by GitHub
parent 82e6b84f5a
commit ee1a6407cb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 216 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

104
litellm/proxy/policy_engine/pipeline_executor.py
View file

@ -6,7 +6,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding.
"""
import time
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from typing import Any, Final, Literal
import litellm
@ -16,7 +16,11 @@ from litellm.integrations.custom_guardrail import (
ModifyResponseException,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import independent_snapshot
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
get_or_create_metadata_bucket,
independent_snapshot,
)
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
@ -25,6 +29,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import (
PipelineStep,
PipelineStepResult,
)
from litellm.types.utils import StandardLoggingGuardrailInformation
try:
from fastapi.exceptions import HTTPException
@ -118,6 +123,7 @@ class PipelineExecutor:
return _allow_result(step_results=step_results, working_data=working_data, request_data=data)
if action == "block":
_carry_working_guardrail_information(working_data=working_data, request_data=data)
return PipelineExecutionResult(
terminal_action="block",
step_results=step_results,
@ -126,6 +132,7 @@ class PipelineExecutor:
)
if action == "modify_response":
_carry_working_guardrail_information(working_data=working_data, request_data=data)
return PipelineExecutionResult(
terminal_action="modify_response",
step_results=step_results,
@ -168,34 +175,33 @@ class PipelineExecutor:
verbose_proxy_logger.warning("Pipeline: guardrail '%s' not found in callbacks", step.guardrail)
return ("error", None, f"Guardrail '{step.guardrail}' not found", None)
# Inject guardrail name into metadata so should_run_guardrail() allows it
if "metadata" not in data:
data["metadata"] = {}
data["metadata"]["guardrails"] = [step.guardrail]
# A scan_raw_request step evaluates the pristine pre-pipeline
# snapshot instead of `data` (which earlier pass_data steps in
# this same pipeline may have already rewritten), same reason
# the normal sequential/parallel guardrail loops do this.
scans_raw_request: Final = callback.scan_raw_request
hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data
independent_snapshot(raw_request_snapshot)
if scans_raw_request and raw_request_snapshot is not None
else data
)
if hook_input is not data:
hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail]
snapshot_entries_before: Final = len(_recorded_guardrail_information(hook_input))
# Use unified_guardrail path if callback implements apply_guardrail
target: CustomLogger = callback
use_unified: Final = "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks
if use_unified:
hook_input["guardrail_to_apply"] = callback
target = UnifiedLLMGuardrails()
try:
# Inject guardrail name into metadata so should_run_guardrail() allows it
if "metadata" not in data:
data["metadata"] = {}
data["metadata"]["guardrails"] = [step.guardrail]
# A scan_raw_request step evaluates the pristine pre-pipeline
# snapshot instead of `data` (which earlier pass_data steps in
# this same pipeline may have already rewritten), same reason
# the normal sequential/parallel guardrail loops do this.
scans_raw_request: Final = callback.scan_raw_request
hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data
independent_snapshot(raw_request_snapshot)
if scans_raw_request and raw_request_snapshot is not None
else data
)
if hook_input is not data:
hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail]
# Use unified_guardrail path if callback implements apply_guardrail
target: CustomLogger = callback
use_unified: Final = (
"apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks
)
if use_unified:
hook_input["guardrail_to_apply"] = callback
target = UnifiedLLMGuardrails()
if mode == "pre_call":
response = await target.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
@ -233,6 +239,12 @@ class PipelineExecutor:
else:
verbose_proxy_logger.error("Pipeline: unexpected error from guardrail '%s': %s", step.guardrail, e)
return ("error", None, str(e), e)
finally:
if hook_input is not data:
_append_guardrail_information(
request_data=data,
entries=_recorded_guardrail_information(hook_input)[snapshot_entries_before:],
)
@staticmethod
def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None:
@ -283,6 +295,40 @@ def _restore_request_guardrails(
return {**working_data, "metadata": stripped} # mutable-ok: request dict
_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information"
def _recorded_guardrail_information(source: Mapping[str, object]) -> list[StandardLoggingGuardrailInformation]:
bucket: Final = source.get(get_metadata_variable_name_from_kwargs(source))
recorded: Final = bucket.get(_GUARDRAIL_INFORMATION_KEY) if isinstance(bucket, dict) else None
return recorded if isinstance(recorded, list) else []
def _append_guardrail_information(
request_data: dict[str, object], # mutable-ok: same request-payload shape as execute_steps' data
entries: Sequence[StandardLoggingGuardrailInformation],
) -> None:
if not entries:
return
_, request_bucket = get_or_create_metadata_bucket(request_data)
existing: Final = request_bucket.get(_GUARDRAIL_INFORMATION_KEY)
if isinstance(existing, list):
existing.extend(entries)
return
request_bucket[_GUARDRAIL_INFORMATION_KEY] = list(entries)
def _carry_working_guardrail_information(
working_data: Mapping[str, object],
request_data: dict[str, object], # mutable-ok: same request-payload shape as execute_steps' data
) -> None:
recorded: Final = _recorded_guardrail_information(working_data)
existing: Final = _recorded_guardrail_information(request_data)
if recorded is existing:
return
_append_guardrail_information(request_data=request_data, entries=[e for e in recorded if e not in existing])
def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str:
"""
Map pipeline step outcome to the configured action.

141
tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py
View file

@ -4,20 +4,26 @@ Tests for the pipeline executor.
Uses mock guardrails to validate pipeline execution without external services.
"""
import copy
from typing import Literal
from unittest.mock import MagicMock
import pytest
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
CustomCodeGuardrail,
)
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.policy_engine.pipeline_types import (
GuardrailPipeline,
PipelineStep,
)
from litellm.types.utils import CallTypesLiteral
try:
from fastapi.exceptions import HTTPException
@ -158,11 +164,146 @@ class ContentCheckGuardrail(CustomGuardrail):
return None
class RecordingGuardrail(CustomGuardrail):
def __init__(self, guardrail_name: str, scan_raw_request: bool = False, block: bool = True):
super().__init__(
guardrail_name=guardrail_name,
event_hook="pre_call",
default_on=True,
scan_raw_request=scan_raw_request,
)
self.block = block
def should_run_guardrail(self, data: dict[str, object], event_type: GuardrailEventHooks) -> bool:
return True
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict[str, object],
call_type: CallTypesLiteral,
) -> dict[str, object]:
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={"detected": ["aws_access_key"]},
request_data=data,
guardrail_status="guardrail_intervened" if self.block else "success",
)
if self.block:
raise HTTPException(status_code=400, detail="Content policy violation")
return copy.deepcopy(data)
# ─────────────────────────────────────────────────────────────────────────────
# Tests
# ─────────────────────────────────────────────────────────────────────────────
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
@pytest.mark.parametrize("scan_raw_request", [False, True])
@pytest.mark.parametrize("on_fail", ["block", "modify_response"])
async def test_terminal_block_carries_guardrail_information_to_request(
monkeypatch: pytest.MonkeyPatch, scan_raw_request: bool, on_fail: Literal["block", "modify_response"]
):
"""
Spend logging and the Guardrails Monitor read standard_logging_guardrail_information
off the caller's request dict. A blocking step records it on the executor's
working copy (or the raw-request snapshot), so the terminal result must carry it
back onto the request or the block is never counted.
"""
guard = RecordingGuardrail(guardrail_name="credentials-api-keys", scan_raw_request=scan_raw_request)
monkeypatch.setattr(litellm, "callbacks", [guard])
data = {
"messages": [{"role": "user", "content": "key AKIAIOSFODNN7EXAMPLE"}],
"metadata": {"user_api_key_hash": "abc"},
}
result = await PipelineExecutor.execute_steps(
steps=[PipelineStep(guardrail="credentials-api-keys", on_fail=on_fail, on_pass="next")],
mode="pre_call",
data=data,
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="baseline-pii-protection",
raw_request_snapshot={"messages": data["messages"], "metadata": {"user_api_key_hash": "abc"}},
)
assert result.terminal_action == on_fail
recorded = data["metadata"]["standard_logging_guardrail_information"]
assert [entry["guardrail_name"] for entry in recorded] == ["credentials-api-keys"]
assert recorded[0]["guardrail_status"] == "guardrail_intervened"
assert data["metadata"]["user_api_key_hash"] == "abc"
assert "guardrails" not in data["metadata"]
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_terminal_block_merges_guardrail_information_without_duplicates(monkeypatch: pytest.MonkeyPatch):
"""A pass_data step that returns a rewritten copy of the request, and a scan_raw_request step
that evaluates a deep copy taken before the pipeline ran, both leave earlier entries in two
dicts at once. Those must be carried back once while every step's own entry is kept."""
first = RecordingGuardrail(guardrail_name="pii-scan", block=False)
second = RecordingGuardrail(guardrail_name="credentials-api-keys", scan_raw_request=True)
monkeypatch.setattr(litellm, "callbacks", [first, second])
earlier = {"guardrail_name": "earlier-guard", "guardrail_status": "success"}
data = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
data["metadata"]["standard_logging_guardrail_information"] = [earlier]
result = await PipelineExecutor.execute_steps(
steps=[
PipelineStep(guardrail="pii-scan", on_fail="block", on_pass="next", pass_data=True),
PipelineStep(guardrail="credentials-api-keys", on_fail="block", on_pass="next"),
],
mode="pre_call",
data=data,
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="baseline-pii-protection",
raw_request_snapshot={
"messages": data["messages"],
"metadata": {"standard_logging_guardrail_information": [dict(earlier)]},
},
)
assert result.terminal_action == "block"
recorded = data["metadata"]["standard_logging_guardrail_information"]
assert [entry["guardrail_name"] for entry in recorded] == ["earlier-guard", "pii-scan", "credentials-api-keys"]
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_repeated_scan_raw_request_step_is_counted_once_per_evaluation(monkeypatch: pytest.MonkeyPatch):
"""Running the same raw-scan guardrail twice yields two identical entries; both must reach the caller,
while the entries the raw snapshot already held before the pipeline ran are not copied again."""
guard = RecordingGuardrail(guardrail_name="credentials-raw", scan_raw_request=True, block=False)
monkeypatch.setattr(litellm, "callbacks", [guard])
earlier = {"guardrail_name": "earlier-guard", "guardrail_status": "success"}
data = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
data["metadata"]["standard_logging_guardrail_information"] = [earlier]
result = await PipelineExecutor.execute_steps(
steps=[
PipelineStep(guardrail="credentials-raw", on_fail="block", on_pass="next"),
PipelineStep(guardrail="credentials-raw", on_fail="block", on_pass="next"),
],
mode="pre_call",
data=data,
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="raw-scan-policy",
raw_request_snapshot={
"messages": data["messages"],
"metadata": {"standard_logging_guardrail_information": [dict(earlier)]},
},
)
assert result.terminal_action == "allow"
assert result.modified_data is not None
recorded = result.modified_data["metadata"]["standard_logging_guardrail_information"]
assert [entry["guardrail_name"] for entry in recorded] == ["earlier-guard", "credentials-raw", "credentials-raw"]
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_escalation_step1_fails_step2_blocks(monkeypatch):