Update SonioxSTTService to use WebsocketSTTService

This commit is contained in:
Mark Backman
2025-12-16 09:46:35 -05:00
parent eeb3b85e39
commit 24b5efb9d8

View File

@@ -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}")