From 9fd061e33e87cff1fea3e6f164e783da9aca057b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Wed, 29 May 2024 10:36:46 -0700 Subject: [PATCH] transports: added new WebsocketServerTransport --- .../transports/network/websocket_server.py | 228 ++++++++++++++++++ 1 file changed, 228 insertions(+) create mode 100644 src/pipecat/transports/network/websocket_server.py diff --git a/src/pipecat/transports/network/websocket_server.py b/src/pipecat/transports/network/websocket_server.py new file mode 100644 index 000000000..81eae933b --- /dev/null +++ b/src/pipecat/transports/network/websocket_server.py @@ -0,0 +1,228 @@ +# +# Copyright (c) 2024, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + + +import asyncio +import websockets + +from typing import Awaitable, Callable +from pydantic.main import BaseModel + +from pipecat.frames.frames import ( + AudioRawFrame, + CancelFrame, + EndFrame, + Frame, + StartFrame, + TTSStartedFrame, + TTSStoppedFrame) +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_transport import BaseTransport, TransportParams + +from loguru import logger + + +class WebsocketServerParams(TransportParams): + audio_frame_size: int = 16000 + serializer: FrameSerializer = ProtobufFrameSerializer() + + +class WebsocketServerCallbacks(BaseModel): + on_connection: Callable[[websockets.WebSocketServerProtocol], Awaitable[None]] + + +class WebsocketServerInputTransport(FrameProcessor): + + def __init__( + self, + host: str, + port: int, + params: WebsocketServerParams, + callbacks: WebsocketServerCallbacks): + super().__init__() + + self._host = host + self._port = port + self._params = params + self._callbacks = callbacks + + self._websocket: websockets.WebSocketServerProtocol | None = None + + self._stop_server_event = asyncio.Event() + + # Create push frame task. This is the task that will push frames in + # order. We also guarantee that all frames are pushed in the same task. + self._create_push_task() + + async def process_frame(self, frame: Frame, direction: FrameDirection): + if isinstance(frame, CancelFrame): + await self._stop() + # We don't queue a CancelFrame since we want to stop ASAP. + await self.push_frame(frame, direction) + elif isinstance(frame, StartFrame): + await self._start() + await self._internal_push_frame(frame, direction) + elif isinstance(frame, EndFrame): + await self._stop() + await self._internal_push_frame(frame, direction) + else: + await self._internal_push_frame(frame, direction) + + async def _server_task_handler(self): + logger.info(f"Starting websocket server on {self._host}:{self._port}") + async with websockets.serve(self._client_handler, self._host, self._port) as server: + await self._stop_server_event.wait() + + async def _client_handler(self, websocket: websockets.WebSocketServerProtocol, path): + logger.info(f"New client connection from {websocket.remote_address}") + if self._websocket: + await self._websocket.close() + logger.warning("Only one client connected, using new connection") + + self._websocket = websocket + + # Notify + await self._callbacks.on_connection(websocket) + + # Handle incoming messages + async for message in websocket: + frame = self._params.serializer.deserialize(message) + await self._internal_push_frame(frame) + + async def _start(self): + loop = self.get_event_loop() + self._server_task = loop.create_task(self._server_task_handler()) + + async def _stop(self): + self._stop_server_event.set() + self._push_frame_task.cancel() + await self._server_task + + # + # Push frames task + # + + def _create_push_task(self): + loop = self.get_event_loop() + self._push_frame_task = loop.create_task(self._push_frame_task_handler()) + self._push_queue = asyncio.Queue() + + async def _internal_push_frame( + self, + frame: Frame | None, + direction: FrameDirection | None = FrameDirection.DOWNSTREAM): + await self._push_queue.put((frame, direction)) + + async def _push_frame_task_handler(self): + running = True + while running: + try: + (frame, direction) = await self._push_queue.get() + await self.push_frame(frame, direction) + running = not isinstance(frame, EndFrame) + except asyncio.CancelledError: + break + + +class WebsocketServerOutputTransport(FrameProcessor): + + def __init__(self, params: WebsocketServerParams): + super().__init__() + + 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: + await self._websocket.close() + 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: + proto = self._params.serializer.serialize(frame) + await self._websocket.send(proto) + + async def process_frame(self, frame: Frame, direction: FrameDirection): + if isinstance(frame, CancelFrame): + # await self.stop() + # We don't queue a CancelFrame since we want to stop ASAP. + 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) + + +class WebsocketServerTransport(BaseTransport): + + def __init__(self, host: str = "localhost", port: int = 8765, + params: WebsocketServerParams = WebsocketServerParams()): + super().__init__() + self._host = host + self._port = port + self._params = params + + self._callbacks = WebsocketServerCallbacks( + on_connection=self._on_connection + ) + self._input: WebsocketServerInputTransport | None = None + self._output: WebsocketServerOutputTransport | None = None + self._websocket: websockets.WebSocketServerProtocol | None = None + + def input(self) -> FrameProcessor: + if not self._input: + self._input = WebsocketServerInputTransport( + self._host, self._port, self._params, self._callbacks) + return self._input + + def output(self) -> FrameProcessor: + if not self._output: + self._output = WebsocketServerOutputTransport(self._params) + return self._output + + async def _on_connection(self, websocket): + if self._output: + await self._output.set_client_connection(websocket) + else: + logger.error("A WebsocketServerTransport output is missing in the pipeline")