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

Merge pull request #35260 from BerriAI/litellm_messages_streaming_post_call_guardrails

Browse source 
fix(proxy): run post_call guardrails on /v1/messages streaming via unified guardrail translation
This commit is contained in:
Mateo Wang 2026-07-30 16:48:15 -07:00 • committed by GitHub
commit 8e287652c6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 367 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -804,6 +804,8 @@ class UnifiedLLMGuardrails(CustomLogger):
user_api_key_dict: UserAPIKeyAuth,
response: Any,
request_data: dict,
guardrail_to_apply: Union[CustomGuardrail, None] = None,
buffer_until_moderated_default: bool = False,
) -> AsyncGenerator[Any, None]:
"""
Passes the entire stream to the guardrail
@ -824,7 +826,8 @@ class UnifiedLLMGuardrails(CustomLogger):
# litellm.integrations.custom_guardrail.
from litellm.integrations.custom_guardrail import ModifyResponseException
guardrail_to_apply: CustomGuardrail = request_data.pop("guardrail_to_apply", None)
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
@ -852,7 +855,7 @@ class UnifiedLLMGuardrails(CustomLogger):
# release the original chunks are replayed as-is, so a
# content-rewriting guardrail (e.g. PII masking) would leak
# unredacted content. Guarded below via mask_response_content.
buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", False)
buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", buffer_until_moderated_default)
if (
buffer_until_moderated

33
litellm/proxy/utils.py
View file

@ -187,6 +187,8 @@ else:
unified_guardrail = UnifiedLLMGuardrails()
NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES: "frozenset[CallTypes]" = frozenset({CallTypes.anthropic_messages})
def print_verbose(print_statement):
"""
@ -1762,6 +1764,20 @@ class ProxyLogging:
cache[sig] = caps
return caps
@staticmethod
def _stream_requires_guardrail_translation(user_api_key_dict: UserAPIKeyAuth) -> bool:
from litellm.litellm_core_utils.api_route_to_call_types import (
get_call_types_for_route,
)
route = user_api_key_dict.request_route
if not route:
return False
call_types = get_call_types_for_route(route)
if not call_types:
return False
return call_types[0] in NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES
@staticmethod
def has_post_call_response_headers_callbacks() -> bool:
return ProxyLogging._callback_capabilities().has_post_call_response_headers
@ -2723,6 +2739,7 @@ class ProxyLogging:
request_data = _check_and_merge_model_level_guardrails(data=request_data, llm_router=llm_router)
current_response = response
stream_needs_translation = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict)
for resolved_callback, kind in caps.iterator_overrides:
if isinstance(resolved_callback, CustomGuardrail):
@ -2731,7 +2748,18 @@ class ProxyLogging:
is not True
):
continue
if kind == "override":
effective_kind = (
"apply_guardrail"
if (
kind == "override"
and stream_needs_translation
and isinstance(resolved_callback, CustomGuardrail)
and resolved_callback.uses_apply_guardrail_interface()
and not resolved_callback.mask_response_content
)
else kind
)
if effective_kind == "override":
current_response = self._wrap_streaming_iterator_with_enrichment(
resolved_callback,
resolved_callback.async_post_call_streaming_iterator_hook(
@ -2742,13 +2770,14 @@ class ProxyLogging:
)
else:
# kind == "apply_guardrail": route through unified_guardrail
request_data["guardrail_to_apply"] = resolved_callback
current_response = self._wrap_streaming_iterator_with_enrichment(
resolved_callback,
unified_guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
request_data=request_data,
response=current_response,
guardrail_to_apply=resolved_callback,
buffer_until_moderated_default=(kind == "override"),
),
)

331
tests/test_litellm/proxy/test_proxy_logging_hook_detection.py
View file

@ -148,3 +148,334 @@ def test_callback_capabilities_cache_invalidates_on_list_change(monkeypatch):
caps = ProxyLogging._callback_capabilities()
assert caps.has_pre_call_override is True
assert pre in caps.resolved_callbacks
def _sse_bytes(event: str, payload: dict) -> bytes:
import json
return f"event: {event}\ndata: {json.dumps(payload)}\n\n".encode()
def _anthropic_stream_chunks(text_parts):
chunks = [
_sse_bytes(
"message_start",
{
"type": "message_start",
"message": {
"model": "claude-sonnet-5",
"id": "msg_1",
"type": "message",
"role": "assistant",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 20, "output_tokens": 1},
},
},
),
_sse_bytes(
"content_block_start",
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
),
]
for part in text_parts:
chunks.append(
_sse_bytes(
"content_block_delta",
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": part}},
)
)
chunks.append(_sse_bytes("content_block_stop", {"type": "content_block_stop", "index": 0}))
chunks.append(
_sse_bytes(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {"input_tokens": 20, "output_tokens": 8},
},
)
)
chunks.append(_sse_bytes("message_stop", {"type": "message_stop"}))
return chunks
def _content_filter_guardrail(action: str, guardrail_cls=None, **guardrail_kwargs):
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
ContentFilterGuardrail,
)
from litellm.types.guardrails import BlockedWord, ContentFilterAction
cls = guardrail_cls or ContentFilterGuardrail
return cls(
guardrail_name="output-filter",
blocked_words=[BlockedWord(keyword="zebra", action=ContentFilterAction(action))],
event_hook="post_call",
default_on=True,
**guardrail_kwargs,
)
def _streaming_logging_obj():
import datetime
import uuid
from litellm.litellm_core_utils.litellm_logging import Logging
return Logging(
model="claude-sonnet-5",
messages=[{"role": "user", "content": "Reply with exactly: the zebra runs"}],
stream=True,
call_type="anthropic_messages",
start_time=datetime.datetime.now(),
litellm_call_id=str(uuid.uuid4()),
function_id="test",
)
def test_stream_requires_guardrail_translation_route_detection():
from litellm.proxy._types import UserAPIKeyAuth
assert (
ProxyLogging._stream_requires_guardrail_translation(
UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages")
)
is True
)
assert (
ProxyLogging._stream_requires_guardrail_translation(
UserAPIKeyAuth(api_key="sk-1234", request_route="/chat/completions")
)
is False
)
assert ProxyLogging._stream_requires_guardrail_translation(UserAPIKeyAuth(api_key="sk-1234")) is False
assert (
ProxyLogging._stream_requires_guardrail_translation(
UserAPIKeyAuth(api_key="sk-1234", request_route="/route/without/call/types")
)
is False
)
@pytest.mark.asyncio
async def test_post_call_stream_guardrail_blocks_anthropic_messages_stream(monkeypatch):
"""
Regression test for https://github.com/BerriAI/litellm/issues/35257.
/v1/messages streams raw Anthropic SSE bytes. A guardrail whose custom
iterator hook only understands OpenAI ModelResponseStream chunks used to
receive those bytes directly and silently pass every chunk through
unscanned. The dispatch must route apply_guardrail-capable guardrails
through unified_guardrail's anthropic translation so blocked output
raises instead of streaming to the client. Because the guardrail's own
iterator hook withheld content until scanned, the rerouted invocation
defaults to buffer_until_moderated, so nothing may reach the client
before the block fires.
"""
from fastapi import HTTPException
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
guardrail = _content_filter_guardrail("BLOCK")
monkeypatch.setattr(litellm, "callbacks", [guardrail])
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
request_data = {
"model": "claude-sonnet-5",
"litellm_logging_obj": _streaming_logging_obj(),
"metadata": {},
}
async def fake_stream():
for chunk in _anthropic_stream_chunks(["the", " zebra runs"]):
yield chunk
delivered = []
with pytest.raises(HTTPException) as exc_info:
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
response=fake_stream(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"),
request_data=request_data,
):
delivered.append(chunk)
detail = exc_info.value.detail
assert detail["guardrail_name"] == "output-filter"
assert detail["keyword"] == "zebra"
assert delivered == []
@pytest.mark.asyncio
async def test_post_call_stream_guardrail_keeps_own_iterator_on_chat_completions(monkeypatch):
"""
On /chat/completions the guardrail's own iterator hook must keep running:
it masks incrementally inside ModelResponseStream chunks, which the
unified block_only path never does. Masked output proves the own-hook
path was used.
"""
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
guardrail = _content_filter_guardrail("MASK")
monkeypatch.setattr(litellm, "callbacks", [guardrail])
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
async def fake_stream():
yield ModelResponseStream(
choices=[StreamingChoices(index=0, delta=Delta(content="the zebra runs"))]
)
yield ModelResponseStream(
choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")]
)
delivered_text = ""
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
response=fake_stream(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/chat/completions"),
request_data={"model": "gpt-4o-mini", "metadata": {}},
):
for choice in chunk.choices:
delivered_text += choice.delta.content or ""
assert "zebra" not in delivered_text
assert delivered_text != ""
@pytest.mark.asyncio
async def test_unified_guardrail_iterator_accepts_explicit_guardrail(monkeypatch):
"""
The dispatch passes each guardrail explicitly instead of through a shared
request_data key, so chaining two unified-routed guardrails cannot drop
all but the last one.
"""
from fastapi import HTTPException
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import unified_guardrail
guardrail = _content_filter_guardrail("BLOCK")
request_data = {
"model": "claude-sonnet-5",
"litellm_logging_obj": _streaming_logging_obj(),
"metadata": {},
}
async def fake_stream():
for chunk in _anthropic_stream_chunks(["the", " zebra runs"]):
yield chunk
with pytest.raises(HTTPException):
async for _ in unified_guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"),
response=fake_stream(),
request_data=request_data,
guardrail_to_apply=guardrail,
):
pass
@pytest.mark.asyncio
async def test_post_call_stream_guardrail_reroutes_inherited_apply_guardrail(monkeypatch):
"""
The reroute predicate must recognize apply_guardrail implementations
inherited from a parent class, not only ones defined on the registered
leaf class. A vendor base class can carry apply_guardrail while the leaf
only overrides the streaming iterator; a leaf-class ``__dict__`` check
would leave that guardrail on the raw Anthropic SSE path unscanned.
"""
from fastapi import HTTPException
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
ContentFilterGuardrail,
)
class _InheritsApplyGuardrail(ContentFilterGuardrail):
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
async for item in response:
yield item
guardrail = _content_filter_guardrail("BLOCK", guardrail_cls=_InheritsApplyGuardrail)
assert "apply_guardrail" not in type(guardrail).__dict__
monkeypatch.setattr(litellm, "callbacks", [guardrail])
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
request_data = {
"model": "claude-sonnet-5",
"litellm_logging_obj": _streaming_logging_obj(),
"metadata": {},
}
async def fake_stream():
for chunk in _anthropic_stream_chunks(["the", " zebra runs"]):
yield chunk
delivered = []
with pytest.raises(HTTPException) as exc_info:
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
response=fake_stream(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"),
request_data=request_data,
):
delivered.append(chunk)
assert exc_info.value.detail["keyword"] == "zebra"
assert delivered == []
@pytest.mark.asyncio
async def test_post_call_stream_masking_guardrail_keeps_own_iterator_on_anthropic(monkeypatch):
"""
A guardrail with mask_response_content=True must stay on its own iterator
hook on /v1/messages. The unified streaming path cannot re-emit rewritten
text on raw Anthropic SSE (block_only drops rewrites and buffered replay
releases the unredacted originals), so rerouting such a guardrail would
deliver content it decided to mask. PANW Prisma AIRS is the concrete
case: its own hook parses the raw bytes and blocks instead of masking.
"""
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
ContentFilterGuardrail,
)
own_hook_streams = []
class _MasksViaOwnRawStreamHook(ContentFilterGuardrail):
apply_guardrail = ContentFilterGuardrail.apply_guardrail
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
own_hook_streams.append(request_data.get("model"))
async for item in response:
yield item
guardrail = _content_filter_guardrail(
"BLOCK", guardrail_cls=_MasksViaOwnRawStreamHook, mask_response_content=True
)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
chunks = _anthropic_stream_chunks(["the", " zebra runs"])
async def fake_stream():
for chunk in chunks:
yield chunk
delivered = []
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
response=fake_stream(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"),
request_data={
"model": "claude-sonnet-5",
"litellm_logging_obj": _streaming_logging_obj(),
"metadata": {},
},
):
delivered.append(chunk)
assert own_hook_streams == ["claude-sonnet-5"]
assert delivered == chunks