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

FastAPIWebsocketTransport: fix to work with text and binary

Browse Source
This commit is contained in:
Aleix Conchillo Flaqué
2024-12-05 23:53:45 -08:00
parent 3cb7829624
commit 842b3de7f5
8 changed files with 152 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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