From ad57bf26bfac1a96ec64c5b40cf5dfffdbc0528d Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 25 Feb 2026 23:26:54 -0800 Subject: [PATCH] feat(realtime guardrails): add end_session_after_n_fails + Endpoint Settings UI step - Replace check_every_n_turns with end_session_after_n_fails: session closes after N violations instead of checking on a turn schedule - Add _violation_count to RealTimeStreaming; increment on each guardrail block - _handle_violation_action now takes explicit end_session bool computed at call site (on_violation=end_session OR violation count hits threshold) - Add end_session_after_n_fails field to BaseLitellmParams, CustomGuardrail, and litellm_content_filter initializer - Remove stale debug print statements - Add 2 new tests: end_session_after_n_fails threshold and on_violation=end_session UI: add "Endpoint Settings (Optional)" as step 5 in the guardrail wizard - Call type dropdown uses /v1/realtime (endpoint path, not marketing name) - Realtime settings live inside a collapsed accordion, closed by default - "End session after X violations" replaces the N-turns field --- litellm/integrations/custom_guardrail.py | 9 + .../litellm_core_utils/realtime_streaming.py | 82 ++++++--- .../litellm_content_filter/__init__.py | 10 +- litellm/types/guardrails.py | 15 ++ .../test_realtime_streaming.py | 112 ++++++++++++ .../src/components/ToolPolicies.tsx | 27 --- .../guardrails/add_guardrail_form.tsx | 161 +++++++++++++++++- 7 files changed, 366 insertions(+), 50 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index bf330944ef8..efbe53bbb8b 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -92,6 +92,9 @@ class CustomGuardrail(CustomLogger): mask_request_content: bool = False, mask_response_content: bool = False, violation_message_template: Optional[str] = None, + end_session_after_n_fails: Optional[int] = None, + on_violation: Optional[str] = None, + realtime_violation_message: Optional[str] = None, **kwargs, ): """ @@ -104,6 +107,9 @@ class CustomGuardrail(CustomLogger): default_on: If True, the guardrail will be run by default on all requests mask_request_content: If True, the guardrail will mask the request content mask_response_content: If True, the guardrail will mask the response content + end_session_after_n_fails: For Realtime API sessions, end the session after this many violations + on_violation: For Realtime API sessions, 'warn' or 'end_session' + realtime_violation_message: Message the bot speaks aloud when a Realtime guardrail fires """ self.guardrail_name = guardrail_name self.supported_event_hooks = supported_event_hooks @@ -114,6 +120,9 @@ class CustomGuardrail(CustomLogger): self.mask_request_content: bool = mask_request_content self.mask_response_content: bool = mask_response_content self.violation_message_template: Optional[str] = violation_message_template + self.end_session_after_n_fails: Optional[int] = end_session_after_n_fails + self.on_violation: Optional[str] = on_violation + self.realtime_violation_message: Optional[str] = realtime_violation_message if supported_event_hooks: ## validate event_hook is in supported_event_hooks diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 907f81f1e8c..472f4c03c83 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -70,6 +70,8 @@ class RealTimeStreaming: self.user_api_key_dict = user_api_key_dict # Buffer for response.text.delta events pending output-guardrail check self._pending_output_text_events: List[str] = [] + # Violation counter for end_session_after_n_fails support + self._violation_count: int = 0 def _should_store_message( self, @@ -227,6 +229,37 @@ class RealTimeStreaming: for cb in litellm.callbacks ) + async def _handle_violation_action( + self, callback: Any, safe_msg: str, end_session: bool = False + ) -> None: + """ + Speak the violation message to the user, then optionally close the session. + + end_session=True is set either when on_violation=='end_session' or when the + session-level end_session_after_n_fails threshold has been reached. + """ + spoken_msg = getattr(callback, "realtime_violation_message", None) or safe_msg + await self.backend_ws.send(json.dumps({"type": "response.cancel"})) + await self.backend_ws.send( + json.dumps( + { + "type": "response.create", + "response": { + "modalities": ["text", "audio"], + "instructions": ( + f"Say exactly and only: \"{spoken_msg}\". " + "Do not add anything else." + ), + }, + } + ) + ) + if end_session: + verbose_logger.warning( + "[realtime guardrail] ending session after violation" + ) + await self.backend_ws.close() + async def run_realtime_guardrails( self, transcript: str, @@ -235,6 +268,10 @@ class RealTimeStreaming: """ Run registered guardrails on a completed speech transcription. + On each violation, increments the session violation counter. + If end_session_after_n_fails is configured and the counter reaches that + threshold, the session is terminated after speaking the violation message. + Returns True if blocked (synthetic warning already sent to client). Returns False if clean (caller should send response.create to the backend). """ @@ -252,6 +289,7 @@ class RealTimeStreaming: is not True ): continue + try: await callback.apply_guardrail( inputs={"texts": [transcript], "images": []}, @@ -276,26 +314,22 @@ class RealTimeStreaming: safe_msg = str(detail) else: safe_msg = str(e) or "I'm sorry, that request was blocked by the content filter." - # Cancel any in-flight response before speaking the warning. - # This handles the race where create_response fired before we could intercept. - await self.backend_ws.send(json.dumps({"type": "response.cancel"})) - # Ask OpenAI to speak the warning — TTS audio plays naturally in the client - await self.backend_ws.send( - json.dumps( - { - "type": "response.create", - "response": { - "modalities": ["text", "audio"], - "instructions": ( - f"Say exactly and only: \"{safe_msg}\". " - "Do not add anything else." - ), - }, - } + + self._violation_count += 1 + end_session_after: Optional[int] = getattr( + callback, "end_session_after_n_fails", None + ) + should_end = ( + getattr(callback, "on_violation", None) == "end_session" + or ( + end_session_after is not None + and self._violation_count >= end_session_after ) ) + await self._handle_violation_action(callback, safe_msg, end_session=should_end) verbose_logger.warning( - "[realtime guardrail] BLOCKED transcript: %r", + "[realtime guardrail] BLOCKED transcript (violation %d): %r", + self._violation_count, transcript[:80], ) return True @@ -494,7 +528,6 @@ class RealTimeStreaming: """ try: event_obj = json.loads(raw_response) - if event_obj.get("type") == "session.created": # If any realtime guardrails are registered, proactively # set create_response=false so the LLM never auto-responds @@ -539,11 +572,22 @@ class RealTimeStreaming: return True if event_obj.get("type") == "response.text.delta": - if self._has_realtime_output_guardrails(): + has_guardrails = self._has_realtime_output_guardrails() + verbose_logger.warning( + "[realtime output guardrail] response.text.delta — _has_realtime_output_guardrails=%s callbacks=%s", + has_guardrails, + [type(c).__name__ for c in litellm.callbacks], + ) + if has_guardrails: self._pending_output_text_events.append(raw_response) return True if event_obj.get("type") == "response.text.done": + verbose_logger.warning( + "[realtime output guardrail] response.text.done — text=%r _has_guardrails=%s", + event_obj.get("text", "")[:60], + self._has_realtime_output_guardrails(), + ) await self._send_output_text_done(event_obj, raw_response) return True diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py index d9a44094ad2..846b0c7689a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py @@ -1,8 +1,9 @@ from typing import TYPE_CHECKING, Optional import litellm -from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import \ - ContentFilterGuardrail +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, +) from litellm.types.guardrails import SupportedGuardrailIntegrations if TYPE_CHECKING: @@ -46,6 +47,11 @@ def initialize_guardrail( competitor_intent_config=getattr( litellm_params, "competitor_intent_config", None ), + end_session_after_n_fails=getattr(litellm_params, "end_session_after_n_fails", None), + on_violation=getattr(litellm_params, "on_violation", None), + realtime_violation_message=getattr( + litellm_params, "realtime_violation_message", None + ), ) litellm.logging_callback_manager.add_litellm_callback(content_filter_guardrail) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index c14bb62fa28..7f640d77d8a 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -643,6 +643,21 @@ class BaseLitellmParams( description="Custom message when a guardrail blocks an action. Supports placeholders like {tool_name}, {rule_id}, and {default_message}.", ) + ################## Realtime API params ################ + ######################################################## + end_session_after_n_fails: Optional[int] = Field( + default=None, + description="For Realtime API sessions: automatically close the session after this many guardrail violations in a single conversation.", + ) + on_violation: Optional[Literal["warn", "end_session"]] = Field( + default=None, + description="For Realtime API sessions: 'warn' speaks the violation message and continues the session; 'end_session' speaks the message and closes the connection.", + ) + realtime_violation_message: Optional[str] = Field( + default=None, + description="The message the bot speaks aloud when a Realtime API guardrail fires. Falls back to violation_message_template if not set.", + ) + # Model Armor params template_id: Optional[str] = Field( default=None, description="The ID of your Model Armor template" diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index ae1a961548e..d806f0e0cac 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -569,3 +569,115 @@ async def test_realtime_session_created_injects_create_response_false(): ) litellm.callbacks = [] # cleanup + + +@pytest.mark.asyncio +async def test_end_session_after_n_fails_closes_connection(): + """ + Test that end_session_after_n_fails=2 closes the backend websocket after + the second guardrail violation in a session. + """ + import litellm + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import GuardrailEventHooks + + class BadWordGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + for text in inputs.get("texts", []): + if "blocked" in text.lower(): + raise ValueError("Content blocked by guardrail.") + return inputs + + guardrail = BadWordGuardrail( + guardrail_name="bad_word_guard", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + end_session_after_n_fails=2, + ) + litellm.callbacks = [guardrail] + + client_ws = MagicMock() + client_ws.send_text = AsyncMock() + + def make_transcript_event(text): + return json.dumps({ + "type": "conversation.item.input_audio_transcription.completed", + "transcript": text, + "item_id": "item_x", + }).encode() + + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[ + make_transcript_event("this is blocked"), # violation 1 — warn + make_transcript_event("also blocked again"), # violation 2 — end session + ConnectionClosed(None, None), + ] + ) + backend_ws.send = AsyncMock() + backend_ws.close = AsyncMock() + + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + await streaming.backend_to_client_send_messages() + + assert backend_ws.close.called, "Expected backend_ws.close() to be called after 2 violations" + assert streaming._violation_count == 2 + + litellm.callbacks = [] # cleanup + + +@pytest.mark.asyncio +async def test_on_violation_end_session_closes_on_first_fail(): + """ + Test that on_violation='end_session' closes the session immediately on the + first violation, regardless of end_session_after_n_fails. + """ + import litellm + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import GuardrailEventHooks + + class TopicGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + for text in inputs.get("texts", []): + if "stock" in text.lower(): + raise ValueError("Topic not allowed: financial advice.") + return inputs + + guardrail = TopicGuardrail( + guardrail_name="topic_guard", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + on_violation="end_session", + ) + litellm.callbacks = [guardrail] + + client_ws = MagicMock() + client_ws.send_text = AsyncMock() + + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[ + json.dumps({ + "type": "conversation.item.input_audio_transcription.completed", + "transcript": "What stock should I buy today?", + "item_id": "item_y", + }).encode(), + ConnectionClosed(None, None), + ] + ) + backend_ws.send = AsyncMock() + backend_ws.close = AsyncMock() + + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + await streaming.backend_to_client_send_messages() + + assert backend_ws.close.called, "Expected session to close immediately with on_violation=end_session" + assert streaming._violation_count == 1 + + litellm.callbacks = [] # cleanup diff --git a/ui/litellm-dashboard/src/components/ToolPolicies.tsx b/ui/litellm-dashboard/src/components/ToolPolicies.tsx index 82785496eae..860093ceadb 100644 --- a/ui/litellm-dashboard/src/components/ToolPolicies.tsx +++ b/ui/litellm-dashboard/src/components/ToolPolicies.tsx @@ -2,19 +2,7 @@ import React, { useCallback, useDeferredValue, useEffect, useState } from "react"; import { Select, Switch, Tooltip } from "antd"; -<<<<<<< cursor/development-environment-setup-13a7 -// @ts-ignore - duplicate import removed -import { - Table, - TableHead, - TableHeaderCell, - TableBody, - TableRow, - TableCell, -} from "@tremor/react"; -======= import { Table, TableHead, TableHeaderCell, TableBody, TableRow, TableCell } from "@tremor/react"; ->>>>>>> main import { TimeCell } from "./view_logs/time_cell"; import { TableHeaderSortDropdown } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown"; import type { SortState } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown"; @@ -60,21 +48,6 @@ const PolicySelect: React.FC<{ minWidth: 110, fontWeight: 500, }} -<<<<<<< cursor/development-environment-setup-13a7 - {...{styles: { - selector: { - backgroundColor: style.bg, - borderColor: style.border, - color: style.color, - borderRadius: 999, - fontSize: 11, - fontWeight: 600, - paddingLeft: 8, - paddingRight: 4, - }, - }} as any} -======= ->>>>>>> main popupMatchSelectWidth={false} options={POLICY_OPTIONS.map((o) => ({ value: o.value, diff --git a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx index 183763667f9..4921969ea0e 100644 --- a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx @@ -118,6 +118,13 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a const [pendingCategorySelection, setPendingCategorySelection] = useState(""); const [competitorIntentEnabled, setCompetitorIntentEnabled] = useState(false); const [competitorIntentConfig, setCompetitorIntentConfig] = useState(null); + + // Endpoint Settings state + const [selectedEndpointType, setSelectedEndpointType] = useState(""); + const [endSessionAfterNFails, setEndSessionAfterNFails] = useState(undefined); + const [onViolation, setOnViolation] = useState<"warn" | "end_session">("warn"); + const [realtimeViolationMessage, setRealtimeViolationMessage] = useState(""); + const [endpointSettingsOpen, setEndpointSettingsOpen] = useState(false); const [toolPermissionConfig, setToolPermissionConfig] = useState({ rules: [], default_action: "deny", @@ -353,6 +360,11 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a on_disallowed_action: "block", violation_message_template: "", }); + setSelectedEndpointType(""); + setEndSessionAfterNFails(undefined); + setOnViolation("warn"); + setRealtimeViolationMessage(""); + setEndpointSettingsOpen(false); setCurrentStep(0); }; @@ -482,6 +494,17 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a } } + // Endpoint Settings (realtime) — apply for content filter guardrails + if (shouldRenderContentFilterConfigSettings(values.provider)) { + if (endSessionAfterNFails !== undefined && endSessionAfterNFails > 0) { + guardrailData.litellm_params.end_session_after_n_fails = endSessionAfterNFails; + } + guardrailData.litellm_params.on_violation = onViolation; + if (realtimeViolationMessage.trim()) { + guardrailData.litellm_params.realtime_violation_message = realtimeViolationMessage.trim(); + } + } + if (guardrailProvider === "tool_permission") { if (toolPermissionConfig.rules.length === 0) { NotificationsManager.fromBackend("Add at least one tool permission rule"); @@ -811,6 +834,137 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a return ; }; + const renderEndpointSettings = () => { + return ( +
+ {/* Section header */} +
+

+ Configure settings for a specific call type. Most guardrails don't need this — skip it + unless you're using a specific endpoint like /v1/realtime. +

+
+ + {/* Endpoint selector */} +
+ + + setEndSessionAfterNFails(e.target.value ? parseInt(e.target.value, 10) : undefined) + } + className="border border-gray-300 rounded px-3 py-1.5 text-sm w-32" + /> +
+ +
+ +
+ {(["warn", "end_session"] as const).map((opt) => ( + + ))} +
+
+ +
+ +

+ What the bot says out loud when this guardrail fires. Falls back to the default + violation message if empty. +

+