wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
FastAPIWebsocketTransport: fix to work with text and binary
This commit is contained in:
Aleix Conchillo Flaqué
2024-12-06 10:46:16 -08:00
committed by GitHub
parent 3cb7829624 842b3de7f5
commit d6f480fa50
8 changed files with 152 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
CHANGELOG.md
View File

@@ -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.

96
examples/websocket-server/index.html
View File

@@ -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) {

11
src/pipecat/serializers/base_serializer.py
View File

@@ -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

6
src/pipecat/serializers/livekit.py
View File

@@ -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

31
src/pipecat/serializers/protobuf.py
View File

@@ -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:

10
src/pipecat/serializers/twilio.py
View File

@@ -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

63
src/pipecat/transports/network/fastapi_websocket.py
View File

@@ -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.

42
src/pipecat/transports/network/websocket_server.py
View File

@@ -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()