Retry OpenAI realtime responses that fail before output

livekit-plugins/livekit-plugins-openai/livekit/plugins/openai/realtime/realtime_model.py

@@ -264,6 +264,9 @@ class _ResponseGeneration:
     """timestamp when the response was created"""
     _first_token_timestamp: float | None = None
     """timestamp when the first token was received"""
+    response_create_params: RealtimeResponseCreateParams | None = None
+    retry_count: int = 0
+    output_started: bool = False
 
     def _close(self) -> None:
         for msg in self.messages.values():
@@ -299,6 +302,26 @@ def _is_fatal_error(error: object | None) -> bool:
     return isinstance(code, str) and code in _FATAL_ERROR_CODES
 
 
+def _response_done_error(*, provider_label: str, event: ResponseDoneEvent) -> APIError | None:
+    if event.response.status != "failed":
+        return None
+
+    if event.response.status_details and hasattr(event.response.status_details, "error"):
+        error_type = getattr(event.response.status_details.error, "type", "unknown")
+        error_body = event.response.status_details.error
+        message = f"{provider_label} response failed with error type: {error_type}"
+    else:
+        error_body = None
+        message = f"{provider_label} response failed with unknown error"
+
+    recoverable = not _is_fatal_error(error_body)
+    return APIError(
+        message=message,
+        body=error_body,
+        retryable=recoverable,
+    )
+
+
 class RealtimeModel(llm.RealtimeModel):
     @overload
     def __init__(
@@ -857,12 +880,15 @@ def __init__(self, realtime_model: RealtimeModel) -> None:
         self.send_event(self._create_session_update_event())
 
         self._response_created_futures: dict[str, asyncio.Future[llm.GenerationCreatedEvent]] = {}
+        self._response_create_params: dict[str, RealtimeResponseCreateParams] = {}
         self._item_delete_future: dict[str, asyncio.Future] = {}
         self._item_create_future: dict[str, asyncio.Future] = {}
 
         # generate_reply event_ids cancelled or timed out before response.created arrived; the
         # response is cancelled by id and discarded when it finally arrives
         self._discarded_event_ids: set[str] = set()
+        self._response_retry_event_ids: set[str] = set()
+        self._response_retry_generations: dict[str, _ResponseGeneration] = {}
 
         # accumulates partial input-audio transcripts per (item_id, content_index)
         self._input_transcript_accumulators: dict[str, dict[int, str]] = {}
@@ -1626,12 +1652,12 @@ def generate_reply(
         if is_given(tools):
             params.tools = self._convert_tools_to_oai(tools)  # type: ignore
 
-        self.send_event(
-            ResponseCreateEvent(type="response.create", event_id=event_id, response=params)
-        )
+        self._response_create_params[event_id] = params
+        self._send_response_create(event_id=event_id, params=params)
 
         def _on_timeout() -> None:
             self._response_created_futures.pop(event_id, None)
+            self._response_create_params.pop(event_id, None)
             if fut and not fut.done():
                 # discard the response if the server still creates it after the timeout
                 self._discarded_event_ids.add(event_id)
@@ -1642,6 +1668,7 @@ def _on_timeout() -> None:
         def _on_fut_done(f: asyncio.Future[llm.GenerationCreatedEvent]) -> None:
             handle.cancel()
             self._response_created_futures.pop(event_id, None)
+            self._response_create_params.pop(event_id, None)
             if f.cancelled():
                 # response.create was already sent; cancel the response server-side
                 self.send_event(ResponseCancelEvent(type="response.cancel"))
@@ -1652,6 +1679,11 @@ def _on_fut_done(f: asyncio.Future[llm.GenerationCreatedEvent]) -> None:
         fut.add_done_callback(_on_fut_done)
         return fut
 
+    def _send_response_create(self, *, event_id: str, params: RealtimeResponseCreateParams) -> None:
+        self.send_event(
+            ResponseCreateEvent(type="response.create", event_id=event_id, response=params)
+        )
+
     @property
     def has_active_generation(self) -> bool:
         return self._current_generation is not None or len(self._response_created_futures) > 0
@@ -1780,6 +1812,17 @@ def _handle_response_created(self, event: ResponseCreatedEvent) -> None:
         if isinstance(event.response.metadata, dict):
             client_event_id = event.response.metadata.get("client_event_id")
 
+        if client_event_id and client_event_id in self._response_retry_event_ids:
+            self._response_retry_event_ids.discard(client_event_id)
+            retry_generation = self._response_retry_generations.pop(client_event_id, None)
+            if not isinstance(retry_generation, _ResponseGeneration) or (
+                retry_generation is not self._current_generation
+            ):
+                return
+            retry_generation._created_timestamp = time.time()
+            retry_generation._first_token_timestamp = None
+            return
+
         if client_event_id and client_event_id in self._discarded_event_ids:
             # interrupted or timed out before the server created it: cancel by id and mark it
             # discarded so its trailing events are skipped, instead of surfacing it
@@ -1795,6 +1838,9 @@ def _handle_response_created(self, event: ResponseCreatedEvent) -> None:
             message_ch=utils.aio.Chan(),
             function_ch=utils.aio.Chan(),
             messages={},
+            response_create_params=self._response_create_params.pop(client_event_id, None)
+            if client_event_id
+            else None,
             _created_timestamp=time.time(),
             _done_fut=asyncio.Future(),
         )
@@ -1819,6 +1865,7 @@ def _handle_response_output_item_added(self, event: ResponseOutputItemAddedEvent
         if isinstance(self._current_generation, _DiscardedGeneration):
             return
         assert self._current_generation is not None, "current_generation is None"
+        self._current_generation.output_started = True
         assert (item_id := event.item.id) is not None, "item.id is None"
         assert (item_type := event.item.type) is not None, "item.type is None"
 
@@ -2077,6 +2124,9 @@ def _handle_response_done(self, event: ResponseDoneEvent) -> None:
 
         assert self._current_generation is not None, "current_generation is None"
 
+        if self._maybe_retry_failed_response(event):
+            return
+
         created_timestamp = self._current_generation._created_timestamp
         first_token_timestamp = self._current_generation._first_token_timestamp
 
@@ -2152,6 +2202,58 @@ def _handle_response_done(self, event: ResponseDoneEvent) -> None:
         self.emit("metrics_collected", metrics)
         self._handle_response_done_but_not_complete(event)
 
+    def _maybe_retry_failed_response(self, event: ResponseDoneEvent) -> bool:
+        if event.response.status != "failed":
+            return False
+        if not isinstance(self._current_generation, _ResponseGeneration):
+            return False
+        if self._current_generation.output_started:
+            return False
+        if self._current_generation.response_create_params is None:
+            return False
+
+        error = _response_done_error(
+            provider_label=self._realtime_model._provider_label, event=event
+        )
+        if error is None:
+            return False
+        if not error.retryable:
+            return False
+
+        max_retries = self._opts.conn_options.max_retry
+        if self._current_generation.retry_count >= max_retries:
+            return False
+
+        retry_count = self._current_generation.retry_count
+        retry_interval = self._opts.conn_options._interval_for_retry(retry_count)
+        self._current_generation.retry_count += 1
+        retry_generation = self._current_generation
+
+        event_id = utils.shortuuid("response_retry_")
+        params = self._current_generation.response_create_params.model_copy(deep=True)
+        params.metadata = {"client_event_id": event_id}
+        self._response_retry_event_ids.add(event_id)
+        self._response_retry_generations[event_id] = retry_generation
+        self._emit_error(error, recoverable=True)
+
+        async def _retry_response() -> None:
+            if retry_interval > 0:
+                await asyncio.sleep(retry_interval)
+            if self._msg_ch.closed or self._current_generation is not retry_generation:
+                self._response_retry_event_ids.discard(event_id)
+                self._response_retry_generations.pop(event_id, None)
+                return
+            logger.warning(
+                "%s realtime response failed before output, retrying in %.1fs",
+                self._realtime_model._provider_label,
+                retry_interval,
+                extra={"attempt": retry_count, "max_retries": max_retries},
+            )
+            self._send_response_create(event_id=event_id, params=params)
+
+        asyncio.create_task(_retry_response(), name="RealtimeSession.retry_response")
+        return True
+
     def _handle_response_done_but_not_complete(self, event: ResponseDoneEvent) -> None:
         """Handle response done but not complete, i.e. cancelled, incomplete or failed.
 
@@ -2165,23 +2267,9 @@ def _handle_response_done_but_not_complete(self, event: ResponseDoneEvent) -> No
 
         provider_label = self._realtime_model._provider_label
         if event.response.status == "failed":
-            if event.response.status_details and hasattr(event.response.status_details, "error"):
-                error_type = getattr(event.response.status_details.error, "type", "unknown")
-                error_body = event.response.status_details.error
-                message = f"{provider_label} response failed with error type: {error_type}"
-            else:
-                error_body = None
-                message = f"{provider_label} response failed with unknown error"
-            # failures are largely undocumented by openai, so we assume optimistically
-            # recoverable unless the code is a known-fatal one (quota / auth / billing),
-            # which is raised so the recv loop breaks and _main_task stops reconnecting
-            recoverable = not _is_fatal_error(error_body)
-            error = APIError(
-                message=message,
-                body=error_body,
-                retryable=recoverable,
-            )
-            if not recoverable:
+            error = _response_done_error(provider_label=provider_label, event=event)
+            assert error is not None
+            if not error.retryable:
                 raise error
             self._emit_error(error, recoverable=True)
         elif event.response.status in {"cancelled", "incomplete"}:

tests/test_realtime/test_openai_realtime_model.py

@@ -1,18 +1,72 @@
 from __future__ import annotations
 
+import asyncio
+import types
 from types import SimpleNamespace
 from typing import cast
 
 import pytest
 
-from livekit.agents import llm
+from livekit.agents import APIConnectOptions, llm
 from livekit.agents._exceptions import APIError
 from livekit.agents.llm.remote_chat_context import RemoteChatContext
 from livekit.plugins.openai.realtime.realtime_model import RealtimeSession, _is_fatal_error
 
 pytestmark = pytest.mark.unit
 
 
+def _create_response_retry_session(*, retry_interval: float = 0) -> RealtimeSession:
+    session = RealtimeSession.__new__(RealtimeSession)
+    session._opts = SimpleNamespace(
+        conn_options=APIConnectOptions(max_retry=1, retry_interval=retry_interval)
+    )
+    session._realtime_model = SimpleNamespace(
+        _provider_label="openai",
+        _label="openai",
+        label="openai",
+        model="gpt-realtime",
+        provider="openai",
+    )
+    session._msg_ch = SimpleNamespace(closed=False)
+    session._response_created_futures = {}
+    session._response_create_params = {}
+    session._response_retry_event_ids = set()
+    session._response_retry_generations = {}
+    session._discarded_event_ids = set()
+    session._current_generation = None
+    session._instructions = None
+    session._sent_events = []
+    session._errors = []
+    session.send_event = types.MethodType(
+        lambda self, event: self._sent_events.append(event),
+        session,
+    )
+    session.emit = types.MethodType(
+        lambda self, event_name, event: self._errors.append((event_name, event)),
+        session,
+    )
+    return session
+
+
+def _response_created(*, client_event_id: str, response_id: str = "resp_1") -> SimpleNamespace:
+    return SimpleNamespace(
+        response=SimpleNamespace(id=response_id, metadata={"client_event_id": client_event_id})
+    )
+
+
+def _response_failed(*, code: str = "rate_limit_exceeded") -> SimpleNamespace:
+    return SimpleNamespace(
+        response=SimpleNamespace(
+            id="resp_1",
+            status="failed",
+            usage=None,
+            status_details=SimpleNamespace(
+                error=SimpleNamespace(type="server_error", code=code),
+            ),
+        )
+    )
+
+
 def test_update_chat_ctx_deletes_empty_remote_items() -> None:
     remote_ctx = RemoteChatContext()
     audio_item = llm.ChatMessage(id="audio_item", role="user", content=[])
@@ -119,3 +173,77 @@ def test_response_done_failed_transient_stays_recoverable() -> None:
     )
     RealtimeSession._handle_response_done_but_not_complete(session, event)
     assert captured["recoverable"] is True
+
+
+async def test_response_done_failed_retries_before_output() -> None:
+    session = _create_response_retry_session()
+
+    fut = session.generate_reply(instructions="say hi")
+    create_event = session._sent_events[-1]
+    client_event_id = create_event.response.metadata["client_event_id"]
+    session._handle_response_created(_response_created(client_event_id=client_event_id))
+
+    generation = await fut
+    assert generation.response_id == "resp_1"
+
+    session._handle_response_done(_response_failed())
+    await asyncio.sleep(0.15)
+
+    assert len(session._sent_events) == 2
+    retry_event = session._sent_events[-1]
+    assert retry_event.type == "response.create"
+    assert retry_event.event_id.startswith("response_retry_")
+    assert retry_event.response.instructions == "say hi"
+    created_events = [event for event in session._errors if event[0] == "generation_created"]
+    session._handle_response_created(
+        _response_created(client_event_id=retry_event.event_id, response_id="resp_2")
+    )
+    assert [
+        event for event in session._errors if event[0] == "generation_created"
+    ] == created_events
+    assert retry_event.event_id not in session._response_retry_event_ids
+    assert session._current_generation is not None
+    assert session._current_generation.retry_count == 1
+    error_event = next(event for event in session._errors if event[0] == "error")
+    assert error_event[1].recoverable is True
+
+
+async def test_response_done_failed_does_not_retry_after_output_started() -> None:
+    session = _create_response_retry_session()
+
+    fut = session.generate_reply(instructions="say hi")
+    create_event = session._sent_events[-1]
+    client_event_id = create_event.response.metadata["client_event_id"]
+    session._handle_response_created(_response_created(client_event_id=client_event_id))
+    await fut
+    assert session._current_generation is not None
+    session._current_generation.output_started = True
+
+    session._handle_response_done(_response_failed())
+    await asyncio.sleep(0.15)
+
+    assert len(session._sent_events) == 1
+    assert session._current_generation is None
+
+
+async def test_response_done_failed_does_not_retry_against_new_generation() -> None:
+    session = _create_response_retry_session(retry_interval=0.05)
+
+    fut = session.generate_reply(instructions="say hi")
+    create_event = session._sent_events[-1]
+    client_event_id = create_event.response.metadata["client_event_id"]
+    session._handle_response_created(_response_created(client_event_id=client_event_id))
+    await fut
+
+    session._handle_response_done(_response_failed())
+    old_generation = session._current_generation
+    session._handle_response_created(
+        _response_created(client_event_id="unrelated_turn", response_id="resp_new")
+    )
+    assert session._current_generation is not old_generation
+
+    await asyncio.sleep(0.5)
+
+    assert len(session._sent_events) == 1
+    assert session._response_retry_event_ids == set()
+    assert session._response_retry_generations == {}