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

fix(guardrails): scan streaming /v1/messages responses instead of raw Anthropic SSE bytes

Browse source 
Post-call guardrails that implement their own async_post_call_streaming_iterator_hook were handed the provider-native stream on /v1/messages, so hooks written against ModelResponseStream skipped every chunk while still recording a successful scan. On routes that do not stream OpenAI-format chunks, dispatch through the translation-aware apply_guardrail path when the guardrail provides one, and warn otherwise.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-07-30 19:37:15 +00:00
parent 71b825a7f0
commit 8baaf6e71b
3 changed files with 215 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
litellm/litellm_core_utils/api_route_to_call_types.py
View file

@ -67,3 +67,31 @@ def get_routes_for_call_type(call_type: CallTypes) -> list:
if call_type in types:
routes.append(route)
return routes
OPENAI_CHUNK_STREAMING_CALL_TYPES = frozenset(
{
CallTypes.completion,
CallTypes.acompletion,
CallTypes.text_completion,
CallTypes.atext_completion,
}
)
def route_streams_openai_format_chunks(route: Optional[str]) -> bool:
"""
Whether streaming chunks emitted by `route` are OpenAI-format
``ModelResponseStream`` objects.
Routes such as ``/v1/messages`` or ``/v1/responses`` relay provider-native
payloads (Anthropic SSE bytes, Responses API events), which hooks written
against ``ModelResponseStream`` cannot read. Unknown routes are treated as
OpenAI-format to preserve existing behavior.
"""
if route is None:
return True
call_types = get_call_types_for_route(route)
if call_types is None:
return True
return any(call_type in OPENAI_CHUNK_STREAMING_CALL_TYPES for call_type in call_types)

44
litellm/proxy/utils.py
View file

@ -105,6 +105,9 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.prometheus import PrometheusLogger
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert
from litellm.litellm_core_utils.api_route_to_call_types import (
route_streams_openai_format_chunks,
)
from litellm.litellm_core_utils.core_helpers import coerce_token_limit
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
@ -2669,6 +2672,8 @@ class ProxyLogging:
current_response = response
streams_openai_chunks = route_streams_openai_format_chunks(user_api_key_dict.request_route)
for resolved_callback, kind in caps.iterator_overrides:
if isinstance(resolved_callback, CustomGuardrail):
if (
@ -2676,7 +2681,13 @@ class ProxyLogging:
is not True
):
continue
if kind == "override":
effective_kind = ProxyLogging._resolve_iterator_override_kind(
resolved_callback=resolved_callback,
kind=kind,
streams_openai_chunks=streams_openai_chunks,
request_route=user_api_key_dict.request_route,
)
if effective_kind == "override":
current_response = self._wrap_streaming_iterator_with_enrichment(
resolved_callback,
resolved_callback.async_post_call_streaming_iterator_hook(
@ -2686,7 +2697,7 @@ class ProxyLogging:
),
)
else:
# kind == "apply_guardrail": route through unified_guardrail
# effective_kind == "apply_guardrail": route through unified_guardrail
request_data["guardrail_to_apply"] = resolved_callback
current_response = self._wrap_streaming_iterator_with_enrichment(
resolved_callback,
@ -2707,6 +2718,35 @@ class ProxyLogging:
# we reach this point the metadata is fully populated.
ProxyLogging._fire_deferred_stream_logging(request_data)
@staticmethod
def _resolve_iterator_override_kind(
resolved_callback: Any,
kind: str,
streams_openai_chunks: bool,
request_route: Optional[str],
) -> str:
"""
Pick how a callback participates in the streaming post-call chain.
A callback's own ``async_post_call_streaming_iterator_hook`` is written
against OpenAI-format ``ModelResponseStream`` chunks, so it cannot scan
the provider-native payloads relayed by routes like ``/v1/messages``
(raw Anthropic SSE bytes). On those routes prefer the translation-aware
``apply_guardrail`` path when the callback provides one; otherwise warn,
since the guardrail will not see the response.
"""
if kind != "override" or streams_openai_chunks:
return kind
if "apply_guardrail" in type(resolved_callback).__dict__:
return "apply_guardrail"
verbose_proxy_logger.warning(
"Guardrail %s only implements async_post_call_streaming_iterator_hook, which expects "
"OpenAI-format chunks; it cannot scan the streaming response on %s",
getattr(resolved_callback, "guardrail_name", None) or type(resolved_callback).__name__,
request_route,
)
return kind
@staticmethod
def _fire_deferred_stream_logging(request_data: dict) -> None:
"""

145
tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py
View file

@ -430,3 +430,148 @@ async def test_post_call_response_headers_hook_swallows_callback_error(proxy_log
data={}, user_api_key_dict=make_user_api_key_auth(), response=response
)
assert out == {}
# ---------------------------------------------------------------------------
# non-OpenAI streaming surfaces (/v1/messages) -- issue #35257
# ---------------------------------------------------------------------------
def _anthropic_sse(event_type: str, data: Dict[str, Any]) -> bytes:
import json
return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
async def _anthropic_stream(text: str):
yield _anthropic_sse(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 1, "output_tokens": 0},
},
},
)
yield _anthropic_sse(
"content_block_start",
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
)
yield _anthropic_sse(
"content_block_delta",
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}},
)
yield _anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0})
yield _anthropic_sse(
"message_delta",
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 3}},
)
yield _anthropic_sse("message_stop", {"type": "message_stop"})
def _content_filter_guardrail():
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
ContentFilterGuardrail,
)
guardrail = ContentFilterGuardrail(
guardrail_name="output-filter",
event_hook="post_call",
blocked_words=[{"keyword": "zebra", "action": "BLOCK"}],
default_on=True,
)
# scan every chunk so the block decision lands before the offending delta
# is forwarded, matching the /chat/completions withholding behavior
guardrail.streaming_sampling_rate = 1
return guardrail
async def _run_v1_messages_stream(proxy_logging, make_user_api_key_auth, text: str):
request_data: Dict[str, Any] = {
"metadata": {"guardrails": ["output-filter"]},
"messages": [{"role": "user", "content": "hi"}],
}
collected: List[Any] = []
blocked = False
try:
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
response=_anthropic_stream(text),
user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"),
request_data=request_data,
):
collected.append(chunk)
except HTTPException:
blocked = True
raw = b"".join(c if isinstance(c, bytes) else str(c).encode() for c in collected).decode()
statuses = [
entry.get("guardrail_status")
for entry in request_data.get("litellm_metadata", {}).get("standard_logging_guardrail_information", [])
]
return raw, blocked, statuses
@pytest.mark.asyncio
async def test_v1_messages_streaming_guardrail_blocks_instead_of_leaking(
proxy_logging, make_user_api_key_auth, monkeypatch
):
"""A guardrail with its own ModelResponseStream-shaped iterator hook must not
be handed raw Anthropic SSE bytes: on /v1/messages the translation-aware
apply_guardrail path runs, so blocked content is withheld and the violation
is recorded rather than silently logged as a successful scan."""
monkeypatch.setattr(litellm, "callbacks", [_content_filter_guardrail()])
raw, blocked, statuses = await _run_v1_messages_stream(proxy_logging, make_user_api_key_auth, "the zebra runs")
assert "zebra" not in raw
assert blocked is True
assert "guardrail_intervened" in statuses
@pytest.mark.asyncio
async def test_v1_messages_streaming_guardrail_passes_clean_content(
proxy_logging, make_user_api_key_auth, monkeypatch
):
monkeypatch.setattr(litellm, "callbacks", [_content_filter_guardrail()])
raw, blocked, statuses = await _run_v1_messages_stream(proxy_logging, make_user_api_key_auth, "the horse runs")
assert "the horse runs" in raw
assert blocked is False
assert "guardrail_intervened" not in statuses
@pytest.mark.asyncio
async def test_chat_completions_streaming_still_uses_callback_iterator_hook(
proxy_logging, make_user_api_key_auth, monkeypatch
):
"""The reroute is scoped to routes that stream provider-native chunks; the
OpenAI chat path must keep using the callback's own iterator hook."""
from litellm.integrations.custom_guardrail import CustomGuardrail
class _Guardrail(CustomGuardrail):
async def async_post_call_streaming_iterator_hook(self, **kwargs): # type: ignore[override]
async for chunk in kwargs["response"]:
yield f"{chunk}*"
async def apply_guardrail(self, **kwargs): # type: ignore[override]
raise AssertionError("apply_guardrail should not run on /chat/completions")
monkeypatch.setattr(litellm, "callbacks", [_Guardrail(guardrail_name="g", event_hook="post_call", default_on=True)])
async def gen():
for ch in ("a", "b"):
yield ch
out: List[str] = []
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
response=gen(),
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
request_data={"metadata": {"guardrails": ["g"]}},
):
out.append(chunk)
assert out == ["a*", "b*"]