From 842b3de7f59a1eb3b66545377b51b17321b53a02 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Thu, 5 Dec 2024 23:53:45 -0800 Subject: [PATCH] FastAPIWebsocketTransport: fix to work with text and binary --- CHANGELOG.md | 3 + examples/websocket-server/index.html | 96 ++++++++++--------- src/pipecat/serializers/base_serializer.py | 11 +++ src/pipecat/serializers/livekit.py | 6 +- src/pipecat/serializers/protobuf.py | 31 ++++-- src/pipecat/serializers/twilio.py | 10 +- .../transports/network/fastapi_websocket.py | 63 ++++++------ .../transports/network/websocket_server.py | 42 ++++---- 8 files changed, 152 insertions(+), 110 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 85f1e20e5..211193fb4 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/examples/websocket-server/index.html b/examples/websocket-server/index.html index 514a4a821..ac93c62d7 100644 --- a/examples/websocket-server/index.html +++ b/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"); - const progressText = document.getElementById("progressText"); - progressText.textContent = "We are ready! Make sure to run the server and then click `Start Audio`."; + Frame = root.lookupType('pipecat.Frame'); + const progressText = document.getElementById('progressText'); + 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) { - const arrayBuffer = await event.data.arrayBuffer(); + function handleWebSocketOpen(event) { + 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) { diff --git a/src/pipecat/serializers/base_serializer.py b/src/pipecat/serializers/base_serializer.py index 96f5fd214..00f2f2a68 100644 --- a/src/pipecat/serializers/base_serializer.py +++ b/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 diff --git a/src/pipecat/serializers/livekit.py b/src/pipecat/serializers/livekit.py index 29d32b861..a14483b15 100644 --- a/src/pipecat/serializers/livekit.py +++ b/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 diff --git a/src/pipecat/serializers/protobuf.py b/src/pipecat/serializers/protobuf.py index 654e87b24..4c5cee0a4 100644 --- a/src/pipecat/serializers/protobuf.py +++ b/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.serializers.base_serializer import FrameSerializer +from pipecat.frames.frames import ( + 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 result + return proto_frame.SerializeToString() 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: diff --git a/src/pipecat/serializers/twilio.py b/src/pipecat/serializers/twilio.py index ebc62e484..a0d02fa2f 100644 --- a/src/pipecat/serializers/twilio.py +++ b/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.serializers.base_serializer import FrameSerializer +from pipecat.frames.frames import AudioRawFrame, Frame, InputAudioRawFrame, StartInterruptionFrame +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 diff --git a/src/pipecat/transports/network/fastapi_websocket.py b/src/pipecat/transports/network/fastapi_websocket.py index 0ca969463..de1cf531e 100644 --- a/src/pipecat/transports/network/fastapi_websocket.py +++ b/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): - 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) @@ -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() - 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 - payload = self._params.serializer.serialize(frame) - if payload: - await self._websocket.send_text(payload) + await self._write_frame(frame) 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. diff --git a/src/pipecat/transports/network/websocket_server.py b/src/pipecat/transports/network/websocket_server.py index 2b88a7334..567dbac7d 100644 --- a/src/pipecat/transports/network/websocket_server.py +++ b/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): - 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) @@ -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() - 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 - proto = self._params.serializer.serialize(frame) - if proto: - await self._websocket.send(proto) + await self._write_frame(frame) self._websocket_audio_buffer = bytes()