From 965ca117bc8a9528d32fa98b6bdfc503ac951811 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 25 Feb 2026 23:49:03 -0800 Subject: [PATCH] feat(realtime guardrails): end_session_after_n_fails + Endpoint Settings wizard step (#22165) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(realtime guardrails): end_session_after_n_fails + Endpoint Settings wizard step Adds per-session violation thresholds and an optional endpoint-settings step to the guardrail wizard for /v1/realtime. Backend: - Add end_session_after_n_fails, on_violation, realtime_violation_message fields to BaseLitellmParams (no DB migration — stored in existing JSON column) - Store same fields on CustomGuardrail instance attrs - Pass through in litellm_content_filter initializer - Track _violation_count per RealTimeStreaming session; close backend_ws when on_violation=end_session OR violation count >= end_session_after_n_fails - Use realtime_violation_message as the spoken text (falls back to guardrail error string if not configured) UI (add_guardrail_form.tsx): - Rename "Default Categories" step to "Topics" - Add step 5 "Endpoint Settings (Optional)" for content filter guardrails - Call type dropdown shows /v1/realtime - Settings are in a collapsed accordion (closed by default) - "End session after X violations" + on_violation radio + spoken message field Tests: 2 new tests in test_realtime_streaming.py - test_end_session_after_n_fails_closes_connection - test_on_violation_end_session_closes_on_first_fail * fix(test): move inline imports to module level in realtime streaming tests * Update ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 9 + .../litellm_core_utils/realtime_streaming.py | 27 ++- .../litellm_content_filter/__init__.py | 8 +- litellm/types/guardrails.py | 15 ++ .../test_realtime_streaming.py | 107 ++++++++++++ .../guardrails/add_guardrail_form.tsx | 156 +++++++++++++++++- 6 files changed, 316 insertions(+), 6 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index bf330944ef..5d11fd6847 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 /v1/realtime sessions, end the session after this many violations + on_violation: For /v1/realtime sessions, 'warn' or 'end_session' + realtime_violation_message: Message the bot speaks aloud when a /v1/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 231f3a975d..6ba1b48c64 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -70,6 +70,8 @@ class RealTimeStreaming: self.session_configuration_request: Optional[str] = None self.user_api_key_dict = user_api_key_dict self.request_data: Dict = request_data or {} + # Violation counter for end_session_after_n_fails support + self._violation_count: int = 0 def _should_store_message( self, @@ -329,6 +331,10 @@ class RealTimeStreaming: safe_msg = str(detail) else: safe_msg = str(e) or "I'm sorry, that request was blocked by the content filter." + + # Use realtime_violation_message if configured; fall back to guardrail error text. + error_msg = getattr(callback, "realtime_violation_message", None) or safe_msg + # Return the error directly to the WebSocket consumer. await self.websocket.send_text( json.dumps( @@ -336,14 +342,31 @@ class RealTimeStreaming: "type": "error", "error": { "type": "guardrail_violation", - "message": safe_msg, + "message": error_msg, "code": "content_policy_violation", }, } ) ) + + 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 + ) + if should_end: + verbose_logger.warning( + "[realtime guardrail] ending session after violation %d", + self._violation_count, + ) + await self.backend_ws.close() + verbose_logger.warning( - "[realtime guardrail] BLOCKED transcript: %r", + "[realtime guardrail] BLOCKED transcript (violation %d): %r", + self._violation_count, transcript[:80], ) 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 d9a44094ad..111f8dc783 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,9 @@ 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 df411f220a..0e71f20700 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -649,6 +649,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 /v1/realtime sessions: automatically close the session after this many guardrail violations.", + ) + on_violation: Optional[Literal["warn", "end_session"]] = Field( + default=None, + description="For /v1/realtime sessions: 'warn' speaks the violation message and continues; '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 /v1/realtime 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 8db626a4d3..bcda3c7bfa 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -6,17 +6,31 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from websockets.exceptions import ConnectionClosed +import litellm + sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import ( OpenAIRealtimeStreamResponseBaseObject, OpenAIRealtimeStreamSessionEvents, ) +def _make_transcript_event(text: str, item_id: str = "item_x") -> bytes: + return json.dumps( + { + "type": "conversation.item.input_audio_transcription.completed", + "transcript": text, + "item_id": item_id, + } + ).encode() + + def test_realtime_streaming_store_message(): # Setup websocket = MagicMock() @@ -749,3 +763,96 @@ async def test_realtime_session_created_no_injection_for_pre_call_only(): 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. + """ + + 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.realtime_input_transcription, + default_on=True, + end_session_after_n_fails=2, + ) + litellm.callbacks = [guardrail] + + client_ws = MagicMock() + client_ws.send_text = AsyncMock() + + 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. + """ + + 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.realtime_input_transcription, + 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=[ + _make_transcript_event("What stock should I buy today?", item_id="item_y"), + 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/guardrails/add_guardrail_form.tsx b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx index f8c8cc17ae..3bdd18f265 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,14 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a const [pendingCategorySelection, setPendingCategorySelection] = useState(""); const [competitorIntentEnabled, setCompetitorIntentEnabled] = useState(false); const [competitorIntentConfig, setCompetitorIntentConfig] = useState(null); + + // Endpoint Settings state (step 5) + 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", @@ -361,6 +369,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); }; @@ -504,6 +517,19 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a } } + // Endpoint Settings (realtime) — content filter only + if (shouldRenderContentFilterConfigSettings(values.provider)) { + if (endSessionAfterNFails !== undefined && endSessionAfterNFails > 0) { + guardrailData.litellm_params.end_session_after_n_fails = endSessionAfterNFails; + } + if (onViolation && selectedEndpointType === "realtime") { + guardrailData.litellm_params.on_violation = onViolation; + } + if (realtimeViolationMessage.trim()) { + guardrailData.litellm_params.realtime_violation_message = realtimeViolationMessage.trim(); + } + } + /****************************** * Add provider-specific params * ---------------------------------- @@ -841,13 +867,15 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a return renderContentFilterConfiguration("keywords"); } return null; + case 4: + return renderEndpointSettings(); default: return null; } }; const renderStepButtons = () => { - const totalSteps = shouldRenderContentFilterConfigSettings(selectedProvider) ? 4 : 2; + const totalSteps = shouldRenderContentFilterConfigSettings(selectedProvider) ? 5 : 2; const isLastStep = currentStep === totalSteps - 1; const isCategoriesStep = shouldRenderContentFilterConfigSettings(selectedProvider) && currentStep === 1; const hasPendingCategory = pendingCategorySelection !== ""; @@ -888,13 +916,137 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a ); }; + const renderEndpointSettings = () => { + return ( +
+
+

+ 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. +

+
+ +
+ + + 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 aloud when this guardrail fires. Falls back to the default + violation message if empty. +

+