mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(guardrails): match policy-pipeline block response to direct guardrail attachment (#31421)
When a guardrail blocked a request through a flow-builder policy pipeline, the proxy discarded the guardrail's own exception and synthesized a generic guardrail_pipeline_error response, so the same guardrail produced a different HTTP response and trace span depending on whether it was attached directly or via a policy. The pipeline now carries the guardrail's original exception and re-raises it verbatim on block, enriching it with the blocking guardrail's name and mode exactly as the direct path does, so the two attachment methods are indistinguishable to clients and tracing. The generic pipeline error remains only as a fallback for blocks with no underlying exception (e.g. a guardrail that could not be found). Resolves LIT-4041
This commit is contained in:
parent
5a395ef909
commit
c14329128b
5 changed files with 251 additions and 13 deletions
|
|
@ -6,7 +6,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding.
|
|||
"""
|
||||
|
||||
import time
|
||||
from typing import Any, List, Optional
|
||||
from typing import Any, List, Literal, Optional
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -64,7 +64,12 @@ class PipelineExecutor:
|
|||
for i, step in enumerate(steps):
|
||||
start_time = time.perf_counter()
|
||||
|
||||
outcome, modified_data, error_detail = await PipelineExecutor._run_step(
|
||||
(
|
||||
outcome,
|
||||
modified_data,
|
||||
error_detail,
|
||||
original_exception,
|
||||
) = await PipelineExecutor._run_step(
|
||||
step=step,
|
||||
mode=mode,
|
||||
data=working_data,
|
||||
|
|
@ -108,6 +113,7 @@ class PipelineExecutor:
|
|||
terminal_action="block",
|
||||
step_results=step_results,
|
||||
error_message=error_detail,
|
||||
original_exception=original_exception,
|
||||
)
|
||||
|
||||
if action == "modify_response":
|
||||
|
|
@ -134,22 +140,30 @@ class PipelineExecutor:
|
|||
data: dict,
|
||||
user_api_key_dict: Any,
|
||||
call_type: str,
|
||||
) -> tuple:
|
||||
) -> tuple[
|
||||
Literal["pass", "fail", "error"],
|
||||
Optional[dict],
|
||||
Optional[str],
|
||||
Optional[Exception],
|
||||
]:
|
||||
"""
|
||||
Run a single pipeline step's guardrail.
|
||||
|
||||
Returns:
|
||||
Tuple of (outcome, modified_data, error_detail) where:
|
||||
Tuple of (outcome, modified_data, error_detail, original_exception):
|
||||
- outcome: "pass", "fail", or "error"
|
||||
- modified_data: dict if guardrail returned modified data, else None
|
||||
- error_detail: error message string if fail/error, else None
|
||||
- original_exception: the exception the guardrail raised, so the
|
||||
pipeline can re-raise it verbatim and match the direct-attachment
|
||||
response/trace, else None
|
||||
"""
|
||||
callback = PipelineExecutor._find_guardrail_callback(step.guardrail)
|
||||
callback = PipelineExecutor.find_guardrail_callback(step.guardrail)
|
||||
if callback is None:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Pipeline: guardrail '{step.guardrail}' not found in callbacks"
|
||||
)
|
||||
return ("error", None, f"Guardrail '{step.guardrail}' not found")
|
||||
return ("error", None, f"Guardrail '{step.guardrail}' not found", None)
|
||||
|
||||
try:
|
||||
# Inject guardrail name into metadata so should_run_guardrail() allows it
|
||||
|
|
@ -182,26 +196,26 @@ class PipelineExecutor:
|
|||
response=data.get("response"), # type: ignore
|
||||
)
|
||||
else:
|
||||
return ("error", None, f"Unsupported pipeline mode: {mode}")
|
||||
return ("error", None, f"Unsupported pipeline mode: {mode}", None)
|
||||
|
||||
# Normal return means pass
|
||||
modified_data = None
|
||||
if response is not None and isinstance(response, dict):
|
||||
modified_data = response
|
||||
return ("pass", modified_data, None)
|
||||
return ("pass", modified_data, None, None)
|
||||
|
||||
except Exception as e:
|
||||
if CustomGuardrail._is_guardrail_intervention(e):
|
||||
error_msg = _extract_error_message(e)
|
||||
return ("fail", None, error_msg)
|
||||
return ("fail", None, error_msg, e)
|
||||
else:
|
||||
verbose_proxy_logger.error(
|
||||
f"Pipeline: unexpected error from guardrail '{step.guardrail}': {e}"
|
||||
)
|
||||
return ("error", None, str(e))
|
||||
return ("error", None, str(e), e)
|
||||
|
||||
@staticmethod
|
||||
def _find_guardrail_callback(guardrail_name: str) -> Optional[CustomGuardrail]:
|
||||
def find_guardrail_callback(guardrail_name: str) -> Optional[CustomGuardrail]:
|
||||
"""Look up an initialized guardrail callback by name from litellm.callbacks."""
|
||||
for callback in litellm.callbacks:
|
||||
if isinstance(callback, CustomGuardrail):
|
||||
|
|
|
|||
|
|
@ -364,6 +364,16 @@ def _enrich_http_exception_with_guardrail_context(
|
|||
detail.setdefault("guardrail_mode", event_hook)
|
||||
|
||||
|
||||
def _exception_changes_request_flow(exc: BaseException) -> bool:
|
||||
"""
|
||||
True for guardrail exceptions the proxy turns into an alternate request flow
|
||||
(a reroute or a passthrough response) rather than a block. A pipeline step
|
||||
configured to block must honor that block, so these are surfaced as the
|
||||
generic pipeline block instead of being re-raised verbatim.
|
||||
"""
|
||||
return isinstance(exc, (SensitiveDataRouteException, ModifyResponseException))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _CallbackCapabilities:
|
||||
"""Cached per-hook capability flags derived from ``litellm.callbacks``.
|
||||
|
|
@ -1352,7 +1362,7 @@ class ProxyLogging:
|
|||
|
||||
@staticmethod
|
||||
def _handle_pipeline_result(
|
||||
result: Any,
|
||||
result: PipelineExecutionResult,
|
||||
data: dict,
|
||||
policy_name: str,
|
||||
) -> dict:
|
||||
|
|
@ -1367,6 +1377,21 @@ class ProxyLogging:
|
|||
return data
|
||||
|
||||
if result.terminal_action == "block":
|
||||
original_exception = result.original_exception
|
||||
if original_exception is not None and not _exception_changes_request_flow(
|
||||
original_exception
|
||||
):
|
||||
blocking_step = result.step_results[-1] if result.step_results else None
|
||||
if blocking_step is not None:
|
||||
callback = PipelineExecutor.find_guardrail_callback(
|
||||
blocking_step.guardrail_name
|
||||
)
|
||||
if callback is not None:
|
||||
_enrich_http_exception_with_guardrail_context(
|
||||
original_exception, callback
|
||||
)
|
||||
raise original_exception
|
||||
|
||||
step_results_serializable = [
|
||||
{
|
||||
"guardrail": sr.guardrail_name,
|
||||
|
|
|
|||
|
|
@ -99,8 +99,11 @@ class PipelineStepResult(BaseModel):
|
|||
class PipelineExecutionResult(BaseModel):
|
||||
"""Result of executing an entire pipeline."""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
terminal_action: str # block | allow | modify_response
|
||||
step_results: List[PipelineStepResult]
|
||||
modified_data: Optional[Dict[str, Any]] = None
|
||||
error_message: Optional[str] = None
|
||||
modify_response_message: Optional[str] = None
|
||||
original_exception: Optional[Exception] = Field(default=None, exclude=True)
|
||||
|
|
|
|||
|
|
@ -209,6 +209,75 @@ async def test_escalation_step1_fails_step2_blocks():
|
|||
litellm.callbacks = original_callbacks
|
||||
|
||||
|
||||
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
|
||||
@pytest.mark.asyncio
|
||||
async def test_block_carries_original_guardrail_exception():
|
||||
"""A blocking step must expose the guardrail's own raised exception on the
|
||||
result so the caller can re-raise it verbatim, giving the policy path the
|
||||
same response/trace as a direct guardrail attachment."""
|
||||
guard = AlwaysFailGuardrail(guardrail_name="moderation-filter")
|
||||
|
||||
pipeline = GuardrailPipeline(
|
||||
mode="pre_call",
|
||||
steps=[
|
||||
PipelineStep(
|
||||
guardrail="moderation-filter", on_fail="block", on_pass="allow"
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
original_callbacks = litellm.callbacks.copy()
|
||||
litellm.callbacks = [guard]
|
||||
|
||||
try:
|
||||
result = await PipelineExecutor.execute_steps(
|
||||
steps=pipeline.steps,
|
||||
mode=pipeline.mode,
|
||||
data={"messages": [{"role": "user", "content": "bad content"}]},
|
||||
user_api_key_dict=MagicMock(),
|
||||
call_type="completion",
|
||||
policy_name="content-safety",
|
||||
)
|
||||
|
||||
assert result.terminal_action == "block"
|
||||
assert isinstance(result.original_exception, HTTPException)
|
||||
assert result.original_exception.status_code == 400
|
||||
assert result.original_exception.detail == "Content policy violation"
|
||||
finally:
|
||||
litellm.callbacks = original_callbacks
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unsupported_mode_yields_error_outcome_without_exception():
|
||||
"""An unexpected hook mode must surface as an error outcome (carrying no
|
||||
original exception), not crash or run the guardrail."""
|
||||
guard = AlwaysPassGuardrail(guardrail_name="filter")
|
||||
|
||||
original_callbacks = litellm.callbacks.copy()
|
||||
litellm.callbacks = [guard]
|
||||
|
||||
try:
|
||||
result = await PipelineExecutor.execute_steps(
|
||||
steps=[PipelineStep(guardrail="filter", on_error="block", on_fail="block")],
|
||||
mode="during_call",
|
||||
data={"messages": [{"role": "user", "content": "hi"}]},
|
||||
user_api_key_dict=MagicMock(),
|
||||
call_type="completion",
|
||||
policy_name="content-safety",
|
||||
)
|
||||
|
||||
assert guard.calls == 0
|
||||
assert result.terminal_action == "block"
|
||||
assert result.step_results[0].outcome == "error"
|
||||
assert (
|
||||
"Unsupported pipeline mode: during_call"
|
||||
in result.step_results[0].error_detail
|
||||
)
|
||||
assert result.original_exception is None
|
||||
finally:
|
||||
litellm.callbacks = original_callbacks
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_passthrough_guardrail_failure_can_pipeline_block():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ import pytest
|
|||
from fastapi import HTTPException
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import SensitiveDataRouteException
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
ModifyResponseException,
|
||||
|
|
@ -359,6 +360,7 @@ async def test_maybe_execute_pipelines_blocks_on_block_terminal_action_raises(
|
|||
fake_result = MagicMock()
|
||||
fake_result.terminal_action = "block"
|
||||
fake_result.step_results = []
|
||||
fake_result.original_exception = None
|
||||
data = {"metadata": {"_guardrail_pipelines": [("policy-1", pipeline)]}, "messages": [], "model": "m"}
|
||||
|
||||
async def fake_execute_steps(**kwargs):
|
||||
|
|
@ -377,6 +379,42 @@ async def test_maybe_execute_pipelines_blocks_on_block_terminal_action_raises(
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_maybe_execute_pipelines_reraises_original_guardrail_exception(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
"""A policy-wrapped guardrail block must surface the guardrail's own
|
||||
exception verbatim, identical to the direct-attachment path."""
|
||||
pipeline = MagicMock()
|
||||
pipeline.mode = "pre_call"
|
||||
pipeline.steps = []
|
||||
original = HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": "Violated OpenAI moderation policy", "moderation_result": {"x": 1}},
|
||||
)
|
||||
fake_result = MagicMock()
|
||||
fake_result.terminal_action = "block"
|
||||
fake_result.step_results = []
|
||||
fake_result.original_exception = original
|
||||
data = {"metadata": {"_guardrail_pipelines": [("policy-1", pipeline)]}, "messages": [], "model": "m"}
|
||||
|
||||
async def fake_execute_steps(**kwargs):
|
||||
return fake_result
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.policy_engine.pipeline_executor.PipelineExecutor.execute_steps",
|
||||
fake_execute_steps,
|
||||
)
|
||||
with pytest.raises(HTTPException) as info:
|
||||
await proxy_logging._maybe_execute_pipelines(
|
||||
data=data,
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
call_type="completion",
|
||||
event_hook="pre_call",
|
||||
)
|
||||
assert info.value is original
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _handle_pipeline_result
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -391,10 +429,11 @@ def test_handle_pipeline_result_allow_with_modifications():
|
|||
assert out == {"a": 1, "b": 2, "c": 3}
|
||||
|
||||
|
||||
def test_handle_pipeline_result_block_raises_http_exception():
|
||||
def test_handle_pipeline_result_block_falls_back_to_generic_when_no_exception():
|
||||
result = MagicMock()
|
||||
result.terminal_action = "block"
|
||||
result.step_results = []
|
||||
result.original_exception = None
|
||||
with pytest.raises(HTTPException) as info:
|
||||
ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p")
|
||||
detail = info.value.detail
|
||||
|
|
@ -410,6 +449,94 @@ def test_handle_pipeline_result_block_raises_http_exception():
|
|||
}
|
||||
|
||||
|
||||
def test_handle_pipeline_result_block_reraises_original_guardrail_exception():
|
||||
"""The policy path must re-raise the guardrail's own exception untouched,
|
||||
not wrap it in a generic ``guardrail_pipeline_error``; this is what makes
|
||||
the response and trace span identical to the direct-attachment path."""
|
||||
original = HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": "Violated OpenAI moderation policy",
|
||||
"moderation_result": {"violated_categories": ["harassment"]},
|
||||
},
|
||||
)
|
||||
result = MagicMock()
|
||||
result.terminal_action = "block"
|
||||
result.step_results = []
|
||||
result.original_exception = original
|
||||
with pytest.raises(HTTPException) as info:
|
||||
ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p")
|
||||
assert info.value is original
|
||||
assert info.value.detail == {
|
||||
"error": "Violated OpenAI moderation policy",
|
||||
"moderation_result": {"violated_categories": ["harassment"]},
|
||||
}
|
||||
|
||||
|
||||
def test_handle_pipeline_result_block_enriches_with_guardrail_name_and_mode():
|
||||
"""The re-raised exception must gain the blocking guardrail's name and mode,
|
||||
matching the enrichment the direct-attachment path applies."""
|
||||
cb = _make_guardrail() # guardrail_name="g", event_hook=pre_call
|
||||
original = HTTPException(status_code=400, detail={"error": "blocked"})
|
||||
result = MagicMock()
|
||||
result.terminal_action = "block"
|
||||
result.step_results = [MagicMock(guardrail_name="g")]
|
||||
result.original_exception = original
|
||||
|
||||
saved = litellm.callbacks
|
||||
litellm.callbacks = [cb]
|
||||
try:
|
||||
with pytest.raises(HTTPException) as info:
|
||||
ProxyLogging._handle_pipeline_result(
|
||||
result=result, data={"model": "m"}, policy_name="p"
|
||||
)
|
||||
finally:
|
||||
litellm.callbacks = saved
|
||||
|
||||
assert info.value is original
|
||||
assert info.value.detail["guardrail_name"] == "g"
|
||||
assert info.value.detail["guardrail_mode"] == GuardrailEventHooks.pre_call
|
||||
|
||||
|
||||
def test_handle_pipeline_result_block_does_not_reraise_sensitive_data_route():
|
||||
"""A step configured to block must enforce the block even when the guardrail
|
||||
raised a reroute exception; re-raising it verbatim would route the request to
|
||||
an alternate model instead of blocking, bypassing the configured policy."""
|
||||
original = SensitiveDataRouteException(
|
||||
route_to_model="on-prem-model",
|
||||
session_id="sess-1",
|
||||
guardrail_name="pii-router",
|
||||
)
|
||||
result = MagicMock()
|
||||
result.terminal_action = "block"
|
||||
result.step_results = [MagicMock(guardrail_name="pii-router")]
|
||||
result.original_exception = original
|
||||
with pytest.raises(HTTPException) as info:
|
||||
ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p")
|
||||
assert info.value.status_code == 400
|
||||
assert info.value.detail["error"]["type"] == "guardrail_pipeline_error"
|
||||
|
||||
|
||||
def test_handle_pipeline_result_block_does_not_reraise_modify_response():
|
||||
"""A step configured to block must enforce the block even when the guardrail
|
||||
raised a passthrough/modify-response exception; re-raising it verbatim would
|
||||
return the guardrail's synthetic response instead of blocking."""
|
||||
original = ModifyResponseException(
|
||||
message="redacted",
|
||||
model="m",
|
||||
request_data={"model": "m"},
|
||||
guardrail_name="masker",
|
||||
)
|
||||
result = MagicMock()
|
||||
result.terminal_action = "block"
|
||||
result.step_results = [MagicMock(guardrail_name="masker")]
|
||||
result.original_exception = original
|
||||
with pytest.raises(HTTPException) as info:
|
||||
ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p")
|
||||
assert info.value.status_code == 400
|
||||
assert info.value.detail["error"]["type"] == "guardrail_pipeline_error"
|
||||
|
||||
|
||||
def test_handle_pipeline_result_modify_response_raises_modify_exception():
|
||||
result = MagicMock()
|
||||
result.terminal_action = "modify_response"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue