From f39e17857e7b49d3828845091c1cfd1bc531ead0 Mon Sep 17 00:00:00 2001 From: Mark Backman Date: Thu, 9 Jan 2025 21:20:45 -0500 Subject: [PATCH] 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 | 33 +----- src/pipecat/services/websocket_service.py | 125 ++++++++++++++++++++++ 2 files changed, 130 insertions(+), 28 deletions(-) create mode 100644 src/pipecat/services/websocket_service.py diff --git a/src/pipecat/services/cartesia.py b/src/pipecat/services/cartesia.py index 4e0bb111c..e821041db 100644 --- a/src/pipecat/services/cartesia.py +++ b/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) diff --git a/src/pipecat/services/websocket_service.py b/src/pipecat/services/websocket_service.py new file mode 100644 index 000000000..2ceeb2a8f --- /dev/null +++ b/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