From 8255770b6ce966a68dd08357633b155abefac559 Mon Sep 17 00:00:00 2001 From: Mark Backman Date: Tue, 16 Dec 2025 09:30:03 -0500 Subject: [PATCH] Update AssemblyAISTTService to use WebsocketSTTService --- src/pipecat/services/assemblyai/stt.py | 165 +++++++++++++++---------- 1 file changed, 103 insertions(+), 62 deletions(-) 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."""