scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(proxy): stream the guardrail block response when the block comes after an SSE keepalive ping (#51114)

Browse source 
* fix(proxy): stream the guardrail block response when the block lands after an SSE keepalive ping

* docs(proxy): keep the keepalive docstring's reasons together

* fix(proxy): stream the /v1/messages block after a keepalive ping and always end a late block stream

A passthrough guardrail block on a streamed /v1/messages request that landed
after the first keepalive ping still ended as an api_error event. The route now
hands its block stream to the keepalive relay, as chat and text completions do.
If a route's block stream itself fails after a ping, the relay sends the usual
late failure frame instead of raising, so the stream always ends.

The passthrough tests now drive the real routes with a real guardrail.

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
paveln83 2026-11-11 13:39:19 -08:00 • committed by GitHub
parent a1f3dcd01b
commit b63656a5de
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 457 additions and 128 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
litellm/proxy/anthropic_endpoints/endpoints.py
View file

@ -7,7 +7,7 @@ from types import MappingProxyType
from typing import Final
from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
from fastapi.responses import JSONResponse
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
@ -55,6 +55,7 @@ from litellm.proxy.litellm_pre_call_utils import (
forwarded_anthropic_oauth_token,
should_auto_drop_params_for_agentic_cli,
)
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
from litellm.types.utils import TokenCountResponse
@ -126,6 +127,32 @@ def _strip_total_tokens_from_anthropic_response(response: object) -> None:
usage.pop("total_tokens", None)
async def _record_passthrough_block(
e: ModifyResponseException, proxy_logging_obj: ProxyLogging, user_api_key_dict: UserAPIKeyAuth
) -> None:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=e.request_data,
)
async def _stream_passthrough_block(
e: ModifyResponseException,
proxy_logging_obj: ProxyLogging,
user_api_key_dict: UserAPIKeyAuth,
headers: Mapping[str, str],
) -> StreamingResponse | JSONResponse:
await _record_passthrough_block(e, proxy_logging_obj, user_api_key_dict)
block_chunks: Final = AnthropicMessagesHandler().build_block_sse_chunks(e)
async def _blocked_stream() -> AsyncGenerator[str, None]:
for chunk in block_chunks:
yield chunk.decode()
return await create_response(generator=_blocked_stream(), media_type="text/event-stream", headers=headers)
@router.post(
"/v1/messages",
tags=["Anthropic `/v1/messages`"],
@ -174,6 +201,9 @@ async def anthropic_response(
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
on_passthrough_block=lambda block: _stream_passthrough_block(
block, proxy_logging_obj, user_api_key_dict, _NO_HEADERS
),
)
# Optionally strip the non-Anthropic `usage.total_tokens` field
# LiteLLM adds internally. Anthropic's official /v1/messages spec
@ -189,30 +219,12 @@ async def anthropic_response(
_strip_total_tokens_from_anthropic_response(result)
return result
except ModifyResponseException as e:
# Guardrail flagged content in passthrough mode - return 200 with violation message
_data: Final = e.request_data
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=_data,
)
block_headers: Final = litellm_call_id_headers(base_llm_response_processor.litellm_call_id) or _NO_HEADERS
block_handler: Final = AnthropicMessagesHandler()
if data.get("stream") is True:
block_chunks: Final = block_handler.build_block_sse_chunks(e)
async def _blocked_stream() -> AsyncGenerator[str, None]:
for chunk in block_chunks:
yield chunk.decode()
return await create_response(
generator=_blocked_stream(), media_type="text/event-stream", headers=block_headers
)
return await _stream_passthrough_block(e, proxy_logging_obj, user_api_key_dict, block_headers)
await _record_passthrough_block(e, proxy_logging_obj, user_api_key_dict)
fastapi_response.headers.update(block_headers)
return block_handler.build_block_response(e)
return AnthropicMessagesHandler().build_block_response(e)
except AnthropicContextManagementError as e:
if e.status_code >= 500:
# Server-side polyfill failures hit the failure hook for spend/alert

26
litellm/proxy/common_request_processing.py
View file

@ -1662,6 +1662,18 @@ async def _relay_late_response(produced: Response, framing: LateSseFraming) -> A
await _aclose_late_response(produced)
async def _late_block_or_failure(
exc: Exception,
on_passthrough_block: Callable[[ModifyResponseException], Awaitable[Response]] | None,
) -> Response | Exception:
if not isinstance(exc, ModifyResponseException) or on_passthrough_block is None:
return exc
try:
return await on_passthrough_block(exc)
except Exception as block_failure: # noqa: BLE001 # the client still needs a frame that ends the stream
return block_failure
async def _sanitized_late_failure(
exc: Exception,
on_late_failure: "Callable[[Exception], Awaitable[HTTPException | None]] | None",
@ -1692,6 +1704,7 @@ async def open_sse_before_first_byte(
media_type: str = "text/event-stream",
on_late_failure: Callable[[Exception], Awaitable[HTTPException | None]] | None = None,
late_framing: LateSseFraming = OPENAI_LATE_SSE_FRAMING,
on_passthrough_block: Callable[[ModifyResponseException], Awaitable[Response]] | None = None,
) -> _LateResponseT | StreamingResponse:
"""Write SSE keepalive comments while the upstream LLM call is still in flight.
@ -1708,6 +1721,8 @@ async def open_sse_before_first_byte(
the first ping reaches the client as an SSE error frame under a 200 rather than
as an HTTP error status, and LiteLLM's own `x-litellm-*` response headers are
not yet known. Both are why this stays off until an operator sets an interval.
A passthrough guardrail block is not a failure: when the route passes
`on_passthrough_block`, the stream it builds for the block is relayed.
"""
interval: Final = coerce_keepalive_interval(ping_interval_seconds)
if interval is None:
@ -1730,14 +1745,19 @@ async def open_sse_before_first_byte(
try:
produced: Final = produce_task.result()
except Exception as exc: # noqa: BLE001 # the status line is already sent; surface it as a frame
late: Final = await _late_block_or_failure(exc, on_passthrough_block)
if isinstance(late, Response):
async for chunk in _relay_late_response(late, late_framing):
yield chunk
return
verbose_proxy_logger.exception(
"request failed after its SSE keepalive had opened the response: %s", exc
"request failed after its SSE keepalive had opened the response: %s", late
)
# The caller's own `except` never sees this, so its failure hook
# would never fire and the failure would go unaudited. The hook
# also gets to sanitize what reaches the client, by returning or
# raising a replacement, so its answer decides the frame.
for frame in late_framing.failure(await _sanitized_late_failure(exc, on_late_failure)):
for frame in late_framing.failure(await _sanitized_late_failure(late, on_late_failure)):
yield frame.encode()
return
async for chunk in _relay_late_response(produced, late_framing):
@ -3156,6 +3176,7 @@ class ProxyBaseLLMRequestProcessing:
contents: list[object] | None = None,
skip_pre_call_logic: bool = False,
sse_framing: SseFraming = OPENAI_SSE_FRAMING,
on_passthrough_block: Callable[[ModifyResponseException], Awaitable[Response]] | None = None,
) -> object:
"""Run the request, sending SSE keepalives while the upstream is still silent.
@ -3207,6 +3228,7 @@ class ProxyBaseLLMRequestProcessing:
ping_interval_seconds=ttft_keepalive_interval(self.data, llm_router),
on_late_failure=_audit_late_failure,
late_framing=ANTHROPIC_LATE_SSE_FRAMING if answers_as_anthropic else sse_framing.late,
on_passthrough_block=on_passthrough_block,
)
def total_timeout_for(

162
litellm/proxy/proxy_server.py
View file

@ -210,8 +210,10 @@ from litellm.rust_bridge.catalog import Target
from litellm.types.utils import (
PRICING_OVERRIDES_KEY,
CallTypes,
Choices,
CountContextManagement,
LlmProviders,
Message,
ModelResponse,
ModelResponseStream,
TextCompletionResponse,
@ -13263,6 +13265,70 @@ def _blocked_response_usage(original_response: object | None) -> "litellm.Usage"
return litellm.Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0)
def _passthrough_block_response(e: ModifyResponseException) -> ModelResponse:
return ModelResponse(
model=e.model,
choices=[Choices(message=Message(content=e.message), finish_reason="content_filter")],
usage=_blocked_response_usage(e.original_response),
)
async def _record_passthrough_block(e: ModifyResponseException, user_api_key_dict: UserAPIKeyAuth) -> None:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=e.request_data,
)
async def _stream_chat_passthrough_block(
e: ModifyResponseException, user_api_key_dict: UserAPIKeyAuth, request: Request
) -> StreamingResponse:
# Read before post_call_failure_hook, which pops it from the request data.
logging_obj: Final = e.request_data.get("litellm_logging_obj")
await _record_passthrough_block(e, user_api_key_dict)
streaming_response: Final = litellm.CustomStreamWrapper(
completion_stream=litellm.utils.ModelResponseIterator(
model_response=_passthrough_block_response(e), convert_to_delta=True
),
model=e.model,
custom_llm_provider="cached_response",
logging_obj=logging_obj,
)
return StreamingResponse(
select_data_generator(
response=streaming_response,
user_api_key_dict=user_api_key_dict,
request_data=e.request_data,
request=request,
),
media_type="text/event-stream",
status_code=200,
)
async def _stream_text_completion_passthrough_block(
e: ModifyResponseException, user_api_key_dict: UserAPIKeyAuth, request: Request
) -> StreamingResponse:
await _record_passthrough_block(e, user_api_key_dict)
streaming_response: Final = litellm.TextCompletionStreamWrapper(
completion_stream=litellm.utils.ModelResponseIterator(
model_response=_passthrough_block_response(e), convert_to_delta=True
),
model=e.model,
)
return StreamingResponse(
select_data_generator(
response=streaming_response,
user_api_key_dict=user_api_key_dict,
request_data=e.request_data,
request=request,
),
media_type="text/event-stream",
status_code=200,
)
@router.post(
"/v1/chat/completions",
dependencies=[Depends(user_api_key_auth)],
@ -13352,50 +13418,17 @@ async def chat_completion(
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
on_passthrough_block=lambda block: _stream_chat_passthrough_block(block, user_api_key_dict, request),
)
if isinstance(result, BaseModel):
return model_dump_with_preserved_fields(result, exclude_unset=True)
else:
return result
except ModifyResponseException as e:
# Guardrail flagged content in passthrough mode - return 200 with violation message
_data = e.request_data
# Capture logging_obj before post_call_failure_hook pops it from _data.
_logging_obj: Final = _data.get("litellm_logging_obj")
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=_data,
)
_chat_response = litellm.ModelResponse()
_chat_response.model = e.model
_chat_response.choices[0].message.content = e.message
_chat_response.choices[0].finish_reason = "content_filter"
# Report the blocked LLM response's real usage (set before the stream
# branch so both paths carry it); zero for pre-call blocks.
_chat_response.usage = _blocked_response_usage(e.original_response)
if data.get("stream", None) is not None and data["stream"] is True:
_iterator = litellm.utils.ModelResponseIterator(model_response=_chat_response, convert_to_delta=True)
_streaming_response = litellm.CustomStreamWrapper(
completion_stream=_iterator,
model=e.model,
custom_llm_provider="cached_response",
logging_obj=_logging_obj,
)
selected_data_generator = select_data_generator(
response=_streaming_response,
user_api_key_dict=user_api_key_dict,
request_data=_data,
request=request,
)
return StreamingResponse(
selected_data_generator,
media_type="text/event-stream",
status_code=200, # Return 200 for passthrough mode
)
return _chat_response
if data.get("stream") is True:
return await _stream_chat_passthrough_block(e, user_api_key_dict, request)
await _record_passthrough_block(e, user_api_key_dict)
return _passthrough_block_response(e)
except RejectedRequestError as e:
_data = e.request_data
await proxy_logging_obj.post_call_failure_hook(
@ -13509,49 +13542,20 @@ async def completion(
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
version=version,
on_passthrough_block=lambda block: _stream_text_completion_passthrough_block(
block, user_api_key_dict, request
),
)
except ModifyResponseException as e:
# Guardrail flagged content in passthrough mode - return 200 with violation message
_data = e.request_data
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=_data,
)
if _data.get("stream", None) is not None and _data["stream"] is True:
_text_response: Final = litellm.ModelResponse()
# Set text attribute dynamically for text completion format
setattr(_text_response.choices[0], "text", e.message)
_text_response.model = e.model
_usage = _blocked_response_usage(e.original_response)
# Set usage attribute dynamically (ModelResponse accepts usage in __init__ but it's not in type definition)
setattr(_text_response, "usage", _usage)
_iterator = litellm.utils.ModelResponseIterator(model_response=_text_response, convert_to_delta=True)
_streaming_response = litellm.TextCompletionStreamWrapper(
completion_stream=_iterator,
model=e.model,
)
selected_data_generator = select_data_generator(
response=_streaming_response,
user_api_key_dict=user_api_key_dict,
request_data=_data,
request=request,
)
return StreamingResponse(
selected_data_generator,
media_type="text/event-stream",
status_code=200, # Return 200 for passthrough mode
)
else:
_response = litellm.TextCompletionResponse()
_response.choices[0].text = e.message
_response.model = e.model
_usage = _blocked_response_usage(e.original_response)
_response.usage = _usage
return _response
if e.request_data.get("stream") is True:
return await _stream_text_completion_passthrough_block(e, user_api_key_dict, request)
await _record_passthrough_block(e, user_api_key_dict)
text_completion_response: Final = litellm.TextCompletionResponse()
text_completion_response.choices[0].text = e.message
text_completion_response.choices[0].finish_reason = "content_filter"
text_completion_response.model = e.model
text_completion_response.usage = _blocked_response_usage(e.original_response)
return text_completion_response
except RejectedRequestError as e:
_data = e.request_data
await proxy_logging_obj.post_call_failure_hook(

1
litellm/utils.py
View file

@ -9228,6 +9228,7 @@ class ModelResponseIterator:
if convert_to_delta is True:
_stream_response: Final = ModelResponseStream()
_stream_response.choices[0].delta.content = model_response.choices[0].message.content
_stream_response.choices[0].finish_reason = model_response.choices[0].finish_reason
self.model_response: ModelResponse | ModelResponseStream = _stream_response
else:
self.model_response = model_response

81
tests/unit/proxy/anthropic_endpoints/test_endpoints.py
View file

@ -2,6 +2,7 @@
Test for anthropic_endpoints/endpoints.py, focusing on handling dictionary objects in streaming responses
"""
import asyncio
import datetime
import json
import logging
@ -85,10 +86,15 @@ _BLOCK_CALL_ID: Final = "messages-block"
_BLOCK_TEXT: Final = "blocked by guardrail"
class _BlockEveryRequest(CustomGuardrail):
class _BlockEveryRequestOnceReleased(CustomGuardrail):
def __init__(self, released: asyncio.Event) -> None:
super().__init__(guardrail_name="deny-list", event_hook=GuardrailEventHooks.pre_call, default_on=True)
self.released: Final = released
async def async_pre_call_hook(
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: str
) -> None:
await self.released.wait()
self.raise_passthrough_exception(violation_message=_BLOCK_TEXT, request_data=data)
@ -117,21 +123,29 @@ def _messages_request(body: Mapping[str, object]) -> Request:
async def _blocked_messages_reply(
body: Mapping[str, object], fastapi_response: Response, rails: tuple[CustomGuardrail, ...]
body: Mapping[str, object],
fastapi_response: Response,
rails: tuple[CustomGuardrail, ...],
keepalive_interval: float | None = None,
) -> tuple[object, str]:
import litellm
import litellm.proxy.anthropic_endpoints.endpoints as ep
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.utils import ProxyLogging
blocker: Final = _BlockEveryRequest(
guardrail_name="deny-list", event_hook=GuardrailEventHooks.pre_call, default_on=True
)
released: Final = asyncio.Event()
if keepalive_interval is None:
released.set()
with (
patch.object(
proxy_server, "proxy_logging_obj", new=ProxyLogging(user_api_key_cache=DualCache())
), # test-quality-ok: module global imported at call time; no injection seam
patch.object(litellm, "callbacks", new=[blocker, *rails]), # test-quality-ok: guardrails register through this global
patch.object(
litellm, "callbacks", new=[_BlockEveryRequestOnceReleased(released), *rails]
), # test-quality-ok: guardrails register through this global
patch.object(
litellm, "sse_keepalive_ping_interval_seconds", new=keepalive_interval
), # test-quality-ok: the operator's keepalive interval is read from this global
):
response = await ep.anthropic_response(
fastapi_response=fastapi_response,
@ -140,12 +154,34 @@ async def _blocked_messages_reply(
)
if not isinstance(response, StreamingResponse):
return response, ""
streamed: Final = b"".join(
[chunk if isinstance(chunk, bytes) else chunk.encode() async for chunk in response.body_iterator]
).decode()
async def chunks() -> AsyncGenerator[bytes, None]:
async for chunk in response.body_iterator:
released.set()
yield chunk if isinstance(chunk, bytes) else str(chunk).encode()
streamed: Final = b"".join([chunk async for chunk in chunks()]).decode()
return response, streamed
def _anthropic_events(streamed: str) -> tuple[tuple[str, Mapping[str, object]], ...]:
frames: Final = tuple(frame for frame in streamed.split("\n\n") if frame.strip() and frame != ": ping")
return tuple(
(frame.split("\n")[0].removeprefix("event: "), json.loads(frame.split("\n")[1].removeprefix("data: ")))
for frame in frames
)
_BLOCK_EVENTS: Final = (
"message_start",
"content_block_start",
"content_block_delta",
"content_block_stop",
"message_delta",
"message_stop",
)
class TestBlockedResponseUsage:
"""Blocked responses report the blocked LLM response's real usage."""
@ -220,26 +256,27 @@ class TestBlockedResponseUsage:
{"model": "claude-haiku-5-5", "stream": True}, Response(), rails
)
frames: Final = tuple(frame for frame in streamed.split("\n\n") if frame.strip())
events: Final = tuple(
(frame.split("\n")[0].removeprefix("event: "), json.loads(frame.split("\n")[1].removeprefix("data: ")))
for frame in frames
)
events: Final = _anthropic_events(streamed)
assert isinstance(response, StreamingResponse)
assert response.media_type == "text/event-stream"
assert response.headers["x-litellm-call-id"] == _BLOCK_CALL_ID
assert [name for name, _ in events] == [
"message_start",
"content_block_start",
"content_block_delta",
"content_block_stop",
"message_delta",
"message_stop",
]
assert [name for name, _ in events] == list(_BLOCK_EVENTS)
assert all(payload["type"] == name for name, payload in events)
assert events[2][1]["delta"] == {"type": "text_delta", "text": _BLOCK_TEXT}
assert events[4][1]["delta"]["stop_reason"] == "end_turn"
@pytest.mark.asyncio
async def test_a_streaming_block_after_a_keepalive_ping_answers_the_block_events(self) -> None:
_, streamed = await _blocked_messages_reply(
{"model": "claude-haiku-5-5", "stream": True}, Response(), (), keepalive_interval=0.05
)
events: Final = _anthropic_events(streamed)
assert streamed.startswith(": ping\n\n")
assert [name for name, _ in events] == list(_BLOCK_EVENTS)
assert events[2][1]["delta"] == {"type": "text_delta", "text": _BLOCK_TEXT}
assert events[4][1]["delta"]["stop_reason"] == "end_turn"
@pytest.mark.asyncio
async def test_json_block_answers_with_the_call_id(self) -> None:
fastapi_response: Final = Response()

63
tests/unit/proxy/test_common_request_processing.py
View file

@ -11608,6 +11608,69 @@ async def test_a_broken_hook_does_not_replace_the_real_error_with_its_own_bug():
assert "audit backend" not in collected[-2].decode()
async def _block_once_released(release: asyncio.Event) -> Response:
await release.wait()
raise ModifyResponseException(
message="blocked by policy", model="gpt-4o", request_data={"stream": True}, guardrail_name="g"
)
async def _drain_releasing(response: StreamingResponse, release: asyncio.Event) -> tuple[bytes, ...]:
async def chunks() -> AsyncGenerator[bytes, None]:
async for chunk in response.body_iterator:
release.set()
yield chunk if isinstance(chunk, bytes) else str(chunk).encode()
return tuple([chunk async for chunk in chunks()])
@pytest.mark.asyncio
async def test_a_late_passthrough_block_streams_the_routes_block_response_instead_of_an_error_frame():
block_frames: Final = ('data: {"choices": [{"delta": {"content": "blocked by policy"}}]}\n\n', "data: [DONE]\n\n")
release: Final = asyncio.Event()
stream_block: Final = AsyncMock(side_effect=lambda _block: _sse_response(block_frames))
record_failure: Final = AsyncMock(return_value=None)
response: Final = await open_sse_before_first_byte(
_block_once_released(release),
ping_interval_seconds=0.05,
on_late_failure=record_failure,
on_passthrough_block=stream_block,
)
collected: Final = await _drain_releasing(response, release)
assert collected[0] == TTFT_PING
assert [chunk for chunk in collected if chunk != TTFT_PING] == [frame.encode() for frame in block_frames]
assert [call.args[0].message for call in stream_block.await_args_list] == ["blocked by policy"]
record_failure.assert_not_awaited()
@pytest.mark.asyncio
async def test_a_late_passthrough_block_on_a_route_without_a_block_stream_still_ends_the_stream():
release: Final = asyncio.Event()
response: Final = await open_sse_before_first_byte(_block_once_released(release), ping_interval_seconds=0.05)
collected: Final = await _drain_releasing(response, release)
assert collected[0] == TTFT_PING
assert "error" in json.loads(collected[-2].decode().removeprefix("data: ").strip())
assert collected[-1] == b"data: [DONE]\n\n"
@pytest.mark.asyncio
async def test_a_late_passthrough_block_whose_block_stream_fails_still_ends_the_stream():
release: Final = asyncio.Event()
broken_block_stream: Final = AsyncMock(side_effect=RuntimeError("block stream unavailable"))
response: Final = await open_sse_before_first_byte(
_block_once_released(release), ping_interval_seconds=0.05, on_passthrough_block=broken_block_stream
)
collected: Final = await _drain_releasing(response, release)
assert collected[0] == TTFT_PING
assert "error" in json.loads(collected[-2].decode().removeprefix("data: ").strip())
assert collected[-1] == b"data: [DONE]\n\n"
@pytest.mark.parametrize(
"exc,expect_traceback",
[

196
tests/unit/proxy/test_modify_response_streaming_passthrough.py
View file

@ -14,16 +14,28 @@ Reading it from the outer body passed ``logging_obj=None`` to
``AttributeError: 'NoneType' object has no attribute 'model_call_details'``.
"""
import asyncio
import json
from collections.abc import AsyncGenerator, Mapping
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import Request, Response
from fastapi.responses import StreamingResponse
from starlette.types import Message
import litellm
import litellm.proxy.proxy_server as proxy_server
from litellm.caching.dual_cache import DualCache
from litellm.exceptions import RejectedRequestError
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.proxy_server import chat_completion
from litellm.proxy.proxy_server import chat_completion, completion
from litellm.proxy.utils import ProxyLogging
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import ModelResponse, ModelResponseStream, TextCompletionResponse
async def _run_streaming_block_and_get_wrapper(exception):
"""Drive chat_completion's streaming guardrail-passthrough handler for the
@ -108,3 +120,181 @@ async def test_streaming_rejected_request_uses_request_data_logging_obj():
mock_csw.assert_called_once()
assert mock_csw.call_args.kwargs["logging_obj"] is sentinel_logging_obj
_BLOCK_TEXT: Final = "blocked by guardrail"
_PING: Final = ": ping"
class _BlockEveryRequestOnceReleased(CustomGuardrail):
def __init__(self, released: asyncio.Event) -> None:
super().__init__(guardrail_name="deny-list", event_hook=GuardrailEventHooks.pre_call, default_on=True)
self.released: Final = released
self.failures: tuple[Exception, ...] = ()
async def async_pre_call_hook(
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: str
) -> None:
await self.released.wait()
self.raise_passthrough_exception(violation_message=_BLOCK_TEXT, request_data=data)
async def async_post_call_failure_hook(
self,
request_data: dict[str, object],
original_exception: Exception,
user_api_key_dict: UserAPIKeyAuth,
traceback_str: str | None = None,
) -> None:
self.failures = (*self.failures, original_exception)
class _RejectEveryRequest(CustomLogger):
async def async_pre_call_hook(
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: str
) -> None:
raise RejectedRequestError(
message="rejected by the pre-call hook", model="gpt-4o", llm_provider="openai", request_data=data
)
def _post(path: str, body: Mapping[str, object]) -> Request:
async def receive() -> Message:
return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False}
return Request(
{
"type": "http",
"method": "POST",
"path": path,
"query_string": b"",
"headers": [(b"content-type", b"application/json")],
},
receive,
)
async def _frames_releasing_the_block(response: StreamingResponse, released: asyncio.Event) -> tuple[str, ...]:
async def chunks() -> AsyncGenerator[str, None]:
async for chunk in response.body_iterator:
released.set()
yield chunk if isinstance(chunk, str) else bytes(chunk).decode()
streamed: Final = "".join([chunk async for chunk in chunks()])
return tuple(frame for frame in streamed.split("\n\n") if frame)
async def _reply(
path: str,
body: Mapping[str, object],
rail: CustomLogger,
released: asyncio.Event,
keepalive_interval: float | None,
) -> tuple[object, tuple[str, ...]]:
endpoint: Final = chat_completion if path == "/v1/chat/completions" else completion
with (
patch.object(
proxy_server, "proxy_logging_obj", new=ProxyLogging(user_api_key_cache=DualCache())
), # test-quality-ok: module global imported at call time; no injection seam
patch.object(litellm, "callbacks", new=[rail]), # test-quality-ok: guardrails register through this global
patch.object(
litellm, "sse_keepalive_ping_interval_seconds", new=keepalive_interval
), # test-quality-ok: the operator's keepalive interval is read from this global
):
response: Final = await endpoint(
request=_post(path, body), fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth()
)
if not isinstance(response, StreamingResponse):
return response, ()
return response, await _frames_releasing_the_block(response, released)
async def _blocked_reply(
path: str, body: Mapping[str, object], keepalive_interval: float | None
) -> tuple[object, tuple[str, ...], _BlockEveryRequestOnceReleased]:
released: Final = asyncio.Event()
if keepalive_interval is None:
released.set()
blocker: Final = _BlockEveryRequestOnceReleased(released)
response, frames = await _reply(path, body, blocker, released, keepalive_interval)
return response, frames, blocker
def _data_frames(frames: tuple[str, ...]) -> tuple[str, ...]:
return tuple(frame.removeprefix("data: ") for frame in frames if frame.startswith("data: {"))
KEEPALIVE_INTERVALS: Final = pytest.mark.parametrize(
"keepalive_interval", [0.05, None], ids=["blocked-after-a-keepalive-ping", "keepalive-off"]
)
@pytest.mark.asyncio
@KEEPALIVE_INTERVALS
async def test_a_blocked_chat_stream_gets_the_block_message_and_a_content_filter_finish(
keepalive_interval: float | None,
) -> None:
_, frames, blocker = await _blocked_reply(
"/v1/chat/completions", {"model": "gpt-4o", "messages": [], "stream": True}, keepalive_interval
)
chunks: Final = tuple(ModelResponseStream.model_validate_json(frame) for frame in _data_frames(frames))
assert (frames[0] == _PING) is (keepalive_interval is not None)
assert frames[-1] == "data: [DONE]"
assert not any('"error"' in frame for frame in frames)
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == _BLOCK_TEXT
assert chunks[-1].choices[0].finish_reason == "content_filter"
assert [type(failure) for failure in blocker.failures] == [ModifyResponseException]
@pytest.mark.asyncio
@KEEPALIVE_INTERVALS
async def test_a_blocked_text_completion_stream_gets_the_block_message(keepalive_interval: float | None) -> None:
_, frames, blocker = await _blocked_reply(
"/v1/completions", {"model": "gpt-4o", "prompt": "hi", "stream": True}, keepalive_interval
)
chunks: Final = tuple(TextCompletionResponse.model_validate_json(frame) for frame in _data_frames(frames))
assert (frames[0] == _PING) is (keepalive_interval is not None)
assert frames[-1] == "data: [DONE]"
assert not any('"error"' in frame for frame in frames)
assert "".join(chunk.choices[0].text or "" for chunk in chunks) == _BLOCK_TEXT
assert chunks[-1].choices[0].finish_reason == "content_filter"
assert [type(failure) for failure in blocker.failures] == [ModifyResponseException]
@pytest.mark.asyncio
async def test_a_blocked_chat_request_without_streaming_gets_the_block_message_and_a_content_filter_finish() -> None:
response, _, blocker = await _blocked_reply(
"/v1/chat/completions", {"model": "gpt-4o", "messages": [], "stream": False}, None
)
assert isinstance(response, ModelResponse)
assert response.choices[0].message.content == _BLOCK_TEXT
assert response.choices[0].finish_reason == "content_filter"
assert [type(failure) for failure in blocker.failures] == [ModifyResponseException]
@pytest.mark.asyncio
async def test_a_blocked_text_completion_without_streaming_gets_the_block_text_and_a_content_filter_finish() -> None:
response, _, blocker = await _blocked_reply(
"/v1/completions", {"model": "gpt-4o", "prompt": "hi", "stream": False}, None
)
assert isinstance(response, TextCompletionResponse)
assert response.choices[0].text == _BLOCK_TEXT
assert response.choices[0].finish_reason == "content_filter"
assert [type(failure) for failure in blocker.failures] == [ModifyResponseException]
@pytest.mark.asyncio
async def test_a_rejected_text_completion_stream_ends_with_a_stop_finish() -> None:
released: Final = asyncio.Event()
released.set()
_, frames = await _reply(
"/v1/completions", {"model": "gpt-4o", "prompt": "hi", "stream": True}, _RejectEveryRequest(), released, None
)
chunks: Final = tuple(TextCompletionResponse.model_validate_json(frame) for frame in _data_frames(frames))
assert frames[-1] == "data: [DONE]"
assert chunks[-1].choices[0].finish_reason == "stop"