diff --git a/litellm/proxy/guardrails/guardrail_hooks/decision_model/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/decision_model/__init__.py index 002b9a2cdc9..a196b5ad303 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/decision_model/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/decision_model/__init__.py @@ -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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py b/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py index c2d022850f9..867b1383682 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py +++ b/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py @@ -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 diff --git a/tests/integration/streaming/test_decision_model_buffering.py b/tests/integration/streaming/test_decision_model_buffering.py new file mode 100644 index 00000000000..a2594fd2b9b --- /dev/null +++ b/tests/integration/streaming/test_decision_model_buffering.py @@ -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]", + ) diff --git a/tests/integration/streaming/test_decision_model_pipeline_buffering.py b/tests/integration/streaming/test_decision_model_pipeline_buffering.py new file mode 100644 index 00000000000..f76e7d76d4a --- /dev/null +++ b/tests/integration/streaming/test_decision_model_pipeline_buffering.py @@ -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 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py index 54a73690840..a792e93344a 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py @@ -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")