diff --git a/changelog/3236.added.md b/changelog/3236.added.md new file mode 100644 index 000000000..ace865739 --- /dev/null +++ b/changelog/3236.added.md @@ -0,0 +1 @@ +- Add Gladia session id to logs for `GladiaSTTService`. diff --git a/changelog/3236.changed.md b/changelog/3236.changed.md new file mode 100644 index 000000000..7c5015ead --- /dev/null +++ b/changelog/3236.changed.md @@ -0,0 +1,6 @@ +- Updated websocket STT services to use the `WebsocketSTTService` base class. This base class manages the websocket connection and handles reconnects. Updated services: + + - `AssemblyAISTTService` + - `AWSTranscribeSTTService` + - `GladiaSTTService` + - `SonioxSTTService` diff --git a/changelog/3236.fixed.md b/changelog/3236.fixed.md new file mode 100644 index 000000000..7814e80f8 --- /dev/null +++ b/changelog/3236.fixed.md @@ -0,0 +1 @@ +- In `GladiaSTTService`, reset the `_bytes_sent` counter on connecting the websocket. This avoids unnecessary audio buffer trimming. diff --git a/src/pipecat/services/assemblyai/stt.py b/src/pipecat/services/assemblyai/stt.py index 24251319f..4467fd134 100644 --- a/src/pipecat/services/assemblyai/stt.py +++ b/src/pipecat/services/assemblyai/stt.py @@ -21,7 +21,6 @@ from pipecat import __version__ as pipecat_version from pipecat.frames.frames import ( CancelFrame, EndFrame, - ErrorFrame, Frame, InterimTranscriptionFrame, StartFrame, @@ -30,7 +29,7 @@ from pipecat.frames.frames import ( UserStoppedSpeakingFrame, ) from pipecat.processors.frame_processor import FrameDirection -from pipecat.services.stt_service import STTService +from pipecat.services.stt_service import WebsocketSTTService from pipecat.transcriptions.language import Language from pipecat.utils.time import time_now_iso8601 from pipecat.utils.tracing.service_decorators import traced_stt @@ -44,15 +43,15 @@ from .models import ( ) try: - import websockets from websockets.asyncio.client import connect as websocket_connect + from websockets.protocol import State except ModuleNotFoundError as e: logger.error(f"Exception: {e}") logger.error('In order to use AssemblyAI, you need to `pip install "pipecat-ai[assemblyai]"`.') raise Exception(f"Missing module: {e}") -class AssemblyAISTTService(STTService): +class AssemblyAISTTService(WebsocketSTTService): """AssemblyAI real-time speech-to-text service. Provides real-time speech transcription using AssemblyAI's WebSocket API. @@ -80,15 +79,14 @@ class AssemblyAISTTService(STTService): vad_force_turn_endpoint: Whether to force turn endpoint on VAD stop. Defaults to True. **kwargs: Additional arguments passed to parent STTService class. """ + super().__init__(sample_rate=connection_params.sample_rate, **kwargs) + self._api_key = api_key self._language = language self._api_endpoint_base_url = api_endpoint_base_url self._connection_params = connection_params self._vad_force_turn_endpoint = vad_force_turn_endpoint - super().__init__(sample_rate=self._connection_params.sample_rate, **kwargs) - - self._websocket = None self._termination_event = asyncio.Event() self._received_termination = False self._connected = False @@ -114,7 +112,7 @@ class AssemblyAISTTService(STTService): frame: Start frame to begin processing. """ await super().start(frame) - self._chunk_size_bytes = int(self._chunk_size_ms * self._sample_rate * 2 / 1000) + self._chunk_size_bytes = int(self._chunk_size_ms * self.sample_rate * 2 / 1000) await self._connect() async def stop(self, frame: EndFrame): @@ -146,10 +144,11 @@ class AssemblyAISTTService(STTService): """ self._audio_buffer.extend(audio) - while len(self._audio_buffer) >= self._chunk_size_bytes: - chunk = bytes(self._audio_buffer[: self._chunk_size_bytes]) - self._audio_buffer = self._audio_buffer[self._chunk_size_bytes :] - await self._websocket.send(chunk) + if self._websocket and self._websocket.state is State.OPEN: + while len(self._audio_buffer) >= self._chunk_size_bytes: + chunk = bytes(self._audio_buffer[: self._chunk_size_bytes]) + self._audio_buffer = self._audio_buffer[self._chunk_size_bytes :] + await self._websocket.send(chunk) yield None @@ -164,7 +163,11 @@ class AssemblyAISTTService(STTService): if isinstance(frame, UserStartedSpeakingFrame): await self.start_ttfb_metrics() elif isinstance(frame, UserStoppedSpeakingFrame): - if self._vad_force_turn_endpoint: + if ( + self._vad_force_turn_endpoint + and self._websocket + and self._websocket.state is State.OPEN + ): await self._websocket.send(json.dumps({"type": "ForceEndpoint"})) await self.start_processing_metrics() @@ -191,7 +194,63 @@ class AssemblyAISTTService(STTService): return self._api_endpoint_base_url async def _connect(self): + """Connect to the AssemblyAI service. + + Establishes websocket connection and starts receive task. + """ + await self._connect_websocket() + + if self._websocket and not self._receive_task: + self._receive_task = self.create_task(self._receive_task_handler(self._report_error)) + + async def _disconnect(self): + """Disconnect from the AssemblyAI service. + + Sends termination message, waits for acknowledgment, and cleans up. + """ + if not self._connected or not self._websocket: + return + try: + self._termination_event.clear() + self._received_termination = False + + if self._websocket.state is State.OPEN: + # Send any remaining audio + if len(self._audio_buffer) > 0: + await self._websocket.send(bytes(self._audio_buffer)) + self._audio_buffer.clear() + + # Send termination message and wait for acknowledgment + try: + await self._websocket.send(json.dumps({"type": "Terminate"})) + + try: + await asyncio.wait_for(self._termination_event.wait(), timeout=5.0) + except asyncio.TimeoutError: + logger.warning("Timed out waiting for termination message from server") + + except Exception as e: + await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) + + except Exception as e: + await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) + finally: + # Clean up tasks and connection + if self._receive_task: + await self.cancel_task(self._receive_task) + self._receive_task = None + + await self._disconnect_websocket() + + async def _connect_websocket(self): + """Establish the websocket connection to AssemblyAI.""" + try: + if self._websocket and self._websocket.state is State.OPEN: + return + + logger.debug("Connecting to AssemblyAI WebSocket") + ws_url = self._build_ws_url() headers = { "Authorization": self._api_key, @@ -202,68 +261,50 @@ class AssemblyAISTTService(STTService): additional_headers=headers, ) self._connected = True - self._receive_task = self.create_task(self._receive_task_handler()) - await self._call_event_handler("on_connected") + logger.debug(f"{self} Connected to AssemblyAI WebSocket") except Exception as e: self._connected = False - await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) + await self.push_error(error_msg=f"Unable to connect to AssemblyAI: {e}", exception=e) raise - async def _disconnect(self): - """Disconnect from AssemblyAI WebSocket and wait for termination message.""" - if not self._connected or not self._websocket: - return - + async def _disconnect_websocket(self): + """Close the websocket connection to AssemblyAI.""" try: - self._termination_event.clear() - self._received_termination = False - - if len(self._audio_buffer) > 0: - await self._websocket.send(bytes(self._audio_buffer)) - self._audio_buffer.clear() - - try: - await self._websocket.send(json.dumps({"type": "Terminate"})) - - try: - await asyncio.wait_for(self._termination_event.wait(), timeout=5.0) - except asyncio.TimeoutError: - logger.warning("Timed out waiting for termination message from server") - - except Exception as e: - await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) - - if self._receive_task: - await self.cancel_task(self._receive_task) - - await self._websocket.close() - + if self._websocket: + logger.debug("Disconnecting from AssemblyAI WebSocket") + await self._websocket.close() except Exception as e: - await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) - + await self.push_error(error_msg=f"Error closing websocket: {e}", exception=e) finally: self._websocket = None self._connected = False - self._receive_task = None await self._call_event_handler("on_disconnected") - async def _receive_task_handler(self): - """Handle incoming WebSocket messages.""" - try: - while self._connected: - try: - message = await self._websocket.recv() - data = json.loads(message) - await self._handle_message(data) - except websockets.exceptions.ConnectionClosedOK: - break - except Exception as e: - await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) - break + def _get_websocket(self): + """Get the current WebSocket connection. - except Exception as e: - await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) + Returns: + The WebSocket connection. + + Raises: + Exception: If WebSocket is not connected. + """ + if self._websocket: + return self._websocket + raise Exception("Websocket not connected") + + async def _receive_messages(self): + """Receive and process websocket messages. + + Continuously processes messages from the websocket connection. + """ + async for message in self._get_websocket(): + try: + data = json.loads(message) + await self._handle_message(data) + except json.JSONDecodeError: + logger.warning(f"Received non-JSON message: {message}") def _parse_message(self, message: Dict[str, Any]) -> BaseMessage: """Parse a raw message into the appropriate message type.""" diff --git a/src/pipecat/services/aws/stt.py b/src/pipecat/services/aws/stt.py index 0878db315..0edbfd389 100644 --- a/src/pipecat/services/aws/stt.py +++ b/src/pipecat/services/aws/stt.py @@ -29,13 +29,12 @@ from pipecat.frames.frames import ( TranscriptionFrame, ) from pipecat.services.aws.utils import build_event_message, decode_event, get_presigned_url -from pipecat.services.stt_service import STTService +from pipecat.services.stt_service import WebsocketSTTService from pipecat.transcriptions.language import Language, resolve_language from pipecat.utils.time import time_now_iso8601 from pipecat.utils.tracing.service_decorators import traced_stt try: - import websockets from websockets.asyncio.client import connect as websocket_connect from websockets.protocol import State except ModuleNotFoundError as e: @@ -44,7 +43,7 @@ except ModuleNotFoundError as e: raise Exception(f"Missing module: {e}") -class AWSTranscribeSTTService(STTService): +class AWSTranscribeSTTService(WebsocketSTTService): """AWS Transcribe Speech-to-Text service using WebSocket streaming. Provides real-time speech transcription using AWS Transcribe's streaming API. @@ -99,9 +98,6 @@ class AWSTranscribeSTTService(STTService): "region": region or os.getenv("AWS_REGION", "us-east-1"), } - self._ws_client = None - self._connection_lock = asyncio.Lock() - self._connecting = False self._receive_task = None def get_service_encoding(self, encoding: str) -> str: @@ -123,29 +119,9 @@ class AWSTranscribeSTTService(STTService): Args: frame: Start frame signaling service initialization. - - Raises: - RuntimeError: If WebSocket connection cannot be established after retries. """ await super().start(frame) - logger.info("Starting AWS Transcribe service...") - retry_count = 0 - max_retries = 3 - - while retry_count < max_retries: - try: - await self._connect() - if self._ws_client and self._ws_client.state is State.OPEN: - logger.info("Successfully established WebSocket connection") - return - logger.warning("WebSocket connection not established after connect") - except Exception as e: - await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) - retry_count += 1 - if retry_count < max_retries: - await asyncio.sleep(1) # Wait before retrying - - raise RuntimeError("Failed to establish WebSocket connection after multiple attempts") + await self._connect() async def stop(self, frame: EndFrame): """Stop the service and disconnect from AWS Transcribe. @@ -174,140 +150,127 @@ class AWSTranscribeSTTService(STTService): Yields: ErrorFrame: If processing fails or connection issues occur. """ - try: - # Ensure WebSocket is connected - if not self._ws_client or self._ws_client.state is State.CLOSED: - logger.debug("WebSocket not connected, attempting to reconnect...") - try: - await self._connect() - except Exception as e: - yield ErrorFrame(error=f"Unknown error occurred: {e}") - return - - # Format the audio data according to AWS event stream format - event_message = build_event_message(audio) - - # Send the formatted event message + if self._websocket and self._websocket.state is State.OPEN: try: - await self._ws_client.send(event_message) + # Format the audio data according to AWS event stream format + event_message = build_event_message(audio) + + # Send the formatted event message + await self._websocket.send(event_message) # Start metrics after first chunk sent await self.start_processing_metrics() await self.start_ttfb_metrics() - except websockets.exceptions.ConnectionClosed as e: - logger.warning(f"Connection closed while sending: {e}") - await self._disconnect() - # Don't yield error here - we'll retry on next frame except Exception as e: - yield ErrorFrame(error=f"Unknown error occurred: {e}") - await self._disconnect() + yield ErrorFrame(error=f"Error sending audio: {e}") - except Exception as e: - yield ErrorFrame(error=f"Unknown error occurred: {e}") - await self._disconnect() + yield None async def _connect(self): - """Connect to AWS Transcribe with connection state management.""" - if self._ws_client and self._ws_client.state is State.OPEN and self._receive_task: - logger.debug(f"{self} Already connected") - return + """Connect to the AWS Transcribe service. - async with self._connection_lock: - if self._connecting: - logger.debug(f"{self} Connection already in progress") - return + Establishes websocket connection and starts receive task. + """ + await self._connect_websocket() - try: - self._connecting = True - logger.debug(f"{self} Starting connection process...") - - if self._ws_client: - await self._disconnect() - - language_code = self.language_to_service_language( - Language(self._settings["language"]) - ) - if not language_code: - raise ValueError(f"Unsupported language: {self._settings['language']}") - - # Generate random websocket key - websocket_key = "".join( - random.choices( - string.ascii_uppercase + string.ascii_lowercase + string.digits, k=20 - ) - ) - - # Add required headers - additional_headers = { - "Origin": "https://localhost", - "Sec-WebSocket-Key": websocket_key, - "Sec-WebSocket-Version": "13", - "Connection": "keep-alive", - } - - # Get presigned URL - presigned_url = get_presigned_url( - region=self._credentials["region"], - credentials={ - "access_key": self._credentials["aws_access_key_id"], - "secret_key": self._credentials["aws_secret_access_key"], - "session_token": self._credentials["aws_session_token"], - }, - language_code=language_code, - media_encoding=self.get_service_encoding( - self._settings["media_encoding"] - ), # Convert to AWS format - sample_rate=self._settings["sample_rate"], - number_of_channels=self._settings["number_of_channels"], - enable_partial_results_stabilization=True, - partial_results_stability="high", - show_speaker_label=self._settings["show_speaker_label"], - enable_channel_identification=self._settings["enable_channel_identification"], - ) - - logger.debug(f"{self} Connecting to WebSocket with URL: {presigned_url[:100]}...") - - # Connect with the required headers and settings - self._ws_client = await websocket_connect( - presigned_url, - additional_headers=additional_headers, - subprotocols=["mqtt"], - ping_interval=None, - ping_timeout=None, - compression=None, - ) - - logger.debug(f"{self} WebSocket connected, starting receive task...") - - # Start receive task - self._receive_task = self.create_task(self._receive_loop()) - - logger.info(f"{self} Successfully connected to AWS Transcribe") - - await self._call_event_handler("on_connected") - except Exception as e: - await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) - await self._disconnect() - raise - - finally: - self._connecting = False + if self._websocket and not self._receive_task: + self._receive_task = self.create_task(self._receive_task_handler(self._report_error)) async def _disconnect(self): - """Disconnect from AWS Transcribe.""" + """Disconnect from the AWS Transcribe service. + + Sends end-stream message and cleans up. + """ if self._receive_task: await self.cancel_task(self._receive_task) self._receive_task = None - try: - if self._ws_client and self._ws_client.state is State.OPEN: - # Send end-stream message + # Send end-stream message before closing + if self._websocket and self._websocket.state is State.OPEN: + try: end_stream = {"message-type": "event", "event": "end"} - await self._ws_client.send(json.dumps(end_stream)) - await self._ws_client.close() + await self._websocket.send(json.dumps(end_stream)) + except Exception as e: + await self.push_error(error_msg=f"Error sending end-stream: {e}", exception=e) + + await self._disconnect_websocket() + + async def _connect_websocket(self): + """Establish the websocket connection to AWS Transcribe.""" + try: + if self._websocket and self._websocket.state is State.OPEN: + return + + logger.debug("Connecting to AWS Transcribe WebSocket") + + language_code = self.language_to_service_language(Language(self._settings["language"])) + if not language_code: + raise ValueError(f"Unsupported language: {self._settings['language']}") + + # Generate random websocket key + websocket_key = "".join( + random.choices( + string.ascii_uppercase + string.ascii_lowercase + string.digits, k=20 + ) + ) + + # Add required headers + additional_headers = { + "Origin": "https://localhost", + "Sec-WebSocket-Key": websocket_key, + "Sec-WebSocket-Version": "13", + "Connection": "keep-alive", + } + + # Get presigned URL + presigned_url = get_presigned_url( + region=self._credentials["region"], + credentials={ + "access_key": self._credentials["aws_access_key_id"], + "secret_key": self._credentials["aws_secret_access_key"], + "session_token": self._credentials["aws_session_token"], + }, + language_code=language_code, + media_encoding=self.get_service_encoding( + self._settings["media_encoding"] + ), # Convert to AWS format + sample_rate=self._settings["sample_rate"], + number_of_channels=self._settings["number_of_channels"], + enable_partial_results_stabilization=True, + partial_results_stability="high", + show_speaker_label=self._settings["show_speaker_label"], + enable_channel_identification=self._settings["enable_channel_identification"], + ) + + logger.debug(f"{self} Connecting to WebSocket with URL: {presigned_url[:100]}...") + + # Connect with the required headers and settings + self._websocket = await websocket_connect( + presigned_url, + additional_headers=additional_headers, + subprotocols=["mqtt"], + ping_interval=None, + ping_timeout=None, + compression=None, + ) + + await self._call_event_handler("on_connected") + logger.info(f"{self} Successfully connected to AWS Transcribe") except Exception as e: - await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) + await self.push_error( + error_msg=f"Unable to connect to AWS Transcribe: {e}", exception=e + ) + raise + + async def _disconnect_websocket(self): + """Close the websocket connection to AWS Transcribe.""" + try: + if self._websocket: + logger.debug("Disconnecting from AWS Transcribe WebSocket") + await self._websocket.close() + except Exception as e: + await self.push_error(error_msg=f"Error closing websocket: {e}", exception=e) finally: - self._ws_client = None + self._websocket = None await self._call_event_handler("on_disconnected") def language_to_service_language(self, language: Language) -> str | None: @@ -471,16 +434,26 @@ class AWSTranscribeSTTService(STTService): ): pass - async def _receive_loop(self): - """Background task to receive and process messages from AWS Transcribe.""" - while True: - if not self._ws_client or self._ws_client.state is State.CLOSED: - logger.warning(f"{self} WebSocket closed in receive loop") - break + def _get_websocket(self): + """Get the current WebSocket connection. + Returns: + The WebSocket connection. + + Raises: + Exception: If WebSocket is not connected. + """ + if self._websocket: + return self._websocket + raise Exception("Websocket not connected") + + async def _receive_messages(self): + """Receive and process websocket messages. + + Continuously processes messages from the websocket connection. + """ + async for response in self._get_websocket(): try: - response = await self._ws_client.recv() - headers, payload = decode_event(response) if headers.get(":message-type") == "event": @@ -527,11 +500,5 @@ class AWSTranscribeSTTService(STTService): else: logger.debug(f"{self} Other message type received: {headers}") logger.debug(f"{self} Payload: {payload}") - except websockets.exceptions.ConnectionClosed as e: - await self.push_error( - error_msg=f"WebSocket connection closed in receive loop", exception=e - ) - break except Exception as e: - await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) - break + logger.warning(f"Error processing message: {e}") diff --git a/src/pipecat/services/gladia/stt.py b/src/pipecat/services/gladia/stt.py index 27b037f36..e73a4d77a 100644 --- a/src/pipecat/services/gladia/stt.py +++ b/src/pipecat/services/gladia/stt.py @@ -23,7 +23,6 @@ from pipecat import __version__ as pipecat_version from pipecat.frames.frames import ( CancelFrame, EndFrame, - ErrorFrame, Frame, InterimTranscriptionFrame, StartFrame, @@ -31,7 +30,7 @@ from pipecat.frames.frames import ( TranslationFrame, ) from pipecat.services.gladia.config import GladiaInputParams -from pipecat.services.stt_service import STTService +from pipecat.services.stt_service import WebsocketSTTService from pipecat.transcriptions.language import Language, resolve_language from pipecat.utils.time import time_now_iso8601 from pipecat.utils.tracing.service_decorators import traced_stt @@ -176,7 +175,7 @@ class _InputParamsDescriptor: return GladiaInputParams -class GladiaSTTService(STTService): +class GladiaSTTService(WebsocketSTTService): """Speech-to-Text service using Gladia's API. This service connects to Gladia's WebSocket API for real-time transcription @@ -202,8 +201,6 @@ class GladiaSTTService(STTService): sample_rate: Optional[int] = None, model: str = "solaria-1", params: Optional[GladiaInputParams] = None, - max_reconnection_attempts: int = 5, - reconnection_delay: float = 1.0, max_buffer_size: int = 1024 * 1024 * 20, # 20MB default buffer **kwargs, ): @@ -222,8 +219,6 @@ class GladiaSTTService(STTService): sample_rate: Audio sample rate in Hz. If None, uses service default. model: Model to use for transcription. Defaults to "solaria-1". params: Additional configuration parameters for Gladia service. - max_reconnection_attempts: Maximum number of reconnection attempts. Defaults to 5. - reconnection_delay: Initial delay between reconnection attempts in seconds. max_buffer_size: Maximum size of audio buffer in bytes. Defaults to 20MB. **kwargs: Additional arguments passed to the STTService parent class. """ @@ -256,16 +251,13 @@ class GladiaSTTService(STTService): self._url = url self.set_model_name(model) self._params = params - self._websocket = None self._receive_task = None self._keepalive_task = None self._settings = {} - # Reconnection settings - self._max_reconnection_attempts = max_reconnection_attempts - self._reconnection_delay = reconnection_delay - self._reconnection_attempts = 0 + # Session management self._session_url = None + self._session_id = None self._connection_active = False # Audio buffer management @@ -274,9 +266,8 @@ class GladiaSTTService(STTService): self._max_buffer_size = max_buffer_size self._buffer_lock = asyncio.Lock() - # Connection management - self._connection_task = None - self._should_reconnect = True + def __str__(self): + return f"{self.name} [{self._session_id}]" def can_generate_metrics(self) -> bool: """Check if the service can generate performance metrics. @@ -355,11 +346,7 @@ class GladiaSTTService(STTService): frame: The start frame triggering service startup. """ await super().start(frame) - if self._connection_task: - return - - self._should_reconnect = True - self._connection_task = self.create_task(self._connection_handler()) + await self._connect() async def stop(self, frame: EndFrame): """Stop the Gladia STT websocket connection. @@ -368,14 +355,8 @@ class GladiaSTTService(STTService): frame: The end frame triggering service shutdown. """ await super().stop(frame) - self._should_reconnect = False await self._send_stop_recording() - - if self._connection_task: - await self.cancel_task(self._connection_task) - self._connection_task = None - - await self._cleanup_connection() + await self._disconnect() async def cancel(self, frame: CancelFrame): """Cancel the Gladia STT websocket connection. @@ -384,13 +365,7 @@ class GladiaSTTService(STTService): frame: The cancel frame triggering service cancellation. """ await super().cancel(frame) - self._should_reconnect = False - - if self._connection_task: - await self.cancel_task(self._connection_task) - self._connection_task = None - - await self._cleanup_connection() + await self._disconnect() async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]: """Run speech-to-text on audio data. @@ -412,88 +387,93 @@ class GladiaSTTService(STTService): trim_size = len(self._audio_buffer) - self._max_buffer_size self._audio_buffer = self._audio_buffer[trim_size:] self._bytes_sent = max(0, self._bytes_sent - trim_size) - logger.warning(f"Audio buffer exceeded max size, trimmed {trim_size} bytes") + logger.warning(f"{self} Audio buffer exceeded max size, trimmed {trim_size} bytes") # Send audio if connected if self._connection_active and self._websocket and self._websocket.state is State.OPEN: try: await self._send_audio(audio) except websockets.exceptions.ConnectionClosed as e: - logger.warning(f"Websocket closed while sending audio chunk: {e}") + logger.warning(f"{self} Websocket closed while sending audio chunk: {e}") self._connection_active = False yield None - async def _connection_handler(self): - """Handle WebSocket connection with automatic reconnection.""" - while self._should_reconnect: - try: - # Initialize session if needed - if not self._session_url: - settings = self._prepare_settings() - response = await self._setup_gladia(settings) - self._session_url = response["url"] - self._reconnection_attempts = 0 - logger.info(f"Session URL : {self._session_url}") + async def _connect(self): + """Connect to the Gladia service. - # Connect with automatic reconnection - async with websocket_connect(self._session_url) as websocket: - try: - self._websocket = websocket - self._connection_active = True - logger.debug(f"{self} Connected to Gladia WebSocket") + Initializes the session if needed and establishes websocket connection. + """ + # Initialize session if needed + if not self._session_url: + settings = self._prepare_settings() + response = await self._setup_gladia(settings) + self._session_url = response["url"] + self._session_id = response["id"] + logger.info(f"{self} Session URL: {self._session_url}") - # Send buffered audio if any - await self._send_buffered_audio() + await self._connect_websocket() - # Start tasks - self._receive_task = self.create_task(self._receive_task_handler()) - self._keepalive_task = self.create_task(self._keepalive_task_handler()) + if self._websocket and not self._receive_task: + self._receive_task = self.create_task(self._receive_task_handler(self._report_error)) - # Wait for tasks to complete - await asyncio.gather(self._receive_task, self._keepalive_task) + if self._websocket and not self._keepalive_task: + self._keepalive_task = self.create_task(self._keepalive_task_handler()) - except websockets.exceptions.ConnectionClosed as e: - logger.warning(f"WebSocket connection closed: {e}") - self._connection_active = False + async def _disconnect(self): + """Disconnect from the Gladia service. - # Clean up tasks - if self._receive_task: - await self.cancel_task(self._receive_task) - if self._keepalive_task: - await self.cancel_task(self._keepalive_task) - - # Attempt reconnect using helper - if not await self._maybe_reconnect(): - break - - except Exception as e: - await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) - self._connection_active = False - - if not self._should_reconnect: - break - - # Reset session URL to get a new one - self._session_url = None - await asyncio.sleep(self._reconnection_delay) - - async def _cleanup_connection(self): - """Clean up connection resources.""" + Cleans up tasks and closes websocket connection. + """ self._connection_active = False if self._keepalive_task: await self.cancel_task(self._keepalive_task) self._keepalive_task = None - if self._websocket: - await self._websocket.close() - self._websocket = None - if self._receive_task: await self.cancel_task(self._receive_task) self._receive_task = None + await self._disconnect_websocket() + + async def _connect_websocket(self): + """Establish the websocket connection to Gladia.""" + try: + if self._websocket and self._websocket.state is State.OPEN: + return + + logger.debug(f"{self}Connecting to Gladia WebSocket") + + self._websocket = await websocket_connect(self._session_url) + self._connection_active = True + + # Reset byte tracking for new connection + async with self._buffer_lock: + self._bytes_sent = 0 + + await self._call_event_handler("on_connected") + + # Send buffered audio if any + await self._send_buffered_audio() + + logger.debug(f"{self} Connected to Gladia WebSocket") + except Exception as e: + await self.push_error(error_msg=f"Unable to connect to Gladia: {e}", exception=e) + raise + + async def _disconnect_websocket(self): + """Close the websocket connection to Gladia.""" + try: + if self._websocket and self._websocket.state is State.OPEN: + logger.debug(f"{self} Disconnecting from Gladia WebSocket") + await self._websocket.close() + except Exception as e: + await self.push_error(error_msg=f"Error closing websocket: {e}", exception=e) + finally: + self._websocket = None + await self._call_event_handler("on_disconnected") + async def _setup_gladia(self, settings: Dict[str, Any]): async with aiohttp.ClientSession() as session: params = {} @@ -510,10 +490,10 @@ class GladiaSTTService(STTService): else: error_text = await response.text() logger.error( - f"Gladia error: {response.status}: {error_text or response.reason}" + f"{self} Gladia error: {response.status}: {error_text or response.reason}" ) raise Exception( - f"Failed to initialize Gladia session: {response.status} - {error_text}" + f"{self} Failed to initialize Gladia session: {response.status} - {error_text}" ) @traced_stt @@ -541,28 +521,26 @@ class GladiaSTTService(STTService): if self._websocket and self._websocket.state is State.OPEN: await self._websocket.send(json.dumps({"type": "stop_recording"})) - async def _keepalive_task_handler(self): - """Send periodic empty audio chunks to keep the connection alive.""" - try: - KEEPALIVE_SLEEP = 20 - while self._connection_active: - # Send keepalive (Gladia times out after 30 seconds) - await asyncio.sleep(KEEPALIVE_SLEEP) - if self._websocket and self._websocket.state is State.OPEN: - # Send an empty audio chunk as keepalive - empty_audio = b"" - await self._send_audio(empty_audio) - else: - logger.debug("Websocket closed, stopping keepalive") - break - except websockets.exceptions.ConnectionClosed: - logger.debug("Connection closed during keepalive") - except Exception as e: - await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) + def _get_websocket(self): + """Get the current WebSocket connection. - async def _receive_task_handler(self): - try: - async for message in self._websocket: + Returns: + The WebSocket connection. + + Raises: + Exception: If WebSocket is not connected. + """ + if self._websocket: + return self._websocket + raise Exception("Websocket not connected") + + async def _receive_messages(self): + """Receive and process websocket messages. + + Continuously processes messages from the websocket connection. + """ + async for message in self._get_websocket(): + try: content = json.loads(message) # Handle audio chunk acknowledgments @@ -617,26 +595,24 @@ class GladiaSTTService(STTService): translation, "", time_now_iso8601(), translated_language ) ) + except json.JSONDecodeError: + logger.warning(f"{self} Received non-JSON message: {message}") + + async def _keepalive_task_handler(self): + """Send periodic empty audio chunks to keep the connection alive.""" + try: + KEEPALIVE_SLEEP = 20 + while self._connection_active: + # Send keepalive (Gladia times out after 30 seconds) + await asyncio.sleep(KEEPALIVE_SLEEP) + if self._websocket and self._websocket.state is State.OPEN: + # Send an empty audio chunk as keepalive + empty_audio = b"" + await self._send_audio(empty_audio) + else: + logger.debug(f"{self} Websocket closed, stopping keepalive") + break except websockets.exceptions.ConnectionClosed: - # Expected when closing the connection - pass + logger.debug(f"{self} Connection closed during keepalive") except Exception as e: await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) - - async def _maybe_reconnect(self) -> bool: - """Handle exponential backoff reconnection logic.""" - if not self._should_reconnect: - return False - self._reconnection_attempts += 1 - if self._reconnection_attempts > self._max_reconnection_attempts: - await self.push_error( - error_msg=f"Max reconnection attempts ({self._max_reconnection_attempts}) reached", - ) - self._should_reconnect = False - return False - delay = self._reconnection_delay * (2 ** (self._reconnection_attempts - 1)) - logger.debug( - f"{self} Reconnecting in {delay} seconds (attempt {self._reconnection_attempts}/{self._max_reconnection_attempts})" - ) - await asyncio.sleep(delay) - return True diff --git a/src/pipecat/services/soniox/stt.py b/src/pipecat/services/soniox/stt.py index 4afcaa7bb..5826691fb 100644 --- a/src/pipecat/services/soniox/stt.py +++ b/src/pipecat/services/soniox/stt.py @@ -17,7 +17,6 @@ from pydantic import BaseModel from pipecat.frames.frames import ( CancelFrame, EndFrame, - ErrorFrame, Frame, InterimTranscriptionFrame, StartFrame, @@ -25,13 +24,12 @@ from pipecat.frames.frames import ( UserStoppedSpeakingFrame, ) from pipecat.processors.frame_processor import FrameDirection -from pipecat.services.stt_service import STTService +from pipecat.services.stt_service import WebsocketSTTService from pipecat.transcriptions.language import Language from pipecat.utils.time import time_now_iso8601 from pipecat.utils.tracing.service_decorators import traced_stt try: - import websockets from websockets.asyncio.client import connect as websocket_connect from websockets.protocol import State except ModuleNotFoundError as e: @@ -134,7 +132,7 @@ def _prepare_language_hints( return list(set(prepared_languages)) -class SonioxSTTService(STTService): +class SonioxSTTService(WebsocketSTTService): """Speech-to-Text service using Soniox's WebSocket API. This service connects to Soniox's WebSocket API for real-time transcription @@ -173,7 +171,6 @@ class SonioxSTTService(STTService): self.set_model_name(params.model) self._params = params self._vad_force_turn_endpoint = vad_force_turn_endpoint - self._websocket = None self._final_transcription_buffer = [] self._last_tokens_received: Optional[float] = None @@ -188,59 +185,7 @@ class SonioxSTTService(STTService): frame: The start frame containing initialization parameters. """ await super().start(frame) - if self._websocket: - return - - self._websocket = await websocket_connect(self._url) - - if not self._websocket: - await self.push_error(error_msg=f"Unable to connect to Soniox API at {self._url}") - - # If vad_force_turn_endpoint is not enabled, we need to enable endpoint detection. - # Either one or the other is required. - enable_endpoint_detection = not self._vad_force_turn_endpoint - - context = self._params.context - if isinstance(context, SonioxContextObject): - context = context.model_dump() - - # Send the initial configuration message. - config = { - "api_key": self._api_key, - "model": self._model_name, - "audio_format": self._params.audio_format, - "num_channels": self._params.num_channels or 1, - "enable_endpoint_detection": enable_endpoint_detection, - "sample_rate": self.sample_rate, - "language_hints": _prepare_language_hints(self._params.language_hints), - "context": context, - "enable_speaker_diarization": self._params.enable_speaker_diarization, - "enable_language_identification": self._params.enable_language_identification, - "client_reference_id": self._params.client_reference_id, - } - - # Send the configuration message. - await self._websocket.send(json.dumps(config)) - - if self._websocket and not self._receive_task: - self._receive_task = self.create_task(self._receive_task_handler()) - if self._websocket and not self._keepalive_task: - self._keepalive_task = self.create_task(self._keepalive_task_handler()) - - async def _cleanup(self): - if self._keepalive_task: - await self.cancel_task(self._keepalive_task) - self._keepalive_task = None - - if self._websocket: - await self._websocket.close() - self._websocket = None - - if self._receive_task: - # Task cannot cancel itself. If task called _cleanup() we expect it to cancel itself. - if self._receive_task != asyncio.current_task(): - await self._receive_task - self._receive_task = None + await self._connect() async def stop(self, frame: EndFrame): """Stop the Soniox STT websocket connection. @@ -253,6 +198,7 @@ class SonioxSTTService(STTService): """ await super().stop(frame) await self._send_stop_recording() + await self._disconnect() async def cancel(self, frame: CancelFrame): """Cancel the Soniox STT websocket connection. @@ -265,7 +211,7 @@ class SonioxSTTService(STTService): frame: The cancel frame. """ await super().cancel(frame) - await self._cleanup() + await self._disconnect() async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]: """Send audio data to Soniox STT Service. @@ -311,28 +257,110 @@ class SonioxSTTService(STTService): # Send stop recording message await self._websocket.send("") - async def _keepalive_task_handler(self): - """Connection has to be open all the time.""" + async def _connect(self): + """Connect to the Soniox service. + + Establishes websocket connection and starts receive and keepalive tasks. + """ + await self._connect_websocket() + + if self._websocket and not self._receive_task: + self._receive_task = self.create_task(self._receive_task_handler(self._report_error)) + + if self._websocket and not self._keepalive_task: + self._keepalive_task = self.create_task(self._keepalive_task_handler()) + + async def _disconnect(self): + """Disconnect from the Soniox service. + + Cleans up tasks and closes websocket connection. + """ + if self._keepalive_task: + await self.cancel_task(self._keepalive_task) + self._keepalive_task = None + + if self._receive_task: + await self.cancel_task(self._receive_task) + self._receive_task = None + + await self._disconnect_websocket() + + async def _connect_websocket(self): + """Establish the websocket connection to Soniox.""" try: - while True: - logger.trace("Sending keepalive message") - if self._websocket and self._websocket.state is State.OPEN: - await self._websocket.send(KEEPALIVE_MESSAGE) - else: - logger.debug("WebSocket connection closed.") - break - await asyncio.sleep(5) + if self._websocket and self._websocket.state is State.OPEN: + return - except websockets.exceptions.ConnectionClosed: - # Expected when closing the connection - logger.debug("WebSocket connection closed, keepalive task stopped.") + logger.debug("Connecting to Soniox STT") + + self._websocket = await websocket_connect(self._url) + + if not self._websocket: + await self.push_error(error_msg=f"Unable to connect to Soniox API at {self._url}") + raise Exception(f"Unable to connect to Soniox API at {self._url}") + + # If vad_force_turn_endpoint is not enabled, we need to enable endpoint detection. + # Either one or the other is required. + enable_endpoint_detection = not self._vad_force_turn_endpoint + + context = self._params.context + if isinstance(context, SonioxContextObject): + context = context.model_dump() + + # Send the initial configuration message. + config = { + "api_key": self._api_key, + "model": self._model_name, + "audio_format": self._params.audio_format, + "num_channels": self._params.num_channels or 1, + "enable_endpoint_detection": enable_endpoint_detection, + "sample_rate": self.sample_rate, + "language_hints": _prepare_language_hints(self._params.language_hints), + "context": context, + "enable_speaker_diarization": self._params.enable_speaker_diarization, + "enable_language_identification": self._params.enable_language_identification, + "client_reference_id": self._params.client_reference_id, + } + + # Send the configuration message. + await self._websocket.send(json.dumps(config)) + + await self._call_event_handler("on_connected") + logger.debug("Connected to Soniox STT") except Exception as e: - await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) + await self.push_error(error_msg=f"Unable to connect to Soniox: {e}", exception=e) + raise - async def _receive_task_handler(self): - if not self._websocket: - return + async def _disconnect_websocket(self): + """Close the websocket connection to Soniox.""" + try: + if self._websocket: + logger.debug("Disconnecting from Soniox STT") + await self._websocket.close() + except Exception as e: + await self.push_error(error_msg=f"Error closing websocket: {e}", exception=e) + finally: + self._websocket = None + await self._call_event_handler("on_disconnected") + def _get_websocket(self): + """Get the current WebSocket connection. + + Returns: + The WebSocket connection. + + Raises: + Exception: If WebSocket is not connected. + """ + if self._websocket: + return self._websocket + raise Exception("Websocket not connected") + + async def _receive_messages(self): + """Receive and process websocket messages. + + Continuously processes messages from the websocket connection. + """ # Transcription frame will be only sent after we get the "endpoint" event. self._final_transcription_buffer = [] @@ -351,8 +379,8 @@ class SonioxSTTService(STTService): await self.stop_processing_metrics() self._final_transcription_buffer = [] - try: - async for message in self._websocket: + async for message in self._get_websocket(): + try: content = json.loads(message) tokens = content["tokens"] @@ -404,7 +432,7 @@ class SonioxSTTService(STTService): # In case of error, still send the final transcript (if any remaining in the buffer). await send_endpoint_transcript() await self.push_error( - error_msg=f"Error: {error_code} (_receive_task_handler) - {error_message}" + error_msg=f"Error: {error_code} (_receive_messages) - {error_message}" ) finished = content.get("finished") @@ -412,11 +440,24 @@ class SonioxSTTService(STTService): # When finished, still send the final transcript (if any remaining in the buffer). await send_endpoint_transcript() logger.debug("Transcription finished.") - await self._cleanup() return - except websockets.exceptions.ConnectionClosed: - # Expected when closing the connection. - pass + except json.JSONDecodeError: + logger.warning(f"Received non-JSON message: {message}") + except Exception as e: + logger.warning(f"Error processing message: {e}") + + async def _keepalive_task_handler(self): + """Connection has to be open all the time.""" + try: + while True: + logger.trace("Sending keepalive message") + if self._websocket and self._websocket.state is State.OPEN: + await self._websocket.send(KEEPALIVE_MESSAGE) + else: + logger.debug("WebSocket connection closed.") + break + await asyncio.sleep(5) + except Exception as e: - await self.push_error(error_msg=f"Error receiving message: {e}", exception=e) + logger.debug(f"Keepalive task stopped: {e}")