From 9b8800ac1d80fa9be378756b2b83bd4f9095e353 Mon Sep 17 00:00:00 2001 From: Martin Schweiger Date: Fri, 23 May 2025 17:32:31 +0800 Subject: [PATCH] update stt.py --- src/pipecat/services/assemblyai/stt.py | 283 ++++++++++++++++--------- 1 file changed, 178 insertions(+), 105 deletions(-) diff --git a/src/pipecat/services/assemblyai/stt.py b/src/pipecat/services/assemblyai/stt.py index 2153d969e..3dfd9a3f3 100644 --- a/src/pipecat/services/assemblyai/stt.py +++ b/src/pipecat/services/assemblyai/stt.py @@ -5,30 +5,40 @@ # import asyncio -from typing import AsyncGenerator, Optional +import json +from typing import Any, AsyncGenerator, Dict +from urllib.parse import urlencode from loguru import logger from pipecat.frames.frames import ( CancelFrame, EndFrame, - ErrorFrame, Frame, InterimTranscriptionFrame, StartFrame, TranscriptionFrame, + UserStartedSpeakingFrame, + UserStoppedSpeakingFrame, ) +from pipecat.processors.frame_processor import FrameDirection from pipecat.services.stt_service import STTService from pipecat.transcriptions.language import Language from pipecat.utils.time import time_now_iso8601 -from pipecat.utils.tracing.service_decorators import traced_stt + +from .models import ( + AssemblyAIConnectionParams, + BaseMessage, + BeginMessage, + TerminationMessage, + TurnMessage, +) try: - import assemblyai as aai - from assemblyai import AudioEncoding + import websockets except ModuleNotFoundError as e: logger.error(f"Exception: {e}") - logger.error("In order to use AssemblyAI, you need to `pip install pipecat-ai[assemblyai]`.") + logger.error("In order to use AssemblyAI, you need to `pip install websockets`.") raise Exception(f"Missing module: {e}") @@ -37,28 +47,28 @@ class AssemblyAISTTService(STTService): self, *, api_key: str, - sample_rate: Optional[int] = None, - encoding: AudioEncoding = AudioEncoding("pcm_s16le"), - language=Language.EN, # Only English is supported for Realtime + language: Language = Language.EN, # AssemblyAI only supports English + api_endpoint_base_url: str = "wss://streaming.assemblyai.com/v3/ws", + connection_params: AssemblyAIConnectionParams = AssemblyAIConnectionParams(), **kwargs, ): - super().__init__(sample_rate=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 - aai.settings.api_key = api_key - self._transcriber: Optional[aai.RealtimeTranscriber] = None + super().__init__(sample_rate=self._connection_params.sample_rate, **kwargs) - self._settings = { - "encoding": encoding, - "language": language, - } + self._websocket = None + self._termination_event = asyncio.Event() + self._received_termination = False + self._connected = False + + self._receive_task = None def can_generate_metrics(self) -> bool: return True - async def set_language(self, language: Language): - logger.info(f"Switching STT language to: [{language}]") - self._settings["language"] = language - async def start(self, frame: StartFrame): await super().start(frame) await self._connect() @@ -72,97 +82,160 @@ class AssemblyAISTTService(STTService): await self._disconnect() async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]: - """Process an audio chunk for STT transcription. + await self._websocket.send(audio) + yield Frame() - This method streams the audio data to AssemblyAI for real-time transcription. - Transcription results are handled asynchronously via callback functions. - - :param audio: Audio data as bytes - :yield: None (transcription frames are pushed via self.push_frame in callbacks) - """ - if self._transcriber: + async def process_frame(self, frame: Frame, direction: FrameDirection): + await super().process_frame(frame, direction) + if isinstance(frame, UserStartedSpeakingFrame): await self.start_ttfb_metrics() + elif isinstance(frame, UserStoppedSpeakingFrame): + # TODO: if the user opts to use VAD, we should send a ForceEndpoint message await self.start_processing_metrics() - self._transcriber.stream(audio) - yield None - @traced_stt - async def _handle_transcription( - self, transcript: str, is_final: bool, language: Optional[Language] = None - ): - """Handle a transcription result with tracing.""" - await self.stop_ttfb_metrics() - await self.stop_processing_metrics() + def _build_ws_url(self) -> str: + """Build WebSocket URL with query parameters using urllib.parse.urlencode.""" + params = { + k: str(v).lower() if isinstance(v, bool) else v + for k, v in self._connection_params.dict().items() + if v is not None + } + if params: + query_string = urlencode(params) + return f"{self._api_endpoint_base_url}?{query_string}" + return self._api_endpoint_base_url async def _connect(self): - """Establish a connection to the AssemblyAI real-time transcription service. - - This method sets up the necessary callback functions and initializes the - AssemblyAI transcriber. - """ - if self._transcriber: - return - - def on_open(session_opened: aai.RealtimeSessionOpened): - """Callback for when the connection to AssemblyAI is opened.""" - logger.info(f"{self}: Connected to AssemblyAI") - - def on_data(transcript: aai.RealtimeTranscript): - """Callback for handling incoming transcription data. - - This function runs in a separate thread from the main asyncio event loop. - It creates appropriate transcription frames and schedules them to be - pushed to the next stage of the pipeline in the main event loop. - """ - if not transcript.text: - return - - timestamp = time_now_iso8601() - is_final = isinstance(transcript, aai.RealtimeFinalTranscript) - language = self._settings["language"] - - if is_final: - frame = TranscriptionFrame(transcript.text, "", timestamp, language) - else: - frame = InterimTranscriptionFrame(transcript.text, "", timestamp, language) - - asyncio.run_coroutine_threadsafe( - self._handle_transcription(transcript.text, is_final, language), - self.get_event_loop(), + try: + ws_url = self._build_ws_url() + self._websocket = await websockets.connect( + ws_url, + extra_headers={"Authorization": self._api_key}, ) - - # Schedule the coroutine to run in the main event loop - # This is necessary because this callback runs in a different thread - asyncio.run_coroutine_threadsafe(self.push_frame(frame), self.get_event_loop()) - - def on_error(error: aai.RealtimeError): - """Callback for handling errors from AssemblyAI. - - Like on_data, this runs in a separate thread and schedules error - handling in the main event loop. - """ - logger.error(f"{self}: An error occurred: {error}") - # Schedule the coroutine to run in the main event loop - asyncio.run_coroutine_threadsafe( - self.push_frame(ErrorFrame(str(error))), self.get_event_loop() - ) - - def on_close(): - """Callback for when the connection to AssemblyAI is closed.""" - logger.info(f"{self}: Disconnected from AssemblyAI") - - self._transcriber = aai.RealtimeTranscriber( - sample_rate=self.sample_rate, - encoding=self._settings["encoding"], - on_data=on_data, - on_error=on_error, - on_open=on_open, - on_close=on_close, - ) - self._transcriber.connect() + self._connected = True + self._receive_task = asyncio.create_task(self._receive_task_handler()) + except Exception as e: + logger.error(f"Failed to connect to AssemblyAI: {e}") + self._connected = False + raise async def _disconnect(self): - """Disconnect from the AssemblyAI service and clean up resources.""" - if self._transcriber: - self._transcriber.close() - self._transcriber = None + """Disconnect from AssemblyAI WebSocket and wait for termination message.""" + if not self._connected or not self._websocket: + return + + try: + self._termination_event.clear() + self._received_termination = False + + 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: + logger.warning(f"Error during termination handshake: {e}") + + finally: + self._connected = False + + try: + if not self._websocket.closed: + await self._websocket.close() + except Exception as e: + logger.warning(f"Error closing WebSocket: {e}") + + if self._receive_task and not self._receive_task.done(): + self._receive_task.cancel() + try: + await self._receive_task + except (asyncio.CancelledError, Exception) as e: + logger.debug(f"Receive task cancelled: {e}") + + 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: + logger.error(f"Error processing WebSocket message: {e}") + break + + except Exception as e: + logger.error(f"Fatal error in receive handler: {e}") + + def _parse_message(self, message: Dict[str, Any]) -> BaseMessage: + """Parse a raw message into the appropriate message type.""" + msg_type = message.get("type") + + if msg_type == "Begin": + return BeginMessage.parse_obj(message) + elif msg_type == "Turn": + return TurnMessage.parse_obj(message) + elif msg_type == "Termination": + return TerminationMessage.parse_obj(message) + else: + raise ValueError(f"Unknown message type: {msg_type}") + + async def _handle_message(self, message: Dict[str, Any]): + """Handle AssemblyAI WebSocket messages.""" + try: + parsed_message = self._parse_message(message) + + if isinstance(parsed_message, BeginMessage): + logger.debug( + f"Session Begin: {parsed_message.id} (expires at {parsed_message.expires_at})" + ) + elif isinstance(parsed_message, TurnMessage): + await self._handle_transcription(parsed_message) + elif isinstance(parsed_message, TerminationMessage): + await self._handle_termination(parsed_message) + except Exception as e: + logger.error(f"Error handling message: {e}") + + async def _handle_termination(self, message: TerminationMessage): + """Handle termination message.""" + self._received_termination = True + self._termination_event.set() + + logger.info( + f"Session Terminated: Audio Duration={message.audio_duration_seconds}s, " + f"Session Duration={message.session_duration_seconds}s" + ) + await self.push_frame(EndFrame()) + + async def _handle_transcription(self, message: TurnMessage): + """Handle transcription results.""" + if not message.transcript: + return + await self.stop_ttfb_metrics() + if message.end_of_turn: + await self.push_frame( + TranscriptionFrame( + message.transcript, + "", # participant + time_now_iso8601(), + self._language, + ) + ) + await self.stop_processing_metrics() + else: + await self.push_frame( + InterimTranscriptionFrame( + message.transcript, + "", # participant + time_now_iso8601(), + self._language, + ) + )