scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123

fix(policy_engine): fail closed on streaming for content-rewriting pipeline steps and untranslatable routes

Browse source
This commit is contained in:
mateo-berri 2026-08-29 13:28:11 -07:00
parent c5bcf3a735
commit fa5a10941e
3 changed files with 155 additions and 59 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -876,6 +876,14 @@ class UnifiedLLMGuardrails(CustomLogger):
choices: Final = _chunk_choices(item)
return any(getattr(choice, "finish_reason", None) is not None for choice in choices)
def resolve_streaming_flag(self, guardrail_to_apply: CustomGuardrail | None, name: str, default: object) -> object:
"""Streaming flag resolution order (later wins): default < guardrail
attribute < guardrail_config dict < this callback's optional_params."""
attribute_value: Final = default if guardrail_to_apply is None else getattr(guardrail_to_apply, name, default)
config: Final = None if guardrail_to_apply is None else getattr(guardrail_to_apply, "guardrail_config", None)
config_value: Final = config.get(name, attribute_value) if isinstance(config, dict) else attribute_value
return self.optional_params.get(name, config_value)
async def async_post_call_streaming_iterator_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -906,17 +914,8 @@ class UnifiedLLMGuardrails(CustomLogger):
if guardrail_to_apply is None:
guardrail_to_apply = request_data.pop("guardrail_to_apply", None)
# Get streaming configuration. Resolution order (later wins): default
# < guardrail attribute < guardrail_config dict < this callback's
# optional_params.
def _streaming_flag(name: str, default: object) -> Any:
value = default
if guardrail_to_apply is not None:
value = getattr(guardrail_to_apply, name, value)
config: Final[Mapping[str, object]] = getattr(guardrail_to_apply, "guardrail_config", {})
if isinstance(config, dict):
value = config.get(name, value)
return self.optional_params.get(name, value)
return self.resolve_streaming_flag(guardrail_to_apply, name, default)
sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5)
# Only apply the guardrail at end of stream (not per chunk).

111
litellm/proxy/utils.py
View file

@ -139,6 +139,7 @@ from litellm.proxy.db.token_auth import (
)
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
resolve_endpoint_translation,
)
from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook
from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
@ -486,11 +487,19 @@ def _merge_pipeline_metadata_writes(data: dict, modified_data: Mapping[str, obje
_merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key))
def _pipeline_step_supports_streaming(guardrail_name: str) -> bool:
def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool:
callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name)
return callback is not None and PipelineExecutor.supports_unified_execution(callback)
def _pipeline_step_rewrites_streamed_content(guardrail_name: str) -> bool:
callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name)
if callback is None:
return False
transform_mode: Final = unified_guardrail.resolve_streaming_flag(callback, "streaming_transform_mode", "block_only")
return callback.mask_response_content or transform_mode == "incremental_diff"
class _PipelineErrorBody(TypedDict):
message: ReadOnly[str]
type: ReadOnly[str]
@ -502,17 +511,20 @@ class _PipelineErrorDetail(TypedDict):
error: ReadOnly[_PipelineErrorBody]
def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None:
def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth) -> None:
"""
Reject up front the requests whose post_call pipelines could never run.
Background responses skip the post_call hooks entirely, so a pipeline
governing one would silently never execute. Streaming responses execute
pipelines against the buffered stream through the endpoint guardrail
translations, which requires every step's guardrail to support the unified
apply_guardrail interface; steps that cannot (native-lifecycle guardrails,
or guardrails not registered at all) keep the 400 rather than letting
ungoverned output stream through.
translation of the request route, releasing the buffered chunks verbatim
on allow. That needs every step's guardrail to support the unified
apply_guardrail interface and to only allow or block (a step that rewrites
streamed content, via mask_response_content or
streaming_transform_mode=incremental_diff, would have its rewrite silently
dropped), and needs the route to have a translation at all; anything else
keeps the 400 rather than letting ungoverned output stream through.
"""
is_stream: Final = data.get("stream") is True
is_background: Final = data.get("background") is True
@ -536,30 +548,61 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None
}
}
raise HTTPException(status_code=400, detail=background_detail)
unsupported_guardrails: Final = tuple(
dict.fromkeys(
step.guardrail
for _policy_name, pipeline in post_call_pipelines
for step in pipeline.steps
if not _pipeline_step_supports_streaming(step.guardrail)
)
step_guardrails: Final = tuple(
dict.fromkeys(step.guardrail for _policy_name, pipeline in post_call_pipelines for step in pipeline.steps)
)
if not unsupported_guardrails:
unsupported_guardrails: Final = tuple(
guardrail for guardrail in step_guardrails if not _pipeline_step_supports_unified_streaming(guardrail)
)
if unsupported_guardrails:
unsupported_detail: Final[_PipelineErrorDetail] = {
"error": {
"message": (
"Policies with post_call guardrail pipelines cannot govern streaming responses "
"because these pipeline guardrails do not support the unified apply_guardrail "
f"interface: {', '.join(unsupported_guardrails)}. Retry with stream=false, or drop "
"them from the pipeline steps so guardrails.add scans them on streamed output."
),
"type": "guardrail_pipeline_error",
"policies": post_call_policies,
"guardrails": unsupported_guardrails,
}
}
raise HTTPException(status_code=400, detail=unsupported_detail)
rewriting_guardrails: Final = tuple(
guardrail for guardrail in step_guardrails if _pipeline_step_rewrites_streamed_content(guardrail)
)
if rewriting_guardrails:
rewriting_detail: Final[_PipelineErrorDetail] = {
"error": {
"message": (
"Policies with post_call guardrail pipelines cannot govern streaming responses "
"because these pipeline guardrails rewrite streamed content (mask_response_content "
"or streaming_transform_mode=incremental_diff), which pipeline steps would release "
f"unmodified: {', '.join(rewriting_guardrails)}. Retry with stream=false, or drop "
"them from the pipeline steps so guardrails.add applies them to streamed output."
),
"type": "guardrail_pipeline_error",
"policies": post_call_policies,
"guardrails": rewriting_guardrails,
}
}
raise HTTPException(status_code=400, detail=rewriting_detail)
route: Final = user_api_key_dict.request_route
if not route or resolve_endpoint_translation(user_api_key_dict, None) is not None:
return
unsupported_detail: Final[_PipelineErrorDetail] = {
route_detail: Final[_PipelineErrorDetail] = {
"error": {
"message": (
"Policies with post_call guardrail pipelines cannot govern streaming responses "
"because these pipeline guardrails do not support the unified apply_guardrail "
f"interface: {', '.join(unsupported_guardrails)}. Retry with stream=false, or move "
"them from pipeline steps to guardrails.add, which scans streamed output."
"Policies with post_call guardrail pipelines cannot govern streaming responses on "
f"route {route} because it has no endpoint guardrail translation to scan the stream "
f"through: {', '.join(post_call_policies)}. Retry with stream=false."
),
"type": "guardrail_pipeline_error",
"policies": post_call_policies,
"guardrails": unsupported_guardrails,
}
}
raise HTTPException(status_code=400, detail=unsupported_detail)
raise HTTPException(status_code=400, detail=route_detail)
def _prompt_block_text(block: object) -> str:
@ -1915,7 +1958,7 @@ class ProxyLogging:
)
try:
_raise_for_streaming_post_call_pipelines(data)
_raise_for_streaming_post_call_pipelines(data, user_api_key_dict)
# Execute guardrail pipelines before the normal callback loop
data, _ = await self._maybe_execute_pipelines(
@ -3431,10 +3474,6 @@ class ProxyLogging:
releases the buffered chunks verbatim; a block or modify_response
terminates with the translation's block chunks or the raised error.
"""
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
resolve_endpoint_translation,
)
buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict
async for item in response:
buffered.append(item)
@ -3444,17 +3483,15 @@ class ProxyLogging:
resolved: Final = resolve_endpoint_translation(user_api_key_dict, buffered[0])
if resolved is None:
policy_names: Final = tuple(policy_name for policy_name, _pipeline in pipelines)
unresolvable_detail: Final[_PipelineErrorDetail] = {
"error": {
"message": (
"Policy pipelines could not govern this streaming response shape; "
f"the response was withheld: {', '.join(policy_names)}."
),
"type": "guardrail_pipeline_error",
"policies": policy_names,
}
}
raise HTTPException(status_code=500, detail=unresolvable_detail)
raise ProxyException(
message=(
"Policy pipelines could not govern this streaming response shape; "
f"the response was withheld: {', '.join(policy_names)}."
),
type="guardrail_pipeline_error",
param=None,
code=500,
)
call_type, endpoint_translation = resolved
for policy_name, pipeline in pipelines:

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

@ -23,6 +23,7 @@ from litellm.integrations.custom_guardrail import (
ModifyResponseException,
)
from litellm.integrations.prometheus import PrometheusLogger
from litellm.proxy._types import ProxyException
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
from litellm.proxy.utils import ProxyLogging, _raise_for_streaming_post_call_pipelines
from litellm.types.guardrails import GuardrailEventHooks
@ -1309,31 +1310,35 @@ async def test_pre_call_hook_rejects_background_request_with_post_call_pipeline(
assert "background=false" in info.value.detail["error"]["message"]
def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_call():
def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_call(make_user_api_key_auth):
post_call = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g", on_fail="block")])
pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g", on_fail="block")])
auth = make_user_api_key_auth(request_route="/custom/stream")
assert (
_raise_for_streaming_post_call_pipelines(
{"stream": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}}
{"stream": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}}, auth
)
is None
)
assert (
_raise_for_streaming_post_call_pipelines(
{"background": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}}
{"background": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}}, auth
)
is None
)
assert _raise_for_streaming_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", post_call)]}}) is None
assert (
_raise_for_streaming_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", post_call)]}}, auth)
is None
)
assert (
_raise_for_streaming_post_call_pipelines(
{"stream": True, "metadata": {"_guardrail_pipelines": [("p", pre_call)]}}
{"stream": True, "metadata": {"_guardrail_pipelines": [("p", pre_call)]}}, auth
)
is None
)
assert _raise_for_streaming_post_call_pipelines({"stream": True}) is None
assert _raise_for_streaming_post_call_pipelines({"background": True}) is None
assert _raise_for_streaming_post_call_pipelines({"stream": True}, auth) is None
assert _raise_for_streaming_post_call_pipelines({"background": True}, auth) is None
# ---------------------------------------------------------------------------
@ -1366,15 +1371,16 @@ async def _async_chunk_iter(chunks: List[Any]):
@pytest.mark.asyncio
@pytest.mark.parametrize("request_route", [None, "/v1/chat/completions"])
async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_supports_unified(
proxy_logging, make_user_api_key_auth, monkeypatch
proxy_logging, make_user_api_key_auth, monkeypatch, request_route
):
seen: Dict[str, Any] = {}
monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)])
data = _post_call_pipeline_data(stream=True)
out = await proxy_logging.pre_call_hook(
user_api_key_dict=make_user_api_key_auth(),
user_api_key_dict=make_user_api_key_auth(request_route=request_route),
data=data,
call_type="completion",
guardrails_only=True,
@ -1422,6 +1428,60 @@ async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_lacks_uni
assert "apply_guardrail" in info.value.detail["error"]["message"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"rewrite_attribute, value",
[
("mask_response_content", True),
("streaming_transform_mode", "incremental_diff"),
("guardrail_config", {"streaming_transform_mode": "incremental_diff"}),
],
)
async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_rewrites_streamed_content(
proxy_logging, make_user_api_key_auth, monkeypatch, rewrite_attribute, value
):
seen: Dict[str, Any] = {}
guardrail = _unified_stream_guardrail(seen)
setattr(guardrail, rewrite_attribute, value)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
data = _post_call_pipeline_data(stream=True)
with pytest.raises(HTTPException) as info:
await proxy_logging.pre_call_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
data=data,
call_type="completion",
guardrails_only=True,
)
assert info.value.status_code == 400
assert info.value.detail["error"]["guardrails"] == ("gr-post",)
assert "rewrite streamed content" in info.value.detail["error"]["message"]
assert seen.get("count") is None
@pytest.mark.asyncio
async def test_pre_call_hook_rejects_streaming_when_route_has_no_guardrail_translation(
proxy_logging, make_user_api_key_auth, monkeypatch
):
seen: Dict[str, Any] = {}
monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)])
data = _post_call_pipeline_data(stream=True)
with pytest.raises(HTTPException) as info:
await proxy_logging.pre_call_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/custom/stream"),
data=data,
call_type="completion",
guardrails_only=True,
)
assert info.value.status_code == 400
assert info.value.detail["error"]["policies"] == ("response-governance",)
assert "/custom/stream" in info.value.detail["error"]["message"]
assert seen.get("count") is None
@pytest.mark.asyncio
async def test_streaming_iterator_hook_pipeline_allow_releases_buffered_chunks(
proxy_logging, make_user_api_key_auth, monkeypatch
@ -1490,10 +1550,10 @@ async def test_streaming_iterator_hook_pipeline_withholds_unresolvable_response_
):
delivered.append(item)
with pytest.raises(HTTPException) as info:
with pytest.raises(ProxyException) as info:
await _drain()
assert delivered == []
assert info.value.status_code == 500
assert "withheld" in info.value.detail["error"]["message"]
assert info.value.code == "500"
assert "withheld" in info.value.message
assert seen.get("count") is None