update stt.py

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,
+                )
+            )