From 66c6a5dc0ff7aa76bda157caba98997e42958564 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Fri, 31 May 2024 00:14:49 -0700 Subject: [PATCH] transports(websocket): base class from BaseOutputTransport --- examples/websocket-server/server.py | 7 +- .../transports/network/websocket_server.py | 92 ++++++------------- 2 files changed, 35 insertions(+), 64 deletions(-) diff --git a/examples/websocket-server/server.py b/examples/websocket-server/server.py index fa9fdc057..cf6ff0289 100644 --- a/examples/websocket-server/server.py +++ b/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) diff --git a/src/pipecat/transports/network/websocket_server.py b/src/pipecat/transports/network/websocket_server.py index 432b88db3..8a1cda629 100644 --- a/src/pipecat/transports/network/websocket_server.py +++ b/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):