transports(websocket): base class from BaseOutputTransport

examples/websocket-server/server.py

@@ -31,7 +31,12 @@ logger.add(sys.stderr, level="DEBUG")
 
 async def main():
     async with aiohttp.ClientSession() as session:
-        transport = WebsocketServerTransport(params=WebsocketServerParams(add_wav_header=True))
+        transport = WebsocketServerTransport(
+            params=WebsocketServerParams(
+                audio_out_enabled=True,
+                add_wav_header=True
+            )
+        )
 
         vad = SileroVAD(audio_passthrough=True)
 

src/pipecat/transports/network/websocket_server.py

@@ -24,6 +24,7 @@ from pipecat.frames.frames import (
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
 from pipecat.serializers.base_serializer import FrameSerializer
 from pipecat.serializers.protobuf import ProtobufFrameSerializer
+from pipecat.transports.base_output import BaseOutputTransport
 from pipecat.transports.base_transport import BaseTransport, TransportParams
 
 from loguru import logger
@@ -134,24 +135,16 @@ class WebsocketServerInputTransport(FrameProcessor):
                 break
 
 
-class WebsocketServerOutputTransport(FrameProcessor):
+class WebsocketServerOutputTransport(BaseOutputTransport):
 
     def __init__(self, params: WebsocketServerParams):
-        super().__init__()
+        super().__init__(params)
 
         self._params = params
 
-        self._websocket = None
-        self._audio_buffer = bytes()
-
         self._websocket: websockets.WebSocketServerProtocol | None = None
 
-        loop = self.get_event_loop()
-        self._send_queue_task = loop.create_task(self._send_queue_task_handler())
-        self._send_queue = asyncio.Queue()
-
         self._audio_buffer = bytes()
-        self._in_tts_audio = False
 
     async def set_client_connection(self, websocket: websockets.WebSocketServerProtocol):
         if self._websocket:
@@ -159,61 +152,34 @@ class WebsocketServerOutputTransport(FrameProcessor):
             logger.warning("Only one client allowed, using new connection")
         self._websocket = websocket
 
-    async def _send_queue_task_handler(self):
-        running = True
-        while running:
-            frame = await self._send_queue.get()
-            if self._websocket and frame:
-                # We send WAV data so we can easily decoded in the browser.
-                if self._params.add_wav_header:
-                    content = io.BytesIO()
-                    ww = wave.open(content, "wb")
-                    ww.setsampwidth(2)
-                    ww.setnchannels(frame.num_channels)
-                    ww.setframerate(frame.sample_rate)
-                    ww.writeframes(frame.audio)
-                    ww.close()
-                    content.seek(0)
-                    wav_frame = AudioRawFrame(
-                        content.read(),
-                        sample_rate=frame.sample_rate,
-                        num_channels=frame.num_channels)
-                    frame = wav_frame
-                proto = self._params.serializer.serialize(frame)
-                await self._websocket.send(proto)
+    def write_raw_audio_frames(self, frames: bytes):
+        self._audio_buffer += frames
+        while len(self._audio_buffer) >= self._params.audio_frame_size:
+            frame = AudioRawFrame(
+                audio=self._audio_buffer[:self._params.audio_frame_size],
+                sample_rate=self._params.audio_out_sample_rate,
+                num_channels=self._params.audio_out_channels
+            )
 
-    async def _stop(self):
-        self._send_queue_task.cancel()
+            if self._params.add_wav_header:
+                content = io.BytesIO()
+                ww = wave.open(content, "wb")
+                ww.setsampwidth(2)
+                ww.setnchannels(frame.num_channels)
+                ww.setframerate(frame.sample_rate)
+                ww.writeframes(frame.audio)
+                ww.close()
+                content.seek(0)
+                wav_frame = AudioRawFrame(
+                    content.read(),
+                    sample_rate=frame.sample_rate,
+                    num_channels=frame.num_channels)
+                frame = wav_frame
 
-    async def process_frame(self, frame: Frame, direction: FrameDirection):
-        if isinstance(frame, CancelFrame):
-            await self._stop()
-            await self.push_frame(frame, direction)
-        elif isinstance(frame, TTSStartedFrame):
-            self._in_tts_audio = True
-        elif isinstance(frame, AudioRawFrame):
-            if self._in_tts_audio:
-                self._audio_buffer += frame.audio
-                while len(self._audio_buffer) >= self._params.audio_frame_size:
-                    frame = AudioRawFrame(
-                        audio=self._audio_buffer[:self._params.audio_frame_size],
-                        sample_rate=self._params.audio_out_sample_rate,
-                        num_channels=self._params.audio_out_channels
-                    )
-                    await self._send_queue.put(frame)
-                    self._audio_buffer = self._audio_buffer[self._params.audio_frame_size:]
-        elif isinstance(frame, TTSStoppedFrame):
-            self._in_tts_audio = False
-            if self._audio_buffer:
-                frame = AudioRawFrame(
-                    audio=self._audio_buffer,
-                    sample_rate=self._params.audio_out_sample_rate,
-                    num_channels=self._params.audio_out_channels
-                )
-                await self._send_queue.put(frame)
-                self._audio_buffer = bytes()
-        else:
-            await self.push_frame(frame, direction)
+            proto = self._params.serializer.serialize(frame)
+            asyncio.run_coroutine_threadsafe(self._websocket.send(proto), self.get_event_loop())
+
+            self._audio_buffer = self._audio_buffer[self._params.audio_frame_size:]
 
 
 class WebsocketServerTransport(BaseTransport):