Merge pull request #791 from pipecat-ai/aleix/fastapi-generic-websocket

FastAPIWebsocketTransport: fix to work with text and binary

CHANGELOG.md

@@ -68,6 +68,9 @@ async def on_audio_data(processor, audio, sample_rate, num_channels):
 
 ### Fixed
 
+- Fixed `FastAPIWebsocketTransport` so it can work with binary data (e.g. using
+  the protobuf serializer).
+
 - Fixed an issue in `CartesiaTTSService` that could cause previous audio to be
   received after an interruption.
 

examples/websocket-server/index.html

@@ -49,13 +49,13 @@
       let startBtn = document.getElementById('startAudioBtn');
       let stopBtn = document.getElementById('stopAudioBtn');
 
-      const proto = protobuf.load("frames.proto", (err, root) => {
+      const proto = protobuf.load('frames.proto', (err, root) => {
           if (err) {
               throw err;
           }
-          Frame = root.lookupType("pipecat.Frame");
+          Frame = root.lookupType('pipecat.Frame');
-          const progressText = document.getElementById("progressText");
+          const progressText = document.getElementById('progressText');
-          progressText.textContent = "We are ready! Make sure to run the server and then click `Start Audio`.";
+          progressText.textContent = 'We are ready! Make sure to run the server and then click `Start Audio`.';
 
           startBtn.disabled = false;
           stopBtn.disabled = true;
@@ -63,18 +63,60 @@
 
       function initWebSocket() {
           ws = new WebSocket('ws://localhost:8765');
+          // This is so `event.data` is already an ArrayBuffer.
+          ws.binaryType = 'arraybuffer';
 
-          ws.addEventListener('open', () => console.log('WebSocket connection established.'));
+          ws.addEventListener('open', handleWebSocketOpen);
           ws.addEventListener('message', handleWebSocketMessage);
           ws.addEventListener('close', (event) => {
-              console.log("WebSocket connection closed.", event.code, event.reason);
+              console.log('WebSocket connection closed.', event.code, event.reason);
               stopAudio(false);
           });
           ws.addEventListener('error', (event) => console.error('WebSocket error:', event));
       }
 
-      async function handleWebSocketMessage(event) {
+      function handleWebSocketOpen(event) {
-          const arrayBuffer = await event.data.arrayBuffer();
+        console.log('WebSocket connection established.', event)
+
+        navigator.mediaDevices.getUserMedia({
+              audio: {
+                  sampleRate: SAMPLE_RATE,
+                  channelCount: NUM_CHANNELS,
+                  autoGainControl: true,
+                  echoCancellation: true,
+                  noiseSuppression: true,
+              }
+          }).then((stream) => {
+              microphoneStream = stream;
+              // 512 is closest thing to 200ms.
+              scriptProcessor = audioContext.createScriptProcessor(512, 1, 1);
+              source = audioContext.createMediaStreamSource(stream);
+              source.connect(scriptProcessor);
+              scriptProcessor.connect(audioContext.destination);
+
+              scriptProcessor.onaudioprocess = (event) => {
+                  if (!ws) {
+                      return;
+                  }
+
+                  const audioData = event.inputBuffer.getChannelData(0);
+                  const pcmS16Array = convertFloat32ToS16PCM(audioData);
+                  const pcmByteArray = new Uint8Array(pcmS16Array.buffer);
+                  const frame = Frame.create({
+                      audio: {
+                          audio: Array.from(pcmByteArray),
+                          sampleRate: SAMPLE_RATE,
+                          numChannels: NUM_CHANNELS
+                      }
+                  });
+                  const encodedFrame = new Uint8Array(Frame.encode(frame).finish());
+                  ws.send(encodedFrame);
+              };
+          }).catch((error) => console.error('Error accessing microphone:', error));
+      }
+
+      function handleWebSocketMessage(event) {
+          const arrayBuffer = event.data;
           if (isPlaying) {
               enqueueAudioFromProto(arrayBuffer);
           }
@@ -127,49 +169,13 @@
           stopBtn.disabled = false;
 
           audioContext = new (window.AudioContext || window.webkitAudioContext)({
-              latencyHint: "interactive",
+              latencyHint: 'interactive',
               sampleRate: SAMPLE_RATE
           });
 
           isPlaying = true;
 
           initWebSocket();
-
-          navigator.mediaDevices.getUserMedia({
-              audio: {
-                  sampleRate: SAMPLE_RATE,
-                  channelCount: NUM_CHANNELS,
-                  autoGainControl: true,
-                  echoCancellation: true,
-                  noiseSuppression: true,
-              }
-          }).then((stream) => {
-              microphoneStream = stream;
-              // 512 is closest thing to 200ms.
-              scriptProcessor = audioContext.createScriptProcessor(512, 1, 1);
-              source = audioContext.createMediaStreamSource(stream);
-              source.connect(scriptProcessor);
-              scriptProcessor.connect(audioContext.destination);
-
-              scriptProcessor.onaudioprocess = (event) => {
-                  if (!ws) {
-                      return;
-                  }
-
-                  const audioData = event.inputBuffer.getChannelData(0);
-                  const pcmS16Array = convertFloat32ToS16PCM(audioData);
-                  const pcmByteArray = new Uint8Array(pcmS16Array.buffer);
-                  const frame = Frame.create({
-                      audio: {
-                          audio: Array.from(pcmByteArray),
-                          sampleRate: SAMPLE_RATE,
-                          numChannels: NUM_CHANNELS
-                      }
-                  });
-                  const encodedFrame = new Uint8Array(Frame.encode(frame).finish());
-                  ws.send(encodedFrame);
-              };
-          }).catch((error) => console.error('Error accessing microphone:', error));
       }
 
       function stopAudio(closeWebsocket) {

src/pipecat/serializers/base_serializer.py

@@ -5,11 +5,22 @@
 #
 
 from abc import ABC, abstractmethod
+from enum import Enum
 
 from pipecat.frames.frames import Frame
 
 
+class FrameSerializerType(Enum):
+    BINARY = "binary"
+    TEXT = "text"
+
+
 class FrameSerializer(ABC):
+    @property
+    @abstractmethod
+    def type(self) -> FrameSerializerType:
+        pass
+
     @abstractmethod
     def serialize(self, frame: Frame) -> str | bytes | None:
         pass

src/pipecat/serializers/livekit.py

@@ -8,7 +8,7 @@ import ctypes
 import pickle
 
 from pipecat.frames.frames import Frame, InputAudioRawFrame, OutputAudioRawFrame
-from pipecat.serializers.base_serializer import FrameSerializer
+from pipecat.serializers.base_serializer import FrameSerializer, FrameSerializerType
 
 from loguru import logger
 
@@ -21,6 +21,10 @@ except ModuleNotFoundError as e:
 
 
 class LivekitFrameSerializer(FrameSerializer):
+    @property
+    def type(self) -> FrameSerializerType:
+        return FrameSerializerType.BINARY
+
     def serialize(self, frame: Frame) -> str | bytes | None:
         if not isinstance(frame, OutputAudioRawFrame):
             return None

src/pipecat/serializers/protobuf.py

@@ -8,8 +8,14 @@ import dataclasses
 
 import pipecat.frames.protobufs.frames_pb2 as frame_protos
 
-from pipecat.frames.frames import AudioRawFrame, Frame, TextFrame, TranscriptionFrame
+from pipecat.frames.frames import (
-from pipecat.serializers.base_serializer import FrameSerializer
+    Frame,
+    InputAudioRawFrame,
+    OutputAudioRawFrame,
+    TextFrame,
+    TranscriptionFrame,
+)
+from pipecat.serializers.base_serializer import FrameSerializer, FrameSerializerType
 
 from loguru import logger
 
@@ -17,15 +23,25 @@ from loguru import logger
 class ProtobufFrameSerializer(FrameSerializer):
     SERIALIZABLE_TYPES = {
         TextFrame: "text",
-        AudioRawFrame: "audio",
+        OutputAudioRawFrame: "audio",
         TranscriptionFrame: "transcription",
     }
-
     SERIALIZABLE_FIELDS = {v: k for k, v in SERIALIZABLE_TYPES.items()}
 
+    DESERIALIZABLE_TYPES = {
+        TextFrame: "text",
+        InputAudioRawFrame: "audio",
+        TranscriptionFrame: "transcription",
+    }
+    DESERIALIZABLE_FIELDS = {v: k for k, v in DESERIALIZABLE_TYPES.items()}
+
     def __init__(self):
         pass
 
+    @property
+    def type(self) -> FrameSerializerType:
+        return FrameSerializerType.BINARY
+
     def serialize(self, frame: Frame) -> str | bytes | None:
         proto_frame = frame_protos.Frame()
         if type(frame) not in self.SERIALIZABLE_TYPES:
@@ -40,8 +56,7 @@ class ProtobufFrameSerializer(FrameSerializer):
             if value and hasattr(proto_attr, field.name):
                 setattr(proto_attr, field.name, value)
 
-        result = proto_frame.SerializeToString()
+        return proto_frame.SerializeToString()
-        return result
 
     def deserialize(self, data: str | bytes) -> Frame | None:
         """Returns a Frame object from a Frame protobuf. Used to convert frames
@@ -64,11 +79,11 @@ class ProtobufFrameSerializer(FrameSerializer):
 
         proto = frame_protos.Frame.FromString(data)
         which = proto.WhichOneof("frame")
-        if which not in self.SERIALIZABLE_FIELDS:
+        if which not in self.DESERIALIZABLE_FIELDS:
             logger.error("Unable to deserialize a valid frame")
             return None
 
-        class_name = self.SERIALIZABLE_FIELDS[which]
+        class_name = self.DESERIALIZABLE_FIELDS[which]
         args = getattr(proto, which)
         args_dict = {}
         for field in proto.DESCRIPTOR.fields_by_name[which].message_type.fields:

src/pipecat/serializers/twilio.py

@@ -10,8 +10,8 @@ import json
 from pydantic import BaseModel
 
 from pipecat.audio.utils import ulaw_to_pcm, pcm_to_ulaw
-from pipecat.frames.frames import AudioRawFrame, Frame, StartInterruptionFrame
+from pipecat.frames.frames import AudioRawFrame, Frame, InputAudioRawFrame, StartInterruptionFrame
-from pipecat.serializers.base_serializer import FrameSerializer
+from pipecat.serializers.base_serializer import FrameSerializer, FrameSerializerType
 
 
 class TwilioFrameSerializer(FrameSerializer):
@@ -23,6 +23,10 @@ class TwilioFrameSerializer(FrameSerializer):
         self._stream_sid = stream_sid
         self._params = params
 
+    @property
+    def type(self) -> FrameSerializerType:
+        return FrameSerializerType.TEXT
+
     def serialize(self, frame: Frame) -> str | bytes | None:
         if isinstance(frame, AudioRawFrame):
             data = frame.audio
@@ -53,7 +57,7 @@ class TwilioFrameSerializer(FrameSerializer):
             deserialized_data = ulaw_to_pcm(
                 payload, self._params.twilio_sample_rate, self._params.sample_rate
             )
-            audio_frame = AudioRawFrame(
+            audio_frame = InputAudioRawFrame(
                 audio=deserialized_data, num_channels=1, sample_rate=self._params.sample_rate
             )
             return audio_frame

src/pipecat/transports/network/fastapi_websocket.py

@@ -8,20 +8,21 @@
 import asyncio
 import io
 import time
+import typing
 import wave
 
 from typing import Awaitable, Callable
 from pydantic.main import BaseModel
 
 from pipecat.frames.frames import (
-    AudioRawFrame,
     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
@@ -68,21 +69,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):
+            if isinstance(frame, InputAudioRawFrame):
-                await self.push_audio_frame(
+                await self.push_audio_frame(frame)
-                    InputAudioRawFrame(
+            else:
-                        audio=frame.audio,
+                await self.push_frame(frame)
-                        sample_rate=frame.sample_rate,
-                        num_channels=frame.num_channels,
-                    )
-                )
 
         await self._callbacks.on_client_disconnected(self._websocket)
 
@@ -110,29 +113,27 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
             await self._write_audio_sleep()
             return
 
-        frame = AudioRawFrame(
+        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()
+            with io.BytesIO() as buffer:
-            ww = wave.open(content, "wb")
+                with wave.open(buffer, "wb") as wf:
-            ww.setsampwidth(2)
+                    wf.setsampwidth(2)
-            ww.setnchannels(frame.num_channels)
+                    wf.setnchannels(frame.num_channels)
-            ww.setframerate(frame.sample_rate)
+                    wf.setframerate(frame.sample_rate)
-            ww.writeframes(frame.audio)
+                    wf.writeframes(frame.audio)
-            ww.close()
+                wav_frame = OutputAudioRawFrame(
-            content.seek(0)
+                    buffer.getvalue(),
-            wav_frame = AudioRawFrame(
+                    sample_rate=frame.sample_rate,
-                content.read(), sample_rate=frame.sample_rate, num_channels=frame.num_channels
+                    num_channels=frame.num_channels,
-            )
+                )
-            frame = wav_frame
+                frame = wav_frame
 
-        payload = self._params.serializer.serialize(frame)
+        await self._write_frame(frame)
-        if payload:
-            await self._websocket.send_text(payload)
 
         self._websocket_audio_buffer = bytes()
 
@@ -142,7 +143,13 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
     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.

src/pipecat/transports/network/websocket_server.py

@@ -13,11 +13,11 @@ from typing import Awaitable, Callable
 from pydantic.main import BaseModel
 
 from pipecat.frames.frames import (
-    AudioRawFrame,
     CancelFrame,
     EndFrame,
     Frame,
     InputAudioRawFrame,
+    OutputAudioRawFrame,
     StartFrame,
     StartInterruptionFrame,
 )
@@ -105,14 +105,8 @@ class WebsocketServerInputTransport(BaseInputTransport):
             if not frame:
                 continue
 
-            if isinstance(frame, AudioRawFrame):
+            if isinstance(frame, InputAudioRawFrame):
-                await self.push_audio_frame(
+                await self.push_audio_frame(frame)
-                    InputAudioRawFrame(
-                        audio=frame.audio,
-                        sample_rate=frame.sample_rate,
-                        num_channels=frame.num_channels,
-                    )
-                )
             else:
                 await self.push_frame(frame)
 
@@ -157,29 +151,27 @@ class WebsocketServerOutputTransport(BaseOutputTransport):
             await self._write_audio_sleep()
             return
 
-        frame = AudioRawFrame(
+        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()
+            with io.BytesIO() as buffer:
-            ww = wave.open(content, "wb")
+                with wave.open(buffer, "wb") as wf:
-            ww.setsampwidth(2)
+                    wf.setsampwidth(2)
-            ww.setnchannels(frame.num_channels)
+                    wf.setnchannels(frame.num_channels)
-            ww.setframerate(frame.sample_rate)
+                    wf.setframerate(frame.sample_rate)
-            ww.writeframes(frame.audio)
+                    wf.writeframes(frame.audio)
-            ww.close()
+                wav_frame = OutputAudioRawFrame(
-            content.seek(0)
+                    buffer.getvalue(),
-            wav_frame = AudioRawFrame(
+                    sample_rate=frame.sample_rate,
-                content.read(), sample_rate=frame.sample_rate, num_channels=frame.num_channels
+                    num_channels=frame.num_channels,
-            )
+                )
-            frame = wav_frame
+                frame = wav_frame
 
-        proto = self._params.serializer.serialize(frame)
+        await self._write_frame(frame)
-        if proto:
-            await self._websocket.send(proto)
 
         self._websocket_audio_buffer = bytes()