tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-16 14:23:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(realtime guardrails): end_session_after_n_fails + Endpoint Settings wizard step (#22165)

Browse Source 
* 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>
This commit is contained in:
Ishaan Jaff
2026-02-25 23:49:03 -08:00
committed by GitHub
co-authored by greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
parent 819581f6bf
commit 965ca117bc
6 changed files with 316 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/integrations/custom_guardrail.py
+9
View File
@@ -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
litellm/litellm_core_utils/realtime_streaming.py
+25 -2
View File
@@ -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
litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/__init__.py
+6 -2
View File
@@ -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)
litellm/types/guardrails.py
+15
View File
@@ -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"
tests/test_litellm/litellm_core_utils/test_realtime_streaming.py
+107
View File
@@ -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
ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx
+154 -2
View File
@@ -118,6 +118,14 @@ const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({ visible, onClose, a
const [pendingCategorySelection, setPendingCategorySelection] = useState<string>("");
const [competitorIntentEnabled, setCompetitorIntentEnabled] = useState(false);
const [competitorIntentConfig, setCompetitorIntentConfig] = useState<any>(null);
// Endpoint Settings state (step 5)
const [selectedEndpointType, setSelectedEndpointType] = useState<string>("");
const [endSessionAfterNFails, setEndSessionAfterNFails] = useState<number | undefined>(undefined);
const [onViolation, setOnViolation] = useState<"warn" | "end_session">("warn");
const [realtimeViolationMessage, setRealtimeViolationMessage] = useState<string>("");
const [endpointSettingsOpen, setEndpointSettingsOpen] = useState<boolean>(false);
const [toolPermissionConfig, setToolPermissionConfig] = useState<ToolPermissionConfig>({
rules: [],
default_action: "deny",
@@ -361,6 +369,11 @@ const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({ 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<AddGuardrailFormProps> = ({ 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<AddGuardrailFormProps> = ({ 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<AddGuardrailFormProps> = ({ visible, onClose, a
);
};
const renderEndpointSettings = () => {
return (
<div className="space-y-6">
<div>
<p className="text-sm text-gray-500">
Configure settings for a specific call type. Most guardrails don't need this — skip it
unless you're using a specific endpoint like <code>/v1/realtime</code>.
</p>
</div>
<div>
<label className="block text-sm font-medium text-gray-700 mb-1">Call type</label>
<Select
placeholder="Select a call type"
value={selectedEndpointType || undefined}
onChange={(v) => {
setSelectedEndpointType(v);
setEndpointSettingsOpen(false);
}}
style={{ width: 260 }}
allowClear
options={[{ value: "realtime", label: "/v1/realtime" }]}
/>
<p className="text-xs text-gray-400 mt-1">More call types coming soon.</p>
</div>
{selectedEndpointType === "realtime" && (
<div className="border border-gray-200 rounded-lg overflow-hidden">
<button
type="button"
onClick={() => setEndpointSettingsOpen((o) => !o)}
className="w-full flex items-center justify-between px-4 py-3 bg-gray-50 hover:bg-gray-100 text-sm font-medium text-gray-700"
>
<span>/v1/realtime settings</span>
<svg
className={`w-4 h-4 text-gray-500 transition-transform ${endpointSettingsOpen ? "rotate-180" : ""}`}
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
strokeWidth={2}
>
<path strokeLinecap="round" strokeLinejoin="round" d="M19 9l-7 7-7-7" />
</svg>
</button>
{endpointSettingsOpen && (
<div className="space-y-5 px-4 py-4 border-t border-gray-200">
<div>
<label className="block text-sm font-medium text-gray-700 mb-1">
End session after X violations
</label>
<p className="text-xs text-gray-400 mb-2">
Automatically close the session after this many guardrail violations. Leave
empty to never auto-close.
</p>
<input
type="number"
min={1}
placeholder="e.g. 3"
value={endSessionAfterNFails ?? ""}
onChange={(e) =>
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"
/>
</div>
<div>
<label className="block text-sm font-medium text-gray-700 mb-2">
On violation
</label>
<div className="space-y-2">
{(["warn", "end_session"] as const).map((opt) => (
<label key={opt} className="flex items-start gap-2 cursor-pointer">
<input
type="radio"
name="on_violation"
value={opt}
checked={onViolation === opt}
onChange={() => setOnViolation(opt)}
className="mt-0.5"
/>
<div>
<span className="text-sm font-medium text-gray-800">
{opt === "warn" ? "Warn" : "End session"}
</span>
<p className="text-xs text-gray-400 m-0">
{opt === "warn"
? "Bot speaks the message, session continues"
: "Bot speaks the message, connection closes immediately"}
</p>
</div>
</label>
))}
</div>
</div>
<div>
<label className="block text-sm font-medium text-gray-700 mb-1">
Message the user hears
</label>
<p className="text-xs text-gray-400 mb-2">
What the bot says aloud when this guardrail fires. Falls back to the default
violation message if empty.
</p>
<textarea
rows={3}
placeholder="e.g. I'm not able to continue this conversation. Please contact us at 1-800-774-2678."
value={realtimeViolationMessage}
onChange={(e) => setRealtimeViolationMessage(e.target.value)}
className="border border-gray-300 rounded px-3 py-2 text-sm w-full resize-none"
/>
</div>
</div>
)}
</div>
)}
</div>
);
};
const getStepConfigs = () => {
if (shouldRenderContentFilterConfigSettings(selectedProvider)) {
return [
{ title: "Basic Info", optional: false },
{ title: "Default Categories", optional: false },
{ title: "Topics", optional: false },
{ title: "Patterns", optional: false },
{ title: "Keywords", optional: false },
{ title: "Endpoint Settings (Optional)", optional: true },
];
}
if (shouldRenderPIIConfigSettings(selectedProvider)) {