scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-13 03:59:12 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 66

fix(guardrails): honor decision model streaming buffering (#46086)

Browse source 
Co-authored-by: moyai-devin-berriai[bot] <336287033+moyai-devin-berriai[bot]@users.noreply.github.com>
This commit is contained in:
moyai-devin-berriai[bot] 2026-10-12 16:28:22 -07:00 • committed by GitHub
parent 231721f20a
commit 0c04d0bdb7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 506 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/proxy/guardrails/guardrail_hooks/decision_model/__init__.py
View file

@ -60,6 +60,15 @@ def _streaming_sampling_rate(litellm_params: LitellmParams) -> int:
return value
def _streaming_buffer_until_moderated(litellm_params: LitellmParams) -> bool | None:
value: Final[object] = getattr(litellm_params, "streaming_buffer_until_moderated", None)
if value is None:
return None
if not isinstance(value, bool):
raise TypeError(f"decision_model guardrail streaming_buffer_until_moderated must be a boolean, got {value!r}")
return value
def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> DecisionModelGuardrail:
import litellm
@ -85,6 +94,7 @@ def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) ->
max_concurrent_decision_calls=litellm_params.max_concurrent_decision_calls or 8,
streaming_end_of_stream_only=_streaming_end_of_stream_only(litellm_params),
streaming_sampling_rate=_streaming_sampling_rate(litellm_params),
streaming_buffer_until_moderated=_streaming_buffer_until_moderated(litellm_params),
)
litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] # callback manager is untyped
_callback

3
litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py
View file

@ -113,6 +113,7 @@ class DecisionModelGuardrail(CustomGuardrail):
max_concurrent_decision_calls: int = 8,
streaming_end_of_stream_only: bool = False,
streaming_sampling_rate: int = 5,
streaming_buffer_until_moderated: bool | None = None,
) -> None:
super().__init__( # pyright: ignore[reportUnknownMemberType] # base init takes untyped **kwargs
guardrail_name=guardrail_name,
@ -131,6 +132,8 @@ class DecisionModelGuardrail(CustomGuardrail):
self._max_concurrent_decision_calls = max_concurrent_decision_calls
self.streaming_end_of_stream_only = streaming_end_of_stream_only
self.streaming_sampling_rate = streaming_sampling_rate
if streaming_buffer_until_moderated is not None:
self.streaming_buffer_until_moderated = streaming_buffer_until_moderated
self._semaphores_by_loop: dict[AbstractEventLoop, Semaphore] = {} # mutable-ok: per-loop semaphore registry
@classmethod

239
tests/integration/streaming/test_decision_model_buffering.py Normal file
View file

@ -0,0 +1,239 @@
import json
import uuid
from pathlib import Path
from typing import Final, Literal
import pytest
import yaml
from integration._support.client import JSON_OBJECT, Gateway, list_value, object_value, string_value
from integration._support.process import owned_proxy
from integration._support.upstream import delete_scenario, register_scenario
from integration.cost_calculation.cost_tracking_case import JsonResponse, SseResponse
from pydantic import JsonValue
@pytest.mark.parametrize("endpoint", ["chat/completions", "messages", "responses"])
@pytest.mark.parametrize(
("buffering", "probability", "judge_status", "fallback"),
[
(True, 1.0, 200, "fail_closed"),
(True, 0.0, 200, "fail_closed"),
(False, 1.0, 200, "fail_closed"),
(None, 1.0, 200, "fail_closed"),
(True, 0.0, 500, "fail_closed"),
(True, 0.0, 500, "fail_open"),
],
ids=["buffered-block", "buffered-allow", "unbuffered-block", "default-block", "outage-closed", "outage-open"],
)
def test_decision_model_streaming_buffering(
gateway: Gateway,
tmp_path: Path,
buffering: bool | None,
probability: float,
endpoint: str,
judge_status: int,
fallback: Literal["fail_closed", "fail_open"],
) -> None:
identity: Final = uuid.uuid4().hex
frames: Final = _frames(identity, endpoint)
chat: Final = register_scenario(
identity,
SseResponse(content_type="text/event-stream", frames=frames),
control_url=gateway.upstream_url,
)
try:
decider: Final = register_scenario(
identity + "-decider",
JsonResponse(
content_type="application/json",
status=judge_status,
body={
"model": "scripted-decider",
"answers": [{"type": "predicate", "name": "bananas", "probability": probability}],
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
},
),
control_url=gateway.upstream_url,
)
try:
config: Final = {
**JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())),
"model_list": [
{
"model_name": "buffering-chat",
"litellm_params": {
"model": "anthropic/scripted-chat" if endpoint == "messages" else "openai/scripted-chat",
"api_key": "synthetic",
"api_base": chat.api_base(),
},
},
{
"model_name": "buffering-decider",
"litellm_params": {
"model": "openai/scripted-decider",
"api_key": "synthetic",
"api_base": decider.api_base(),
},
},
],
"guardrails": [
{
"guardrail_name": "banana-policy",
"litellm_params": {
"guardrail": "decision_model",
"mode": "post_call",
"default_on": True,
"decision_model": "buffering-decider",
"unreachable_fallback": fallback,
**({"streaming_buffer_until_moderated": buffering} if buffering is not None else {}),
"checks": [
{
"name": "bananas",
"instructions": "Does the text mention bananas?",
"action": "block",
"threshold": 0.5,
}
],
},
}
],
"general_settings": {"master_key": gateway.key},
}
config_path: Final = tmp_path / "decision-buffering.yaml"
config_path.write_text(yaml.safe_dump(config))
with owned_proxy(gateway, tmp_path, {}, config=config_path, workers=2) as candidate:
response: Final = candidate.request(
"POST",
f"/v1/{endpoint}",
{
"model": "buffering-chat",
**(
{"input": "Tell me about bananas"}
if endpoint == "responses"
else {"messages": [{"role": "user", "content": "Tell me about bananas"}]}
),
**({"max_tokens": 100} if endpoint == "messages" else {}),
"stream": True,
"cache": {"no-cache": True, "no-store": True},
},
)
events: Final = tuple(
JSON_OBJECT.validate_json(line.removeprefix("data: "))
for line in response.text.splitlines()
if line.startswith("data: {")
)
content: Final = "".join(_event_text(event, endpoint) for event in events)
if judge_status != 200 and fallback == "fail_closed":
assert "Decision model guardrail failed to respond" in response.text, response.text
assert content == "", response.text
elif probability == 0:
assert content == "Bananas are fruit.", response.text
assert '"error":' not in response.text, response.text
else:
assert "Violated decision model guardrail policy" in response.text, response.text
if buffering:
assert content == "", response.text
else:
assert content and "Bananas are fruit.".startswith(content), response.text
finally:
delete_scenario(decider)
finally:
delete_scenario(chat)
def _event_text(event: dict[str, JsonValue], endpoint: str) -> str:
if endpoint == "responses":
return string_value(event.get("delta", "")) if event.get("type") == "response.output_text.delta" else ""
if endpoint == "messages":
return string_value(object_value(event.get("delta", {})).get("text", ""))
choices: Final = list_value(event.get("choices", [{}]))
return string_value(object_value(object_value(choices[0]).get("delta", {})).get("content", ""))
def _frames(identity: str, endpoint: str) -> tuple[str, ...]:
pieces: Final = ("Ban", "anas", " are", " fruit", ".")
if endpoint == "messages":
events: Final = (
{
"type": "message_start",
"message": {
"id": identity,
"type": "message",
"role": "assistant",
"content": [],
"model": "scripted-chat",
"stop_reason": None,
"usage": {"input_tokens": 1},
},
},
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
*(
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}
for text in pieces
),
{"type": "content_block_stop", "index": 0},
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 5}},
{"type": "message_stop"},
)
return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}" for event in events)
if endpoint == "responses":
message: Final = {
"id": "msg_" + identity,
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "".join(pieces), "annotations": []}],
}
response: Final = {
"id": "resp_" + identity,
"object": "response",
"created_at": 1,
"status": "completed",
"model": "scripted-chat",
"output": [message],
"usage": {"input_tokens": 1, "output_tokens": 5, "total_tokens": 6},
}
response_events: Final = (
{
"type": "response.created",
"sequence_number": 0,
"response": {**response, "status": "in_progress", "output": []},
},
{
"type": "response.output_item.added",
"sequence_number": 1,
"output_index": 0,
"item": {**message, "content": []},
},
*(
{
"type": "response.output_text.delta",
"sequence_number": index + 2,
"item_id": message["id"],
"output_index": 0,
"content_index": 0,
"delta": text,
}
for index, text in enumerate(pieces)
),
{"type": "response.output_item.done", "sequence_number": 7, "output_index": 0, "item": message},
{"type": "response.completed", "sequence_number": 8, "response": response},
)
return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}" for event in response_events)
return (
*(
f"data: {
json.dumps(
{
'id': identity,
'object': 'chat.completion.chunk',
'created': 1,
'model': 'scripted-chat',
'choices': [{'index': 0, 'delta': {'content': text}, 'finish_reason': None}],
}
)
}"
for text in pieces
),
"data: [DONE]",
)

244
tests/integration/streaming/test_decision_model_pipeline_buffering.py Normal file
View file

@ -0,0 +1,244 @@
from __future__ import annotations
import json
import threading
import uuid
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import JSON_OBJECT, Gateway, list_value, object_value
from integration._support.process import owned_proxy
from integration._support.upstream import delete_scenario, register_scenario
from integration._support.wire import Reply, Request, wire_server
from tests.integration.cost_calculation.cost_tracking_case import SseResponse
def _frame(identity: str, content: str, finish: str | None = None) -> str:
return "data: " + json.dumps(
{
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4o-mini",
"choices": [
{
"index": 0,
"delta": {"content": content},
"finish_reason": finish,
}
],
}
)
def _content(
line: str,
content_received: threading.Event,
done_received: threading.Event,
) -> str:
if not line.startswith("data:"):
return ""
payload: Final = line.removeprefix("data:").strip()
if payload == "[DONE]":
done_received.set()
return ""
event: Final = JSON_OBJECT.validate_json(payload)
assert "error" not in event, event
pieces: Final = tuple(
object_value(object_value(choice)["delta"]).get("content", "")
for choice in list_value(event.get("choices", []))
)
assert all(isinstance(piece, str) for piece in pieces), event
text: Final = "".join(piece for piece in pieces if isinstance(piece, str))
if text:
content_received.set()
return text
def _read_stream(
candidate: Gateway,
model: str,
content_received: threading.Event,
done_received: threading.Event,
) -> str:
with candidate.client.stream(
"POST",
"/v1/chat/completions",
headers={"Authorization": f"Bearer {candidate.key}"},
json={
"model": model,
"messages": [{"role": "user", "content": "synthetic control"}],
"stream": True,
},
timeout=30,
) as response:
assert response.status_code == 200, response.read()
return "".join(_content(line, content_received, done_received) for line in response.iter_lines())
@pytest.mark.parametrize("buffering", [None, True, False], ids=["omitted", "enabled", "disabled"])
def test_detect_only_decision_pipeline_respects_buffering(
gateway: Gateway,
tmp_path: Path,
buffering: bool | None,
) -> None:
identity: Final = f"decision-buffer-{uuid.uuid4().hex}"
chat_model: Final = f"{identity}-chat"
judge_model: Final = f"{identity}-judge"
secret: Final = f"LIT9453_SECRET_{identity}"
judge_entered: Final = threading.Event()
release_judge: Final = threading.Event()
content_received: Final = threading.Event()
done_received: Final = threading.Event()
def judge(request: Request) -> Reply:
assert request.method == "POST", request
assert request.target == "/v1/decisions", request
body: Final = JSON_OBJECT.validate_json(request.body)
assert secret in json.dumps(body.get("input")), body
judge_entered.set()
assert release_judge.wait(timeout=20), "Judge gate was never released"
return Reply(
body=json.dumps(
{
"model": "gpt-6-luna",
"answers": [
{
"type": "predicate",
"name": "unsafe_output",
"probability": 0.99,
}
],
"usage": {
"input_tokens": 10,
"output_tokens": 1,
"total_tokens": 11,
},
}
).encode()
)
with gateway.scenario() as scenario:
upstream: Final = register_scenario(
identity,
SseResponse(
content_type="text/event-stream",
frames=(
_frame(identity, secret),
_frame(identity, "", "stop"),
"data: [DONE]",
),
frame_delay_ms=20,
),
control_url=gateway.upstream_url,
)
scenario.cleanups.callback(delete_scenario, upstream)
with wire_server(judge) as judge_server:
base: Final = JSON_OBJECT.validate_python(
yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
)
config: Final = {
**base,
"model_list": [
{
"model_name": chat_model,
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_base": f"{upstream.api_base()}/v1",
"api_key": "synthetic-chat-key",
},
},
{
"model_name": judge_model,
"litellm_params": {
"model": "openai/gpt-6-luna",
"api_base": f"{judge_server.url}/v1",
"api_key": "synthetic-judge-key",
},
},
],
"guardrails": [
{
"guardrail_name": identity,
"litellm_params": {
"guardrail": "decision_model",
"mode": "post_call",
"default_on": False,
"decision_model": judge_model,
"timeout": 25,
"streaming_end_of_stream_only": True,
**({"streaming_buffer_until_moderated": buffering} if buffering is not None else {}),
"checks": [
{
"name": "unsafe_output",
"instructions": ("Does the response contain the synthetic secret?"),
"action": "block",
"threshold": 0.5,
}
],
},
}
],
"policies": {
identity: {
"guardrails": {"add": [identity]},
"pipeline": {
"mode": "post_call",
"steps": [
{
"guardrail": identity,
"on_pass": "allow",
"on_fail": "next",
"on_error": "next",
}
],
},
}
},
"policy_attachments": [{"policy": identity, "models": [chat_model]}],
}
config_path: Final = tmp_path / "decision-pipeline-buffering.yaml"
config_path.write_text(yaml.safe_dump(config))
with (
owned_proxy(
gateway,
tmp_path,
{},
config=config_path,
workers=2,
) as candidate,
ThreadPoolExecutor(max_workers=1) as reader,
):
result: Final = reader.submit(
_read_stream,
candidate,
chat_model,
content_received,
done_received,
)
try:
assert judge_entered.wait(timeout=15), "The completed chat response never reached the judge"
if buffering is False:
assert content_received.wait(timeout=5), (
"Pipeline did not release explicitly unbuffered content"
)
else:
assert not content_received.wait(timeout=1), (
"Pipeline released content while its judge was held"
)
assert not result.done(), "Client stream ended before the judge verdict"
finally:
release_judge.set()
assert result.result(timeout=15) == secret
assert done_received.is_set(), "Missing terminal SSE [DONE]"
requests: Final = judge_server.drain()
assert len(requests) == 1, requests
observed: Final = JSON_OBJECT.validate_json(requests[0].body)
assert secret in json.dumps(observed.get("input")), observed

11
tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py
View file

@ -381,14 +381,22 @@ def test_initialize_applies_streaming_scan_settings(mock_mgr):
default_params: Final = _litellm_params()
default_instance: Final = initialize_guardrail(default_params, _guardrail(default_params))
configured_params: Final = _litellm_params(streaming_end_of_stream_only=True, streaming_sampling_rate=3)
configured_params: Final = _litellm_params(
streaming_end_of_stream_only=True, streaming_sampling_rate=3, streaming_buffer_until_moderated=True
)
configured_instance: Final = initialize_guardrail(configured_params, _guardrail(configured_params))
disabled_params: Final = _litellm_params(streaming_buffer_until_moderated=False)
disabled_instance: Final = initialize_guardrail(disabled_params, _guardrail(disabled_params))
unified: Final = UnifiedLLMGuardrails()
assert unified.resolve_streaming_flag(default_instance, "streaming_end_of_stream_only", False) is False
assert unified.resolve_streaming_flag(default_instance, "streaming_sampling_rate", 5) == 5
assert unified.resolve_streaming_flag(configured_instance, "streaming_end_of_stream_only", False) is True
assert unified.resolve_streaming_flag(configured_instance, "streaming_sampling_rate", 5) == 3
assert unified.resolve_streaming_flag(default_instance, "streaming_buffer_until_moderated", False) is False
assert unified.resolve_streaming_flag(default_instance, "streaming_buffer_until_moderated", True) is True
assert unified.resolve_streaming_flag(configured_instance, "streaming_buffer_until_moderated", False) is True
assert unified.resolve_streaming_flag(disabled_instance, "streaming_buffer_until_moderated", True) is False
@pytest.mark.parametrize(
@ -397,6 +405,7 @@ def test_initialize_applies_streaming_scan_settings(mock_mgr):
({"streaming_end_of_stream_only": "yes"}, TypeError, "streaming_end_of_stream_only"),
({"streaming_sampling_rate": 0}, ValueError, "streaming_sampling_rate"),
({"streaming_sampling_rate": True}, TypeError, "streaming_sampling_rate"),
({"streaming_buffer_until_moderated": "yes"}, TypeError, "streaming_buffer_until_moderated"),
],
)
@patch("litellm.logging_callback_manager")