Add a WebsocketService base class to retry, ensure that retries reset after a successful connection, update Cartesia to use the new WebsocketService

src/pipecat/services/cartesia.py

@@ -12,7 +12,6 @@ from typing import AsyncGenerator, List, Optional, Union
 
 from loguru import logger
 from pydantic import BaseModel
-from tenacity import AsyncRetrying, RetryCallState, stop_after_attempt, wait_exponential
 
 from pipecat.frames.frames import (
     BotStoppedSpeakingFrame,
@@ -30,6 +29,7 @@ from pipecat.frames.frames import (
 )
 from pipecat.processors.frame_processor import FrameDirection
 from pipecat.services.ai_services import TTSService, WordTTSService
+from pipecat.services.websocket_service import WebsocketService
 from pipecat.transcriptions.language import Language
 
 # See .env.example for Cartesia configuration needed
@@ -76,7 +76,7 @@ def language_to_cartesia_language(language: Language) -> str | None:
     return result
 
 
-class CartesiaTTSService(WordTTSService):
+class CartesiaTTSService(WordTTSService, WebsocketService):
     class InputParams(BaseModel):
         language: Optional[Language] = Language.EN
         speed: Optional[Union[str, float]] = ""
@@ -106,12 +106,14 @@ class CartesiaTTSService(WordTTSService):
         # if we're interrupted. Cartesia gives us word-by-word timestamps. We
         # can use those to generate text frames ourselves aligned with the
         # playout timing of the audio!
-        super().__init__(
+        WordTTSService.__init__(
+            self,
             aggregate_sentences=True,
             push_text_frames=False,
             sample_rate=sample_rate,
             **kwargs,
         )
+        WebsocketService.__init__(self)
 
         self._api_key = api_key
         self._cartesia_version = cartesia_version
@@ -131,7 +133,6 @@ class CartesiaTTSService(WordTTSService):
         self.set_model_name(model)
         self.set_voice(voice_id)
 
-        self._websocket = None
         self._context_id = None
         self._receive_task = None
 
@@ -275,30 +276,6 @@ class CartesiaTTSService(WordTTSService):
             else:
                 logger.error(f"{self} error, unknown message type: {msg}")
 
-    async def _reconnect_websocket(self, retry_state: RetryCallState):
-        logger.warning(f"{self} reconnecting (attempt: {retry_state.attempt_number})")
-        await self._disconnect_websocket()
-        await self._connect_websocket()
-
-    async def _receive_task_handler(self):
-        while True:
-            try:
-                async for attempt in AsyncRetrying(
-                    stop=stop_after_attempt(3),
-                    wait=wait_exponential(multiplier=1, min=4, max=10),
-                    before_sleep=self._reconnect_websocket,
-                    reraise=True,
-                ):
-                    with attempt:
-                        await self._receive_messages()
-            except asyncio.CancelledError:
-                break
-            except Exception as e:
-                message = f"{self} error receiving messages: {e}"
-                logger.error(message)
-                await self.push_error(ErrorFrame(message, fatal=True))
-                break
-
     async def process_frame(self, frame: Frame, direction: FrameDirection):
         await super().process_frame(frame, direction)
 

src/pipecat/services/websocket_service.py

@@ -0,0 +1,125 @@
+#
+# Copyright (c) 2025, Daily
+#
+# SPDX-License-Identifier: BSD 2-Clause License
+#
+
+import asyncio
+from abc import ABC, abstractmethod
+from typing import Optional
+
+import websockets
+from loguru import logger
+
+from pipecat.frames.frames import ErrorFrame
+
+
+class WebsocketService(ABC):
+    """Base class for websocket-based services with reconnection logic."""
+
+    def __init__(self):
+        """Initialize websocket attributes."""
+        self._websocket: Optional[websockets.WebSocketClientProtocol] = None
+
+    async def _verify_connection(self) -> bool:
+        """Verify websocket connection is working.
+
+        Returns:
+            bool: True if connection is verified working, False otherwise
+        """
+        try:
+            if not self._websocket:
+                return False
+            await self._websocket.ping()
+            return True
+        except Exception as e:
+            logger.error(f"{self} connection verification failed: {e}")
+            return False
+
+    async def _reconnect_websocket(self, attempt_number: int) -> bool:
+        """Reconnect the websocket.
+
+        Args:
+            attempt_number: Current retry attempt number
+
+        Returns:
+            bool: True if reconnection and verification successful, False otherwise
+        """
+        logger.warning(f"{self} reconnecting (attempt: {attempt_number})")
+        await self._disconnect_websocket()
+        await self._connect_websocket()
+        return await self._verify_connection()
+
+    def calculate_wait_time(
+        self, attempt: int, min_wait: float = 4, max_wait: float = 10, multiplier: float = 1
+    ) -> float:
+        """Calculate exponential backoff wait time.
+
+        Args:
+            attempt: Current attempt number (1-based)
+            min_wait: Minimum wait time in seconds
+            max_wait: Maximum wait time in seconds
+            multiplier: Base multiplier for exponential calculation
+
+        Returns:
+            Wait time in seconds
+        """
+        try:
+            exp = 2 ** (attempt - 1) * multiplier
+            result = max(0, min(exp, max_wait))
+            return max(min_wait, result)
+        except (ValueError, ArithmeticError):
+            return max_wait
+
+    async def _receive_task_handler(self):
+        """Handles WebSocket message receiving with automatic retry logic."""
+        retry_count = 0
+        MAX_RETRIES = 3
+
+        while True:
+            try:
+                await self._receive_messages()
+                logger.debug(f"{self} connection established successfully")
+                retry_count = 0  # Reset counter on successful message receive
+
+            except asyncio.CancelledError:
+                break
+
+            except Exception as e:
+                retry_count += 1
+                if retry_count >= MAX_RETRIES:
+                    message = f"{self} error receiving messages: {e}"
+                    logger.error(message)
+                    await self.push_error(ErrorFrame(message, fatal=True))
+                    break
+
+                logger.warning(f"{self} connection error, will retry: {e}")
+
+                try:
+                    if await self._reconnect_websocket(retry_count):
+                        retry_count = 0  # Reset counter on successful reconnection
+                    wait_time = self.calculate_wait_time(retry_count)
+                    await asyncio.sleep(wait_time)
+                except Exception as reconnect_error:
+                    logger.error(f"{self} reconnection failed: {reconnect_error}")
+                    continue
+
+    @abstractmethod
+    async def _connect_websocket(self):
+        """Implement service-specific websocket connection logic."""
+        pass
+
+    @abstractmethod
+    async def _disconnect_websocket(self):
+        """Implement service-specific websocket disconnection logic."""
+        pass
+
+    @abstractmethod
+    async def _receive_messages(self):
+        """Implement service-specific message receiving logic."""
+        pass
+
+    @abstractmethod
+    async def push_error(self, error: ErrorFrame):
+        """Implement service-specific error handling."""
+        pass