Update SonioxSTTService to use WebsocketSTTService

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:
+        await self._connect()
-            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
 
     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):
+    async def _connect(self):
-        """Connection has to be open all the time."""
+        """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:
+            if self._websocket and self._websocket.state is State.OPEN:
-                logger.trace("Sending keepalive message")
+                return
-                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 websockets.exceptions.ConnectionClosed:
+            logger.debug("Connecting to Soniox STT")
-            # Expected when closing the connection
+
-            logger.debug("WebSocket connection closed, keepalive task stopped.")
+            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):
+    async def _disconnect_websocket(self):
-        if not self._websocket:
+        """Close the websocket connection to Soniox."""
-            return
+        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._get_websocket():
-            async for message in self._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:
+            except json.JSONDecodeError:
-            # Expected when closing the connection.
+                logger.warning(f"Received non-JSON message: {message}")
-            pass
+            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}")