MAINT: Unify OpenAI realtime event routing (#2319)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: be81055b-39ed-494c-94ca-c7656200a344

Author: romanlutz

pyrit/prompt_target/openai/_openai_realtime_dispatcher.py

@@ -3,7 +3,6 @@
 
 """Concrete OpenAI Realtime event dispatcher for streaming sessions."""
 
-import base64
 import logging
 from typing import Any, ClassVar
 
@@ -13,6 +12,10 @@
     RealtimeTargetResult,
     RealtimeTurnState,
 )
+from pyrit.prompt_target.openai._openai_realtime_event_router import (
+    _OpenAIRealtimeEventKind,
+    _OpenAIRealtimeEventRouter,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -32,18 +35,19 @@ class _OpenAIRealtimeDispatcher(RealtimeEventDispatcher):
     async def _route_event_async(self, *, event: Any, state: RealtimeTurnState | None) -> None:
         """Route an OpenAI Realtime event to the active turn or to an input-side callback."""
         event_type = getattr(event, "type", "")
+        event_kind = _OpenAIRealtimeEventRouter.classify_event(event_type)
 
         # Capture audio_start_ms from speech_started for the next committed event.
         # The server reports it reliably here but omits it from the commit event itself.
         # Do not return — the downstream state-aware branch still needs to fire the
         # barge-in cancel when speech starts mid-response.
-        if event_type == "input_audio_buffer.speech_started":
+        if event_kind is _OpenAIRealtimeEventKind.SPEECH_STARTED:
             speech_start = getattr(event, "audio_start_ms", None)
             if speech_start is not None:
                 self._pending_speech_start_ms = speech_start
 
         # Input-side events fire callbacks regardless of whether a turn is registered.
-        if event_type == "input_audio_buffer.committed":
+        if event_kind is _OpenAIRealtimeEventKind.INPUT_COMMITTED:
             item_id = getattr(event, "item_id", None)
             if item_id is None:
                 return
@@ -63,32 +67,33 @@ async def _route_event_async(self, *, event: Any, state: RealtimeTurnState | Non
         if state is None or state.completion.done():
             return
 
-        if event_type == "response.created":
+        _OpenAIRealtimeEventRouter.collect_response_delta(
+            event=event,
+            event_kind=event_kind,
+            audio_buffer=state.delivered_audio,
+            transcripts=state.delivered_transcripts,
+        )
+
+        if event_kind is _OpenAIRealtimeEventKind.RESPONSE_CREATED:
             state.is_responding = True
             response = getattr(event, "response", None)
             if response is not None:
                 state.last_response_id = getattr(response, "id", None)
             return
 
-        if event_type in ("response.output_item.added", "response.output_item.created"):
+        if event_kind is _OpenAIRealtimeEventKind.OUTPUT_ITEM:
             item = getattr(event, "item", None)
             if item is not None:
                 state.current_item_id = getattr(item, "id", None)
             return
 
-        if event_type in ("response.audio.delta", "response.output_audio.delta"):
-            delta = getattr(event, "delta", "")
-            if delta:
-                state.delivered_audio.extend(base64.b64decode(delta))
+        if event_kind is _OpenAIRealtimeEventKind.AUDIO_DELTA:
             return
 
-        if event_type in ("response.audio_transcript.delta", "response.output_audio_transcript.delta"):
-            delta = getattr(event, "delta", "")
-            if delta:
-                state.delivered_transcripts.append(delta)
+        if event_kind is _OpenAIRealtimeEventKind.TRANSCRIPT_DELTA:
             return
 
-        if event_type == "response.done":
+        if event_kind is _OpenAIRealtimeEventKind.RESPONSE_DONE:
             response = getattr(event, "response", None)
             done_response_id = getattr(response, "id", None) if response is not None else None
             if state.last_response_id is not None and done_response_id != state.last_response_id:
@@ -103,7 +108,7 @@ async def _route_event_async(self, *, event: Any, state: RealtimeTurnState | Non
             )
             return
 
-        if event_type == "input_audio_buffer.speech_started" and state.is_responding:
+        if event_kind is _OpenAIRealtimeEventKind.SPEECH_STARTED and state.is_responding:
             await self._cancel_async(state=state)
             state.is_responding = False
             state.completion.set_result(
@@ -115,7 +120,7 @@ async def _route_event_async(self, *, event: Any, state: RealtimeTurnState | Non
             )
             return
 
-        if event_type == "error":
+        if event_kind is _OpenAIRealtimeEventKind.ERROR:
             error = getattr(event, "error", None)
             code = getattr(error, "code", None) if error is not None else None
             message = getattr(error, "message", "unknown") if error is not None else "unknown"

pyrit/prompt_target/openai/_openai_realtime_event_router.py

@@ -0,0 +1,106 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""Shared OpenAI Realtime event classification and response accumulation."""
+
+import base64
+from enum import Enum, auto
+from typing import Any, ClassVar
+
+
+class _OpenAIRealtimeEventKind(Enum):
+    """Provider event categories shared by atomic and streaming receive policies."""
+
+    RESPONSE_DONE = auto()
+    ERROR = auto()
+    AUDIO_DELTA = auto()
+    AUDIO_DONE = auto()
+    TRANSCRIPT_DELTA = auto()
+    OUTPUT_TEXT_DONE = auto()
+    RESPONSE_CREATED = auto()
+    OUTPUT_ITEM = auto()
+    SPEECH_STARTED = auto()
+    INPUT_COMMITTED = auto()
+    LIFECYCLE = auto()
+    OTHER = auto()
+
+
+class _OpenAIRealtimeEventRouter:
+    """Classify provider events and apply response deltas to caller-owned buffers."""
+
+    _KINDS_BY_EVENT_TYPE: ClassVar[dict[str, _OpenAIRealtimeEventKind]] = {
+        "response.done": _OpenAIRealtimeEventKind.RESPONSE_DONE,
+        "error": _OpenAIRealtimeEventKind.ERROR,
+        "response.audio.delta": _OpenAIRealtimeEventKind.AUDIO_DELTA,
+        "response.output_audio.delta": _OpenAIRealtimeEventKind.AUDIO_DELTA,
+        "response.audio.done": _OpenAIRealtimeEventKind.AUDIO_DONE,
+        "response.output_audio.done": _OpenAIRealtimeEventKind.AUDIO_DONE,
+        "response.audio_transcript.delta": _OpenAIRealtimeEventKind.TRANSCRIPT_DELTA,
+        "response.output_audio_transcript.delta": _OpenAIRealtimeEventKind.TRANSCRIPT_DELTA,
+        "response.output_text.done": _OpenAIRealtimeEventKind.OUTPUT_TEXT_DONE,
+        "response.created": _OpenAIRealtimeEventKind.RESPONSE_CREATED,
+        "response.output_item.added": _OpenAIRealtimeEventKind.OUTPUT_ITEM,
+        "response.output_item.created": _OpenAIRealtimeEventKind.OUTPUT_ITEM,
+        "input_audio_buffer.speech_started": _OpenAIRealtimeEventKind.SPEECH_STARTED,
+        "input_audio_buffer.committed": _OpenAIRealtimeEventKind.INPUT_COMMITTED,
+    }
+    _LIFECYCLE_EVENT_TYPES: ClassVar[frozenset[str]] = frozenset(
+        {
+            "session.created",
+            "session.updated",
+            "conversation.created",
+            "conversation.item.created",
+            "conversation.item.added",
+            "conversation.item.done",
+            "input_audio_buffer.speech_stopped",
+            "conversation.item.input_audio_transcription.completed",
+            "response.output_item.done",
+            "response.content_part.added",
+            "response.content_part.done",
+            "response.audio_transcript.done",
+            "response.output_audio_transcript.done",
+            "response.output_text.delta",
+            "rate_limits.updated",
+        }
+    )
+    _LIFECYCLE_KINDS: ClassVar[frozenset[_OpenAIRealtimeEventKind]] = frozenset(
+        {
+            _OpenAIRealtimeEventKind.RESPONSE_CREATED,
+            _OpenAIRealtimeEventKind.OUTPUT_ITEM,
+            _OpenAIRealtimeEventKind.SPEECH_STARTED,
+            _OpenAIRealtimeEventKind.INPUT_COMMITTED,
+            _OpenAIRealtimeEventKind.LIFECYCLE,
+        }
+    )
+
+    @classmethod
+    def classify_event(cls, event_type: str) -> _OpenAIRealtimeEventKind:
+        """Return the normalized category for a provider event type."""
+        event_kind = cls._KINDS_BY_EVENT_TYPE.get(event_type)
+        if event_kind is not None:
+            return event_kind
+        if event_type in cls._LIFECYCLE_EVENT_TYPES:
+            return _OpenAIRealtimeEventKind.LIFECYCLE
+        return _OpenAIRealtimeEventKind.OTHER
+
+    @classmethod
+    def is_lifecycle_event(cls, event_kind: _OpenAIRealtimeEventKind) -> bool:
+        """Return whether atomic receiving should log the event as lifecycle-only."""
+        return event_kind in cls._LIFECYCLE_KINDS
+
+    @staticmethod
+    def collect_response_delta(
+        *,
+        event: Any,
+        event_kind: _OpenAIRealtimeEventKind,
+        audio_buffer: bytearray,
+        transcripts: list[str],
+    ) -> None:
+        """Apply an audio or transcript delta to caller-owned response buffers."""
+        delta = getattr(event, "delta", "")
+        if not delta:
+            return
+        if event_kind is _OpenAIRealtimeEventKind.AUDIO_DELTA:
+            audio_buffer.extend(base64.b64decode(delta))
+        elif event_kind is _OpenAIRealtimeEventKind.TRANSCRIPT_DELTA:
+            transcripts.append(delta)

pyrit/prompt_target/openai/openai_realtime_target.py

@@ -24,6 +24,10 @@
 from pyrit.prompt_target.common.target_capabilities import TargetCapabilities
 from pyrit.prompt_target.common.target_configuration import TargetConfiguration
 from pyrit.prompt_target.common.utils import limit_requests_per_minute
+from pyrit.prompt_target.openai._openai_realtime_event_router import (
+    _OpenAIRealtimeEventKind,
+    _OpenAIRealtimeEventRouter,
+)
 from pyrit.prompt_target.openai._openai_realtime_streaming_session import (
     _OpenAIRealtimeStreamingSession,
 )
@@ -575,6 +579,7 @@ async def receive_events_async(self, conversation_id: str) -> RealtimeTargetResu
         connection = self._get_connection(conversation_id=conversation_id)
 
         result = RealtimeTargetResult()
+        audio_buffer = bytearray()
         audio_done_received = False
         current_turn_event_count = 0
         grace_period_sec = 1.0  # Wait 1 second after audio.done before soft-finishing
@@ -595,7 +600,7 @@ async def receive_events_async(self, conversation_id: str) -> RealtimeTargetResu
                     if audio_done_received:
                         logger.warning(
                             f"Soft-finishing: No response.done {grace_period_sec}s after audio.done. "
-                            f"Audio bytes: {len(result.audio_bytes)}"
+                            f"Audio bytes: {len(audio_buffer)}"
                         )
                         break
                     # Should not happen if timeout is None, but re-raise if it does
@@ -606,22 +611,30 @@ async def receive_events_async(self, conversation_id: str) -> RealtimeTargetResu
                     break
                 except Exception as conn_err:
                     # Handle websockets connection errors as soft-finish if we have audio
-                    if "ConnectionClosed" in str(type(conn_err).__name__) and result.audio_bytes:
+                    if "ConnectionClosed" in str(type(conn_err).__name__) and audio_buffer:
                         logger.warning(
                             f"Connection closed without response.done (likely API issue). "
-                            f"Audio bytes received: {len(result.audio_bytes)}. Soft-finishing."
+                            f"Audio bytes received: {len(audio_buffer)}. Soft-finishing."
                         )
                         break
                     # Re-raise if not a connection close or no audio received
                     raise
 
                 event_type = event.type
+                event_kind = _OpenAIRealtimeEventRouter.classify_event(event_type)
                 current_turn_event_count += 1
                 logger.debug(f"Processing event type: {event_type}")
-
-                if event_type == "response.done":
+                audio_size_before = len(audio_buffer)
+                _OpenAIRealtimeEventRouter.collect_response_delta(
+                    event=event,
+                    event_kind=event_kind,
+                    audio_buffer=audio_buffer,
+                    transcripts=result.transcripts,
+                )
+
+                if event_kind is _OpenAIRealtimeEventKind.RESPONSE_DONE:
                     self._handle_response_done_event(event=event, result=result)
-                    if result.audio_bytes or current_turn_event_count > 1:
+                    if audio_buffer or current_turn_event_count > 1:
                         # Legitimate response.done: either we have audio, or other events
                         # (e.g. response.created) preceded it, confirming it belongs to this turn.
                         logger.debug("Received response.done - finishing normally")
@@ -635,53 +648,27 @@ async def receive_events_async(self, conversation_id: str) -> RealtimeTargetResu
                         "likely a stale event from a prior turn's soft-finish. Skipping."
                     )
 
-                elif event_type == "error":
+                elif event_kind is _OpenAIRealtimeEventKind.ERROR:
                     error_message = event.error.message if hasattr(event.error, "message") else str(event.error)
                     error_type = event.error.type if hasattr(event.error, "type") else "unknown"
                     logger.error(f"Received 'error' event: [{error_type}] {error_message}")
                     raise RuntimeError(f"Server error: [{error_type}] {error_message}")
 
-                elif event_type in ["response.audio.delta", "response.output_audio.delta"]:
-                    audio_data = base64.b64decode(event.delta)
-                    result.audio_bytes += audio_data
-                    logger.debug(f"Decoded {len(audio_data)} bytes of audio data")
+                elif event_kind is _OpenAIRealtimeEventKind.AUDIO_DELTA:
+                    logger.debug(f"Decoded {len(audio_buffer) - audio_size_before} bytes of audio data")
 
-                elif event_type in ["response.audio.done", "response.output_audio.done"]:
+                elif event_kind is _OpenAIRealtimeEventKind.AUDIO_DONE:
                     logger.debug(f"Received audio.done - will soft-finish in {grace_period_sec}s if no response.done")
                     audio_done_received = True
 
-                elif event_type in ["response.audio_transcript.delta", "response.output_audio_transcript.delta"]:
-                    # Capture transcript deltas as they arrive (needed when response.done never comes)
-                    if hasattr(event, "delta") and event.delta:
-                        result.transcripts.append(event.delta)
+                elif event_kind is _OpenAIRealtimeEventKind.TRANSCRIPT_DELTA:
+                    if getattr(event, "delta", ""):
                         logger.debug(f"Captured transcript delta: {event.delta[:50]}...")
 
-                elif event_type in ["response.output_text.done"]:
+                elif event_kind is _OpenAIRealtimeEventKind.OUTPUT_TEXT_DONE:
                     logger.debug("Received text.done")
 
-                # Handle lifecycle events that we can safely log
-                elif event_type in [
-                    "session.created",
-                    "session.updated",
-                    "conversation.created",
-                    "conversation.item.created",
-                    "conversation.item.added",
-                    "conversation.item.done",
-                    "input_audio_buffer.committed",
-                    "input_audio_buffer.speech_started",
-                    "input_audio_buffer.speech_stopped",
-                    "conversation.item.input_audio_transcription.completed",
-                    "response.created",
-                    "response.output_item.added",
-                    "response.output_item.created",
-                    "response.output_item.done",
-                    "response.content_part.added",
-                    "response.content_part.done",
-                    "response.audio_transcript.done",
-                    "response.output_audio_transcript.done",
-                    "response.output_text.delta",
-                    "rate_limits.updated",
-                ]:
+                elif _OpenAIRealtimeEventRouter.is_lifecycle_event(event_kind):
                     logger.debug(f"Lifecycle event '{event_type}'")
 
                 else:
@@ -691,6 +678,7 @@ async def receive_events_async(self, conversation_id: str) -> RealtimeTargetResu
             logger.error(f"An unexpected error occurred for conversation {conversation_id}: {e}")
             raise
 
+        result.audio_bytes = bytes(audio_buffer)
         logger.debug(
             f"Completed receive_events with {len(result.transcripts)} transcripts "
             f"and {len(result.audio_bytes)} bytes of audio"