Add (relatively spartan) reconnection logic to GeminiMultimodalLiveLLMService, leveraging the Gemini Live session resumption mechanism

src/pipecat/services/gemini_multimodal_live/gemini.py

@@ -14,6 +14,7 @@ voice transcription, streaming responses, and tool usage.
 import base64
 import io
 import json
+import random
 import time
 from dataclasses import dataclass
 from enum import Enum
@@ -98,6 +99,7 @@ try:
         Part,
         ProactivityConfig,
         RealtimeInputConfig,
+        SessionResumptionConfig,
         SlidingWindow,
         SpeechConfig,
         StartSensitivity,
@@ -618,6 +620,12 @@ class GeminiMultimodalLiveLLMService(LLMService):
         )
         self._vad_params = params.vad
 
+        # Reconnection tracking
+        self._consecutive_failures = 0
+        self._max_consecutive_failures = 3
+        self._connection_established_threshold = 10.0  # seconds
+        self._connection_start_time = None
+
         self._settings = {
             "frequency_penalty": params.frequency_penalty,
             "max_tokens": params.max_tokens,
@@ -645,6 +653,9 @@ class GeminiMultimodalLiveLLMService(LLMService):
         self._search_result_buffer = ""
         self._accumulated_grounding_metadata = None
 
+        # Session resumption
+        self._session_resumption_handle: Optional[str] = None
+
     def _create_client(self, api_key: str, http_options: Optional[HttpOptions] = None):
         self._client = Client(api_key=api_key, http_options=http_options)
 
@@ -850,14 +861,19 @@ class GeminiMultimodalLiveLLMService(LLMService):
         else:
             await self.push_frame(frame, direction)
 
-    async def _connect(self):
+    async def _connect(self, session_resumption_handle: Optional[str] = None):
         """Establish client connection to Gemini Live API."""
         if self._session:
             # Here we assume that if we have a client, we are connected. We
             # handle disconnections in the send/recv code paths.
             return
 
-        logger.info("Connecting to Gemini service")
+        if session_resumption_handle:
+            logger.info(
+                f"Connecting to Gemini service with session_resumption_handle: {session_resumption_handle}"
+            )
+        else:
+            logger.info("Connecting to Gemini service")
         try:
             # Assemble basic configuration
             config = LiveConnectConfig(
@@ -879,6 +895,7 @@ class GeminiMultimodalLiveLLMService(LLMService):
                 ),
                 input_audio_transcription=AudioTranscriptionConfig(),
                 output_audio_transcription=AudioTranscriptionConfig(),
+                session_resumption=SessionResumptionConfig(handle=session_resumption_handle),
             )
 
             # Add context window compression to configuration, if enabled
@@ -965,12 +982,18 @@ class GeminiMultimodalLiveLLMService(LLMService):
         async with self._client.aio.live.connect(model=self._model_name, config=config) as session:
             logger.info("Connected to Gemini service")
 
+            # Mark connection start time
+            self._connection_start_time = time.time()
+
             await self._handle_session_ready(session)
 
             while True:
                 try:
                     turn = self._session.receive()
                     async for message in turn:
+                        # Reset failure counter if connection has been stable
+                        self._check_and_reset_failure_counter()
+
                         if message.server_content and message.server_content.model_turn:
                             await self._handle_msg_model_turn(message)
                         elif (
@@ -988,13 +1011,67 @@ class GeminiMultimodalLiveLLMService(LLMService):
                             await self._handle_msg_grounding_metadata(message)
                         elif message.tool_call:
                             await self._handle_msg_tool_call(message)
+                        elif message.session_resumption_update:
+                            self._handle_msg_resumption_update(message)
                 except Exception as e:
                     if not self._disconnecting:
-                        await self.push_error(
+                        should_reconnect = await self._handle_connection_error(e)
-                            ErrorFrame(error=f"{self} Error in receive loop: {e}", fatal=True)
+                        if should_reconnect:
-                        )
+                            await self._reconnect()
+                            return  # Exit this connection handler, _reconnect will start a new one
                     break
 
+    def _check_and_reset_failure_counter(self):
+        """Check if connection has been stable long enough to reset the failure counter.
+
+        If the connection has been active for longer than the established threshold
+        and there are accumulated failures, reset the counter to 0.
+        """
+        if (
+            self._connection_start_time
+            and self._consecutive_failures > 0
+            and time.time() - self._connection_start_time >= self._connection_established_threshold
+        ):
+            logger.info(
+                f"Connection stable for {self._connection_established_threshold}s, "
+                f"resetting failure counter from {self._consecutive_failures} to 0"
+            )
+            self._consecutive_failures = 0
+
+    async def _handle_connection_error(self, error: Exception) -> bool:
+        """Handle a connection error and determine if reconnection should be attempted.
+
+        Args:
+            error: The exception that caused the connection error.
+
+        Returns:
+            True if reconnection should be attempted, False if a fatal error should be pushed.
+        """
+        self._consecutive_failures += 1
+        logger.warning(
+            f"Connection error (failure {self._consecutive_failures}/{self._max_consecutive_failures}): {error}"
+        )
+
+        if self._consecutive_failures >= self._max_consecutive_failures:
+            logger.error(
+                f"Max consecutive failures ({self._max_consecutive_failures}) reached, "
+                "treating as fatal error"
+            )
+            await self.push_error(
+                ErrorFrame(error=f"{self} Error in receive loop: {error}", fatal=True)
+            )
+            return False
+        else:
+            logger.info(
+                f"Attempting reconnection ({self._consecutive_failures}/{self._max_consecutive_failures})"
+            )
+            return True
+
+    async def _reconnect(self):
+        """Reconnect to Gemini Live API."""
+        await self._disconnect()
+        await self._connect(session_resumption_handle=self._session_resumption_handle)
+
     async def _disconnect(self):
         """Disconnect from Gemini Live API and clean up resources."""
         logger.info("Disconnecting from Gemini service")
@@ -1013,10 +1090,7 @@ class GeminiMultimodalLiveLLMService(LLMService):
 
     async def _send_user_audio(self, frame):
         """Send user audio frame to Gemini Live API."""
-        if self._audio_input_paused:
+        if self._audio_input_paused or self._disconnecting or not self._session:
-            return
-
-        if not self._session:
             return
 
         # Send all audio to Gemini
@@ -1051,7 +1125,7 @@ class GeminiMultimodalLiveLLMService(LLMService):
         Args:
             text: The text to send as user input.
         """
-        if not self._session:
+        if self._disconnecting or not self._session:
             return
 
         try:
@@ -1061,10 +1135,7 @@ class GeminiMultimodalLiveLLMService(LLMService):
 
     async def _send_user_video(self, frame):
         """Send user video frame to Gemini Live API."""
-        if self._video_input_paused:
+        if self._video_input_paused or self._disconnecting or not self._session:
-            return
-
-        if not self._session:
             return
 
         now = time.time()
@@ -1085,6 +1156,9 @@ class GeminiMultimodalLiveLLMService(LLMService):
 
     async def _create_initial_response(self):
         """Create initial response based on context history."""
+        if self._disconnecting:
+            return
+
         if not self._session:
             self._run_llm_when_session_ready = True
             return
@@ -1112,6 +1186,9 @@ class GeminiMultimodalLiveLLMService(LLMService):
 
     async def _create_single_response(self, messages_list):
         """Create a single response from a list of messages."""
+        if self._disconnecting or not self._session:
+            return
+
         # Create a throwaway context just for the purpose of getting messages
         # in the right format
         context = GeminiMultimodalLiveContext.upgrade(OpenAILLMContext(messages=messages_list))
@@ -1132,6 +1209,9 @@ class GeminiMultimodalLiveLLMService(LLMService):
     @traced_gemini_live(operation="llm_tool_result")
     async def _tool_result(self, tool_result_message):
         """Send tool result back to the API."""
+        if self._disconnecting or not self._session:
+            return
+
         # For now we're shoving the name into the tool_call_id field, so this
         # will work until we revisit that.
         id = tool_result_message.get("tool_call_id")
@@ -1407,6 +1487,11 @@ class GeminiMultimodalLiveLLMService(LLMService):
 
         await self.start_llm_usage_metrics(tokens)
 
+    def _handle_msg_resumption_update(self, message: LiveServerMessage):
+        update = message.session_resumption_update
+        if update.resumable and update.new_handle:
+            self._session_resumption_handle = update.new_handle
+
     async def _handle_send_error(self, error: Exception):
         # In server-to-server contexts, a WebSocket error should be quite rare.
         # Given how hard it is to recover from a send-side error with proper