From e31e87aabdb9badd6feb6b92bf7a5052a1e0462e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= <aleix@daily.co>
Date: Wed, 29 May 2024 23:48:55 -0700
Subject: [PATCH] transport(websocket): update audio_frame_size

---
 src/pipecat/services/ai_services.py           |  4 +--
 .../transports/network/websocket_server.py    | 28 +++++++++++++++++--
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/src/pipecat/services/ai_services.py b/src/pipecat/services/ai_services.py
index a7f74ccc3..fc879d887 100644
--- a/src/pipecat/services/ai_services.py
+++ b/src/pipecat/services/ai_services.py
@@ -196,7 +196,7 @@ class ImageGenService(AIService):
         super().__init__()
 
     # Renders the image. Returns an Image object.
-    @ abstractmethod
+    @abstractmethod
     async def run_image_gen(self, prompt: str) -> AsyncGenerator[Frame, None]:
         pass
 
@@ -215,7 +215,7 @@ class VisionService(AIService):
         super().__init__()
         self._describe_text = None
 
-    @ abstractmethod
+    @abstractmethod
     async def run_vision(self, frame: VisionImageRawFrame) -> AsyncGenerator[Frame, None]:
         pass
 
diff --git a/src/pipecat/transports/network/websocket_server.py b/src/pipecat/transports/network/websocket_server.py
index 24560ea2a..f765af615 100644
--- a/src/pipecat/transports/network/websocket_server.py
+++ b/src/pipecat/transports/network/websocket_server.py
@@ -6,6 +6,8 @@
 
 
 import asyncio
+import io
+import wave
 import websockets
 
 from typing import Awaitable, Callable
@@ -28,7 +30,8 @@ from loguru import logger
 
 
 class WebsocketServerParams(TransportParams):
-    audio_frame_size: int = 16000
+    add_wav_header: bool = False
+    audio_frame_size: int = 6400  # 200ms
     serializer: FrameSerializer = ProtobufFrameSerializer()
 
 
@@ -94,6 +97,8 @@ class WebsocketServerInputTransport(FrameProcessor):
             frame = self._params.serializer.deserialize(message)
             await self._internal_push_frame(frame)
 
+        logger.info(f"Client {websocket.remote_address} disconnected")
+
     async def _start(self):
         loop = self.get_event_loop()
         self._server_task = loop.create_task(self._server_task_handler())
@@ -159,13 +164,30 @@ class WebsocketServerOutputTransport(FrameProcessor):
         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)
 
+    async def _stop(self):
+        self._send_queue_task.cancel()
+
     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._stop()
             await self.push_frame(frame, direction)
         elif isinstance(frame, TTSStartedFrame):
             self._in_tts_audio = True