diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 4d5258b76fd..231f3a975dc 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -70,9 +70,6 @@ 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 {} - # 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, @@ -523,12 +520,6 @@ class RealTimeStreaming: msg_obj = json.loads(message) msg_type = msg_obj.get("type") - # 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": # Check user text messages for prompt injection item = msg_obj.get("item", {}) @@ -545,7 +536,6 @@ class RealTimeStreaming: combined_text ) if blocked: - self._swallow_pending_response_creates += 1 continue # don't forward to backend except (json.JSONDecodeError, AttributeError): 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 47adad01ba4..8db626a4d33 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -591,12 +591,6 @@ 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 @@ -755,65 +749,3 @@ async def test_realtime_session_created_no_injection_for_pre_call_only(): 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