From 24b5efb9d8374fedc118ec2baadf29c6f01d09f1 Mon Sep 17 00:00:00 2001 From: Mark Backman Date: Tue, 16 Dec 2025 09:46:35 -0500 Subject: [PATCH] Update SonioxSTTService to use WebsocketSTTService --- src/pipecat/services/soniox/stt.py | 209 +++++++++++++++++------------ 1 file changed, 125 insertions(+), 84 deletions(-) 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}")