From f3300e6cc6e6d2bb39db958657f46e21c3b58192 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 25 Feb 2026 23:14:56 -0800 Subject: [PATCH] fix(realtime): guardrails with pre_call/post_call mode now work on realtime WebSocket; return error directly to consumer --- .../litellm_core_utils/realtime_streaming.py | 105 +++++------- litellm/llms/openai/realtime/handler.py | 2 + litellm/realtime_api/main.py | 6 + .../test_realtime_streaming.py | 150 ++++++++++++------ 4 files changed, 155 insertions(+), 108 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 5d7a5bfe318..3866a401845 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -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): diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index c2fccfc7289..05915e36a69 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -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() diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index e4c8f648190..045d1ef99bc 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -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 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 aaaab95ce6f..c8159f11fea 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -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