From 758d5482ac515b4d90ebc5918b2e1836194ae509 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 25 Feb 2026 23:32:15 -0800 Subject: [PATCH] fix(realtime guardrails): address code review feedback - Restore session.update injection for audio/VAD path, but only when realtime_input_transcription guardrails are configured (not pre_call). Forward session.created to the client first so no error arrives before the client sees the session. - Change _swallow_next_response_create bool to int counter so consecutive blocked items are handled correctly. - Extract _build_litellm_metadata() helper to eliminate duplicated metadata-building logic across OpenAI/Azure/XAI provider branches. - Plumb litellm_metadata and user_api_key_dict to Azure and XAI handlers so guardrails work for those providers too. - Add tests for session.update injection, no-inject for pre_call-only, and consecutive-block counter. --- .../litellm_core_utils/realtime_streaming.py | 71 ++++++- litellm/llms/azure/realtime/handler.py | 10 +- litellm/realtime_api/main.py | 20 +- .../test_realtime_streaming.py | 192 ++++++++++++++++++ 4 files changed, 278 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 3866a401845..4d5258b76fd 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -70,9 +70,9 @@ 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 {} - # 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 + # Counter of pending response.create messages to swallow (incremented on each + # text-input block so consecutive blocks are handled correctly). + self._swallow_pending_response_creates: int = 0 def _should_store_message( self, @@ -257,6 +257,24 @@ class RealTimeStreaming: for cb in litellm.callbacks ) + def _has_audio_transcription_guardrails(self) -> bool: + """Return True if any callback needs to run on audio transcriptions (VAD path). + + When this returns True, we inject a session.update to disable the LLM's + auto-response so the guardrail can gate it first. + """ + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import GuardrailEventHooks + + return any( + isinstance(cb, CustomGuardrail) + and cb.should_run_guardrail( + data=self.request_data, + event_type=GuardrailEventHooks.realtime_input_transcription, + ) + for cb in litellm.callbacks + ) + async def run_realtime_guardrails( self, transcript: str, @@ -366,6 +384,23 @@ class RealTimeStreaming: ) for event in events: event_str = json.dumps(event) + ## For audio/VAD guardrail path: forward session.created first, then inject. + if ( + isinstance(event, dict) + and event.get("type") == "session.created" + and self._has_audio_transcription_guardrails() + ): + self.store_message(event_str) + await self.websocket.send_text(event_str) + await self._send_to_backend( + json.dumps( + { + "type": "session.update", + "session": {"turn_detection": {"create_response": False}}, + } + ) + ) + continue ## GUARDRAIL: run on transcription events in provider_config path too if ( isinstance(event, dict) @@ -396,6 +431,27 @@ class RealTimeStreaming: try: event_obj = json.loads(raw_response) + # For audio/VAD guardrail path: once the session is ready, tell the backend + # not to auto-respond after VAD detects end-of-speech. We send the + # session.created to the client FIRST so the client is always in sync, then + # inject the session.update so a potential error from the backend doesn't + # arrive before the client sees session.created. + if ( + event_obj.get("type") == "session.created" + and self._has_audio_transcription_guardrails() + ): + self.store_message(raw_response) + await self.websocket.send_text(raw_response) + await self._send_to_backend( + json.dumps( + { + "type": "session.update", + "session": {"turn_detection": {"create_response": False}}, + } + ) + ) + return True + if ( event_obj.get("type") == "conversation.item.input_audio_transcription.completed" @@ -467,9 +523,10 @@ 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 + # Swallow the client's response.create if any items were blocked + # (use a counter so consecutive blocks are handled correctly). + if msg_type == "response.create" and self._swallow_pending_response_creates > 0: + self._swallow_pending_response_creates -= 1 continue # block response already sent by guardrail if msg_type == "conversation.item.create": @@ -488,7 +545,7 @@ class RealTimeStreaming: combined_text ) if blocked: - self._swallow_next_response_create = True + self._swallow_pending_response_creates += 1 continue # don't forward to backend except (json.JSONDecodeError, AttributeError): diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index e533978e07a..8f4291ec271 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -6,13 +6,13 @@ This requires websockets, and is currently only supported on LiteLLM Proxy. from typing import Any, Optional, cast +from litellm._logging import verbose_proxy_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from ....litellm_core_utils.realtime_streaming import RealTimeStreaming from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..azure import AzureChatCompletion -from litellm._logging import verbose_proxy_logger # BACKEND_WS_URL = "ws://localhost:8080/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01" @@ -77,6 +77,8 @@ class AzureOpenAIRealtime(AzureChatCompletion): client: Optional[Any] = None, timeout: Optional[float] = None, realtime_protocol: Optional[str] = None, + user_api_key_dict: Optional[Any] = None, + litellm_metadata: Optional[dict] = None, ): import websockets from websockets.asyncio.client import ClientConnection @@ -101,7 +103,11 @@ class AzureOpenAIRealtime(AzureChatCompletion): ssl=ssl_context, ) as backend_ws: realtime_streaming = RealTimeStreaming( - websocket, cast(ClientConnection, backend_ws), logging_obj + websocket, + 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 045d1ef99bc..df49d4c54b2 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -32,6 +32,15 @@ vertex_llm_base = VertexBase() base_llm_http_handler = BaseLLMHTTPHandler() +def _build_litellm_metadata(kwargs: dict) -> dict: + """Build the litellm_metadata dict for guardrail checking (internal only, not forwarded to provider).""" + metadata: dict = {**(kwargs.get("litellm_metadata") or {})} + guardrails = (kwargs.get("metadata") or {}).get("guardrails") or kwargs.get("guardrails") or [] + if guardrails: + metadata["guardrails"] = guardrails + return metadata + + @wrapper_client async def _arealtime( model: str, @@ -134,6 +143,8 @@ async def _arealtime( timeout=timeout, logging_obj=litellm_logging_obj, realtime_protocol=realtime_protocol, + user_api_key_dict=kwargs.get("user_api_key_dict"), + litellm_metadata=_build_litellm_metadata(kwargs), ) elif _custom_llm_provider == "openai": api_base = ( @@ -150,11 +161,6 @@ 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, @@ -165,7 +171,7 @@ async def _arealtime( timeout=timeout, query_params=query_params, user_api_key_dict=kwargs.get("user_api_key_dict"), - litellm_metadata=_litellm_metadata, + litellm_metadata=_build_litellm_metadata(kwargs), ) elif _custom_llm_provider == "bedrock": # Extract AWS parameters from kwargs @@ -223,6 +229,8 @@ async def _arealtime( client=None, timeout=timeout, query_params=query_params, + user_api_key_dict=kwargs.get("user_api_key_dict"), + litellm_metadata=_build_litellm_metadata(kwargs), ) elif _custom_llm_provider == "vertex_ai": vertex_credentials = ( 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 c8159f11fea..47adad01ba4 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -591,6 +591,12 @@ async def test_realtime_text_input_guardrail_blocks_and_returns_error(): f"Blocked item should not be forwarded to backend, got: {forwarded_items}" ) + # ASSERT: counter was incremented and is now back to 0 after we confirmed the block + # (the loop stopped before a response.create came in, so it stays at 1) + assert streaming._swallow_pending_response_creates == 1, ( + "Counter should be 1 since no response.create arrived to consume it" + ) + litellm.callbacks = [] # cleanup @@ -623,5 +629,191 @@ async def test_realtime_text_input_guardrail_uses_pre_call_mode(): assert streaming._has_realtime_guardrails() is True, ( "pre_call guardrail should be recognized as a realtime guardrail" ) + # pre_call guardrail should NOT trigger the audio/VAD session.update injection + assert streaming._has_audio_transcription_guardrails() is False, ( + "pre_call guardrail should not trigger audio transcription guardrail path" + ) + + litellm.callbacks = [] # cleanup + + +@pytest.mark.asyncio +async def test_realtime_session_created_injects_session_update_for_audio_guardrail(): + """ + Test that when an audio transcription guardrail is configured, a session.created + event from the backend triggers a session.update injection (create_response: false) + AFTER forwarding session.created to the client. This prevents the LLM from + auto-responding before the guardrail can run on the transcript. + """ + import litellm + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import GuardrailEventHooks + + class AudioGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + guardrail = AudioGuardrail( + guardrail_name="audio-guardrail", + event_hook=GuardrailEventHooks.realtime_input_transcription, + default_on=True, + ) + litellm.callbacks = [guardrail] + + client_ws = MagicMock() + client_ws.send_text = AsyncMock() + + session_created_event = json.dumps( + {"type": "session.created", "session": {"id": "sess_abc"}} + ).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() + + # session.created must be forwarded to the client + sent_to_client = [ + json.loads(c.args[0]) for c in client_ws.send_text.call_args_list if c.args + ] + session_created_events = [e for e in sent_to_client if e.get("type") == "session.created"] + assert len(session_created_events) == 1, ( + f"session.created should be forwarded to client, got: {sent_to_client}" + ) + + # session.update must be sent to the backend AFTER session.created was forwarded + 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 one session.update injected to backend, got: {sent_to_backend}" + ) + assert session_updates[0]["session"]["turn_detection"]["create_response"] is False + + litellm.callbacks = [] # cleanup + + +@pytest.mark.asyncio +async def test_realtime_session_created_no_injection_for_pre_call_only(): + """ + Test that when only a pre_call guardrail is configured (no audio transcription), + session.created does NOT trigger the session.update injection. + """ + import litellm + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import GuardrailEventHooks + + class PreCallGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + guardrail = PreCallGuardrail( + guardrail_name="pre-call-only", + 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", "session": {"id": "sess_xyz"}} + ).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() + + # No session.update should be injected + 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) == 0, ( + f"pre_call guardrail should NOT inject session.update, got: {sent_to_backend}" + ) + + litellm.callbacks = [] # cleanup + + +@pytest.mark.asyncio +async def test_swallow_pending_response_creates_counter_consecutive_blocks(): + """ + Test that consecutive blocked items correctly increment the counter so that + each subsequent response.create is swallowed, not just the first. + """ + from fastapi import HTTPException + + import litellm + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import GuardrailEventHooks + + class AlwaysBlock(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise HTTPException(status_code=403, detail={"error": "blocked"}) + + guardrail = AlwaysBlock( + guardrail_name="always-block", + 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() + + item_create = json.dumps({ + "type": "conversation.item.create", + "item": {"role": "user", "content": [{"type": "input_text", "text": "bad text"}]}, + }) + response_create = json.dumps({"type": "response.create"}) + + # Two blocked items, each followed by a response.create + client_ws.receive_text = AsyncMock( + side_effect=[ + item_create, + response_create, + item_create, + response_create, + Exception("done"), + ] + ) + + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + await streaming.client_ack_messages() + + # Neither item_create nor response_create should have been forwarded + assert backend_ws.send.call_count == 0, ( + f"No messages should be forwarded to backend when all are blocked, " + f"got {backend_ws.send.call_count} sends" + ) + # Counter should be back to 0 (both response.creates consumed the increments) + assert streaming._swallow_pending_response_creates == 0 litellm.callbacks = [] # cleanup