scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

fix(realtime): guardrails with pre_call/post_call mode now work on realtime WebSocket; return error directly to consumer

Browse source
This commit is contained in:
Ishaan Jaffer 2026-02-25 23:14:56 -08:00
parent 8192819bda
commit f3300e6cc6
4 changed files with 155 additions and 108 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

105
litellm/litellm_core_utils/realtime_streaming.py
View file

@ -43,6 +43,7 @@ class RealTimeStreaming:
provider_config: Optional[BaseRealtimeConfig] = None,
model: str = "",
user_api_key_dict: Optional[Any] = None,
request_data: Optional[Dict] = None,
):
self.websocket = websocket
self.backend_ws = backend_ws
@ -68,6 +69,10 @@ class RealTimeStreaming:
self.current_delta_type: Optional[ALL_DELTA_TYPES] = None
self.session_configuration_request: Optional[str] = None
self.user_api_key_dict = user_api_key_dict
self.request_data: Dict = request_data or {}
# Set to True after a text-input guardrail block so we can swallow the client's
# subsequent response.create (which would conflict with the block response).
self._swallow_next_response_create: bool = False
def _should_store_message(
self,
@ -231,15 +236,23 @@ class RealTimeStreaming:
await self.backend_ws.send(message)
def _has_realtime_guardrails(self) -> bool:
"""Return True if any callback is registered for realtime_input_transcription."""
"""Return True if any callback is registered for realtime guardrail event types."""
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.types.guardrails import GuardrailEventHooks
_realtime_event_types = [
GuardrailEventHooks.realtime_input_transcription,
GuardrailEventHooks.pre_call,
GuardrailEventHooks.post_call,
]
return any(
isinstance(cb, CustomGuardrail)
and cb.should_run_guardrail(
data={},
event_type=GuardrailEventHooks.realtime_input_transcription,
and any(
cb.should_run_guardrail(
data=self.request_data,
event_type=et,
)
for et in _realtime_event_types
)
for cb in litellm.callbacks
)
@ -258,17 +271,25 @@ class RealTimeStreaming:
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.types.guardrails import GuardrailEventHooks
_realtime_event_types = [
GuardrailEventHooks.realtime_input_transcription,
GuardrailEventHooks.pre_call,
GuardrailEventHooks.post_call,
]
_check_data = {**self.request_data, "transcript": transcript}
_already_run: set = set()
for callback in litellm.callbacks:
if not isinstance(callback, CustomGuardrail):
continue
if (
callback.should_run_guardrail(
data={"transcript": transcript},
event_type=GuardrailEventHooks.realtime_input_transcription,
)
is not True
if id(callback) in _already_run:
continue
if not any(
callback.should_run_guardrail(data=_check_data, event_type=et)
for et in _realtime_event_types
):
continue
_already_run.add(id(callback))
try:
await callback.apply_guardrail(
inputs={"texts": [transcript], "images": []},
@ -293,20 +314,15 @@ 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._send_to_backend(json.dumps({"type": "response.cancel"}))
# Ask the model to speak the warning — TTS audio plays naturally in the client
await self._send_to_backend(
# Return the error directly to the WebSocket consumer.
await self.websocket.send_text(
json.dumps(
{
"type": "response.create",
"response": {
"modalities": ["text", "audio"],
"instructions": (
f"Say exactly and only: \"{safe_msg}\". "
"Do not add anything else."
),
"type": "error",
"error": {
"type": "guardrail_violation",
"message": safe_msg,
"code": "content_policy_violation",
},
}
)
@ -348,23 +364,6 @@ class RealTimeStreaming:
if isinstance(transformed_response, list)
else [transformed_response]
)
for event in events:
## GUARDRAIL: inject create_response=false on session.created
if isinstance(event, dict) and event.get("type") == "session.created":
if self._has_realtime_guardrails():
await self._send_to_backend(
json.dumps(
{
"type": "session.update",
"session": {
"turn_detection": {
"type": "server_vad",
"create_response": False,
}
},
}
)
)
for event in events:
event_str = json.dumps(event)
## GUARDRAIL: run on transcription events in provider_config path too
@ -397,28 +396,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
# before our guardrail has a chance to run.
if self._has_realtime_guardrails():
await self._send_to_backend(
json.dumps(
{
"type": "session.update",
"session": {
"turn_detection": {
"type": "server_vad",
"create_response": False,
}
},
}
)
)
verbose_logger.debug(
"[realtime guardrail] injected create_response=false into session"
)
if (
event_obj.get("type")
== "conversation.item.input_audio_transcription.completed"
@ -490,6 +467,11 @@ class RealTimeStreaming:
msg_obj = json.loads(message)
msg_type = msg_obj.get("type")
# Swallow the client's response.create if we just blocked an item.
if msg_type == "response.create" and self._swallow_next_response_create:
self._swallow_next_response_create = False
continue # block response already sent by guardrail
if msg_type == "conversation.item.create":
# Check user text messages for prompt injection
item = msg_obj.get("item", {})
@ -506,6 +488,7 @@ class RealTimeStreaming:
combined_text
)
if blocked:
self._swallow_next_response_create = True
continue # don't forward to backend
except (json.JSONDecodeError, AttributeError):

2
litellm/llms/openai/realtime/handler.py
View file

@ -99,6 +99,7 @@ class OpenAIRealtime(OpenAIChatCompletion):
timeout: Optional[float] = None,
query_params: Optional[RealtimeQueryParams] = None,
user_api_key_dict: Optional[Any] = None,
litellm_metadata: Optional[dict] = None,
**kwargs: Any,
):
import websockets
@ -142,6 +143,7 @@ class OpenAIRealtime(OpenAIChatCompletion):
cast(ClientConnection, backend_ws),
logging_obj,
user_api_key_dict=user_api_key_dict,
request_data={"litellm_metadata": litellm_metadata or {}},
)
await realtime_streaming.bidirectional_forward()

6
litellm/realtime_api/main.py
View file

@ -150,6 +150,11 @@ async def _arealtime(
or get_secret_str("OPENAI_API_KEY")
)
# Build metadata for guardrail checking.
_litellm_metadata: dict = {**(kwargs.get("litellm_metadata") or {})}
_guardrails = (kwargs.get("metadata") or {}).get("guardrails") or kwargs.get("guardrails") or []
if _guardrails:
_litellm_metadata["guardrails"] = _guardrails
await openai_realtime.async_realtime(
model=model,
websocket=websocket,
@ -160,6 +165,7 @@ async def _arealtime(
timeout=timeout,
query_params=query_params,
user_api_key_dict=kwargs.get("user_api_key_dict"),
litellm_metadata=_litellm_metadata,
)
elif _custom_llm_provider == "bedrock":
# Extract AWS parameters from kwargs

150
tests/test_litellm/litellm_core_utils/test_realtime_streaming.py
View file

@ -416,32 +416,32 @@ async def test_realtime_guardrail_blocks_prompt_injection():
streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj)
await streaming.backend_to_client_send_messages()
# ASSERT 1: no bare response.create was sent to backend (injection blocked).
# The only response.create allowed is the warning one (has "instructions" field).
# ASSERT 1: no response.create was sent to backend (injection blocked).
sent_to_backend = [
json.loads(c.args[0])
for c in backend_ws.send.call_args_list
if c.args
]
bare_response_creates = [
response_creates = [
e for e in sent_to_backend
if e.get("type") == "response.create"
and "instructions" not in e.get("response", {})
]
assert len(bare_response_creates) == 0, (
f"Guardrail should prevent bare response.create for injected content, "
f"but got: {bare_response_creates}"
assert len(response_creates) == 0, (
f"Guardrail should prevent response.create for injected content, "
f"but got: {response_creates}"
)
# ASSERT 2: warning response.create was sent to backend (to speak the block message)
warning_creates = [
e for e in sent_to_backend
if e.get("type") == "response.create"
and "instructions" in e.get("response", {})
# ASSERT 2: error event was sent directly to the client WebSocket
sent_to_client = [
json.loads(c.args[0]) for c in client_ws.send_text.call_args_list
if c.args
]
assert len(warning_creates) > 0, (
f"Backend should receive a response.create with warning instructions, "
f"but got: {sent_to_backend}"
error_events = [e for e in sent_to_client if e.get("type") == "error"]
assert len(error_events) == 1, (
f"Expected one error event sent to client, got: {sent_to_client}"
)
assert error_events[0]["error"]["type"] == "guardrail_violation", (
f"Expected guardrail_violation error type, got: {error_events[0]}"
)
litellm.callbacks = [] # cleanup
@ -514,11 +514,91 @@ async def test_realtime_guardrail_allows_clean_transcript():
@pytest.mark.asyncio
async def test_realtime_session_created_injects_create_response_false():
async def test_realtime_text_input_guardrail_blocks_and_returns_error():
"""
Test that when session.created arrives from the backend and realtime guardrails
are registered, the proxy injects a session.update with create_response=False
so the LLM never auto-responds before the guardrail runs.
Test that when conversation.item.create arrives with text that triggers a guardrail,
the proxy blocks it (doesn't forward to backend) and returns an error event directly
to the client WebSocket.
"""
from fastapi import HTTPException
import litellm
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.types.guardrails import GuardrailEventHooks
class BlockingGuardrail(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
texts = inputs.get("texts", [])
for text in texts:
if "@" in text:
raise HTTPException(
status_code=403,
detail={"error": "email address detected"},
)
return inputs
guardrail = BlockingGuardrail(
guardrail_name="email-blocker",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
)
litellm.callbacks = [guardrail]
client_ws = MagicMock()
client_ws.send_text = AsyncMock()
backend_ws = MagicMock()
backend_ws.send = AsyncMock()
backend_ws.recv = AsyncMock(side_effect=ConnectionClosed(None, None))
logging_obj = MagicMock()
logging_obj.pre_call = MagicMock()
streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj)
item_create_msg = json.dumps({
"type": "conversation.item.create",
"item": {
"role": "user",
"content": [{"type": "input_text", "text": "My email is test@example.com"}],
},
})
# Simulate the client sending a conversation.item.create with an email
client_ws.receive_text = AsyncMock(
side_effect=[
item_create_msg,
Exception("connection closed"), # stop the loop
]
)
await streaming.client_ack_messages()
# ASSERT: error event was sent to client
assert client_ws.send_text.called, "Expected error to be sent to client websocket"
sent_texts = [json.loads(c.args[0]) for c in client_ws.send_text.call_args_list]
error_events = [e for e in sent_texts if e.get("type") == "error"]
assert len(error_events) == 1, f"Expected one error event, got: {sent_texts}"
assert error_events[0]["error"]["type"] == "guardrail_violation"
# ASSERT: blocked item was NOT forwarded to the backend
sent_to_backend = [c.args[0] for c in backend_ws.send.call_args_list if c.args]
forwarded_items = [
json.loads(m) for m in sent_to_backend
if isinstance(m, str) and json.loads(m).get("type") == "conversation.item.create"
]
assert len(forwarded_items) == 0, (
f"Blocked item should not be forwarded to backend, got: {forwarded_items}"
)
litellm.callbacks = [] # cleanup
@pytest.mark.asyncio
async def test_realtime_text_input_guardrail_uses_pre_call_mode():
"""
Test that _has_realtime_guardrails returns True for a guardrail configured with
pre_call mode (not just realtime_input_transcription).
"""
import litellm
from litellm.integrations.custom_guardrail import CustomGuardrail
@ -529,43 +609,19 @@ async def test_realtime_session_created_injects_create_response_false():
return inputs
guardrail = DummyGuardrail(
guardrail_name="dummy",
event_hook=GuardrailEventHooks.realtime_input_transcription,
guardrail_name="pre-call-guardrail",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
)
litellm.callbacks = [guardrail]
client_ws = MagicMock()
client_ws.send_text = AsyncMock()
session_created_event = json.dumps({"type": "session.created"}).encode()
backend_ws = MagicMock()
backend_ws.recv = AsyncMock(
side_effect=[
session_created_event,
ConnectionClosed(None, None),
]
)
backend_ws.send = 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: proxy injected session.update with create_response=False to backend
sent_to_backend = [
json.loads(c.args[0]) for c in backend_ws.send.call_args_list if c.args
]
session_updates = [e for e in sent_to_backend if e.get("type") == "session.update"]
assert len(session_updates) == 1, (
f"Expected proxy to inject session.update, got: {sent_to_backend}"
)
td = session_updates[0]["session"]["turn_detection"]
assert td["create_response"] is False, (
f"Expected create_response=False, got: {td}"
assert streaming._has_realtime_guardrails() is True, (
"pre_call guardrail should be recognized as a realtime guardrail"
)
litellm.callbacks = [] # cleanup