Merge branch 'main' into vl_feature_websocket_fastapi_timeout

src/pipecat/transports/network/fastapi_websocket.py

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2024, Daily
+# Copyright (c) 2025, Daily
 #
 # SPDX-License-Identifier: BSD 2-Clause License
 #
@@ -8,28 +8,26 @@
 import asyncio
 import io
 import time
+import typing
 import wave
-
 from typing import Awaitable, Callable
-from pydantic.main import BaseModel
+
+from loguru import logger
+from pydantic import BaseModel
 
 from pipecat.frames.frames import (
-    AudioRawFrame,
-    CancelFrame,
-    EndFrame,
     Frame,
     InputAudioRawFrame,
+    OutputAudioRawFrame,
     StartFrame,
     StartInterruptionFrame,
 )
 from pipecat.processors.frame_processor import FrameDirection
-from pipecat.serializers.base_serializer import FrameSerializer
+from pipecat.serializers.base_serializer import FrameSerializer, FrameSerializerType
 from pipecat.transports.base_input import BaseInputTransport
 from pipecat.transports.base_output import BaseOutputTransport
 from pipecat.transports.base_transport import BaseTransport, TransportParams
 
-from loguru import logger
-
 try:
     from fastapi import WebSocket
     from starlette.websockets import WebSocketState
@@ -73,21 +71,23 @@ class FastAPIWebsocketInputTransport(BaseInputTransport):
         await self._callbacks.on_client_connected(self._websocket)
         self._receive_task = self.get_event_loop().create_task(self._receive_messages())
 
+    def _iter_data(self) -> typing.AsyncIterator[bytes | str]:
+        if self._params.serializer.type == FrameSerializerType.BINARY:
+            return self._websocket.iter_bytes()
+        else:
+            return self._websocket.iter_text()
+
     async def _receive_messages(self):
-        async for message in self._websocket.iter_text():
+        async for message in self._iter_data():
             frame = self._params.serializer.deserialize(message)
 
             if not frame:
                 continue
 
-            if isinstance(frame, AudioRawFrame):
-                await self.push_audio_frame(
-                    InputAudioRawFrame(
-                        audio=frame.audio,
-                        sample_rate=frame.sample_rate,
-                        num_channels=frame.num_channels,
-                    )
-                )
+            if isinstance(frame, InputAudioRawFrame):
+                await self.push_audio_frame(frame)
+            else:
+                await self.push_frame(frame)
 
         await self._callbacks.on_client_disconnected(self._websocket)
 
@@ -120,30 +120,50 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
             self._next_send_time = 0
 
     async def write_raw_audio_frames(self, frames: bytes):
-        frame = AudioRawFrame(
+        if self._websocket.client_state != WebSocketState.CONNECTED:
+            # Simulate audio playback with a sleep.
+            await self._write_audio_sleep()
+            return
+
+        frame = OutputAudioRawFrame(
             audio=frames,
             sample_rate=self._params.audio_out_sample_rate,
             num_channels=self._params.audio_out_channels,
         )
 
         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
+            with io.BytesIO() as buffer:
+                with wave.open(buffer, "wb") as wf:
+                    wf.setsampwidth(2)
+                    wf.setnchannels(frame.num_channels)
+                    wf.setframerate(frame.sample_rate)
+                    wf.writeframes(frame.audio)
+                wav_frame = OutputAudioRawFrame(
+                    buffer.getvalue(),
+                    sample_rate=frame.sample_rate,
+                    num_channels=frame.num_channels,
+                )
+                frame = wav_frame
 
+        await self._write_frame(frame)
+
+        self._websocket_audio_buffer = bytes()
+
+        # Simulate audio playback with a sleep.
+        await self._write_audio_sleep()
+
+    async def _write_frame(self, frame: Frame):
         payload = self._params.serializer.serialize(frame)
         if payload and self._websocket.client_state == WebSocketState.CONNECTED:
-            await self._websocket.send_text(payload)
+            await self._send_data(payload)
 
+    def _send_data(self, data: str | bytes):
+        if self._params.serializer.type == FrameSerializerType.BINARY:
+            return self._websocket.send_bytes(data)
+        else:
+            return self._websocket.send_text(data)
+
+    async def _write_audio_sleep(self):
         # Simulate a clock.
         current_time = time.monotonic()
         sleep_duration = max(0, self._next_send_time - current_time)
@@ -153,13 +173,6 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
         else:
             self._next_send_time += self._send_interval
 
-        self._websocket_audio_buffer = bytes()
-
-    async def _write_frame(self, frame: Frame):
-        payload = self._params.serializer.serialize(frame)
-        if payload and self._websocket.client_state == WebSocketState.CONNECTED:
-            await self._websocket.send_text(payload)
-
 
 class FastAPIWebsocketTransport(BaseTransport):
     def __init__(