Websocket transport

src/dailyai/pipeline/frame_processor.py

@@ -23,8 +23,6 @@ class FrameProcessor:
         self, frame: Frame
     ) -> AsyncGenerator[Frame, None]:
         """Process a single frame and yield 0 or more frames."""
-        if isinstance(frame, ControlFrame):
-            yield frame
         yield frame
 
     @abstractmethod

src/dailyai/pipeline/frames.proto

@@ -0,0 +1,25 @@
+syntax = "proto3";
+
+package dailyai_proto;
+
+message TextFrame {
+    string text = 1;
+}
+
+message AudioFrame {
+    bytes data = 1;
+}
+
+message TranscriptionFrame {
+    string text = 1;
+    string participantId = 2;
+    string timestamp = 3;
+}
+
+message Frame {
+    oneof frame {
+        TextFrame text = 1;
+        AudioFrame audio = 2;
+        TranscriptionFrame transcription = 3;
+    }
+}

src/dailyai/pipeline/frames.py

@@ -2,6 +2,7 @@ from dataclasses import dataclass
 from typing import Any, List
 
 from dailyai.services.openai_llm_context import OpenAILLMContext
+import dailyai.pipeline.protobufs.frames_pb2 as frame_protos
 
 
 class Frame:
@@ -107,6 +108,22 @@ class TranscriptionQueueFrame(TextFrame):
     participantId: str
     timestamp: str
 
+    def __str__(self):
+        return f"{self.__class__.__name__}, text: '{self.text}' participantId: {self.participantId}, timestamp: {self.timestamp}"
+
+
+class TTSStartFrame(ControlFrame):
+    """Used to indicate the beginning of a TTS response. Following AudioFrames
+    are part of the TTS response until an TTEndFrame. These frames can be used
+    for aggregating audio frames in a transport to optimize the size of frames
+    sent to the session, without needing to control this in the TTS service."""
+    pass
+
+
+class TTSEndFrame(ControlFrame):
+    """Indicates the end of a TTS response."""
+    pass
+
 
 @dataclass()
 class LLMMessagesQueueFrame(Frame):

src/dailyai/pipeline/pipeline.py

@@ -24,7 +24,7 @@ class Pipeline:
         queues. If this pipeline is run by a transport, its sink and source queues
         will be overridden.
         """
-        self.processors: List[FrameProcessor] = processors
+        self._processors: List[FrameProcessor] = processors
 
         self.source: asyncio.Queue[Frame] = source or asyncio.Queue()
         self.sink: asyncio.Queue[Frame] = sink or asyncio.Queue()
@@ -40,6 +40,9 @@ class Pipeline:
         has processed a frame, its output will be placed on this queue."""
         self.sink = sink
 
+    def add_processor(self, processor: FrameProcessor):
+        self._processors.append(processor)
+
     async def get_next_source_frame(self) -> AsyncGenerator[Frame, None]:
         """Convenience function to get the next frame from the source queue. This
         lets us consistently have an AsyncGenerator yield frames, from either the
@@ -80,7 +83,7 @@ class Pipeline:
             while True:
                 initial_frame = await self.source.get()
                 async for frame in self._run_pipeline_recursively(
-                    initial_frame, self.processors
+                    initial_frame, self._processors
                 ):
                     await self.sink.put(frame)
 
@@ -91,7 +94,7 @@ class Pipeline:
         except asyncio.CancelledError:
             # this means there's been an interruption, do any cleanup necessary
             # here.
-            for processor in self.processors:
+            for processor in self._processors:
                 await processor.interrupted()
             pass
 

src/dailyai/pipeline/protobufs/frames_pb2.py

@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: frames.proto
+# Protobuf Python Version: 4.25.3
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0c\x66rames.proto\x12\rdailyai_proto\"\x19\n\tTextFrame\x12\x0c\n\x04text\x18\x01 \x01(\t\"\x1a\n\nAudioFrame\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\"L\n\x12TranscriptionFrame\x12\x0c\n\x04text\x18\x01 \x01(\t\x12\x15\n\rparticipantId\x18\x02 \x01(\t\x12\x11\n\ttimestamp\x18\x03 \x01(\t\"\xa2\x01\n\x05\x46rame\x12(\n\x04text\x18\x01 \x01(\x0b\x32\x18.dailyai_proto.TextFrameH\x00\x12*\n\x05\x61udio\x18\x02 \x01(\x0b\x32\x19.dailyai_proto.AudioFrameH\x00\x12:\n\rtranscription\x18\x03 \x01(\x0b\x32!.dailyai_proto.TranscriptionFrameH\x00\x42\x07\n\x05\x66rameb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'frames_pb2', _globals)
+if _descriptor._USE_C_DESCRIPTORS == False:
+  DESCRIPTOR._options = None
+  _globals['_TEXTFRAME']._serialized_start=31
+  _globals['_TEXTFRAME']._serialized_end=56
+  _globals['_AUDIOFRAME']._serialized_start=58
+  _globals['_AUDIOFRAME']._serialized_end=84
+  _globals['_TRANSCRIPTIONFRAME']._serialized_start=86
+  _globals['_TRANSCRIPTIONFRAME']._serialized_end=162
+  _globals['_FRAME']._serialized_start=165
+  _globals['_FRAME']._serialized_end=327
+# @@protoc_insertion_point(module_scope)

src/dailyai/serializers/abstract_frame_serializer.py

@@ -0,0 +1,16 @@
+from abc import abstractmethod
+
+from dailyai.pipeline.frames import Frame
+
+
+class FrameSerializer:
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def serialize(self, frame: Frame) -> bytes:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def deserialize(self, data: bytes) -> Frame:
+        raise NotImplementedError

src/dailyai/serializers/protobuf_serializer.py

@@ -0,0 +1,64 @@
+import dataclasses
+from typing import Text
+from dailyai.pipeline.frames import AudioFrame, Frame, TextFrame, TranscriptionQueueFrame
+import dailyai.pipeline.protobufs.frames_pb2 as frame_protos
+from dailyai.serializers.abstract_frame_serializer import FrameSerializer
+
+
+class ProtobufFrameSerializer(FrameSerializer):
+    SERIALIZABLE_TYPES = {
+        TextFrame: "text",
+        AudioFrame: "audio",
+        TranscriptionQueueFrame: "transcription"
+    }
+
+    SERIALIZABLE_FIELDS = {v: k for k, v in SERIALIZABLE_TYPES.items()}
+
+    def __init__(self):
+        pass
+
+    def serialize(self, frame: Frame) -> bytes:
+        proto_frame = frame_protos.Frame()
+        if type(frame) not in self.SERIALIZABLE_TYPES:
+            raise ValueError(
+                f"Frame type {type(frame)} is not serializable. You may need to add it to ProtobufFrameSerializer.SERIALIZABLE_FIELDS.")
+
+        # ignoring linter errors; we check that type(frame) is in this dict above
+        proto_optional_name = self.SERIALIZABLE_TYPES[type(frame)]  # type: ignore
+        for field in dataclasses.fields(frame):  # type: ignore
+            setattr(getattr(proto_frame, proto_optional_name), field.name,
+                    getattr(frame, field.name))
+
+        return proto_frame.SerializeToString()
+
+    def deserialize(self, data: bytes) -> Frame:
+        """Returns a Frame object from a Frame protobuf. Used to convert frames
+        passed over the wire as protobufs to Frame objects used in pipelines
+        and frame processors.
+
+        >>> serializer = ProtobufFrameSerializer()
+        >>> serializer.deserialize(
+        ...     serializer.serialize(AudioFrame(data=b'1234567890')))
+        AudioFrame(data=b'1234567890')
+
+        >>> serializer.deserialize(
+        ...     serializer.serialize(TextFrame(text='hello world')))
+        TextFrame(text='hello world')
+
+        >>> serializer.deserialize(serializer.serialize(TranscriptionQueueFrame(
+        ...     text="Hello there!", participantId="123", timestamp="2021-01-01")))
+        TranscriptionQueueFrame(text='Hello there!', participantId='123', timestamp='2021-01-01')
+        """
+
+        proto = frame_protos.Frame.FromString(data)
+        which = proto.WhichOneof("frame")
+        if which not in self.SERIALIZABLE_FIELDS:
+            raise ValueError(
+                "Proto does not contain a valid frame. You may need to add a new case to ProtobufFrameSerializer.deserialize.")
+
+        class_name = self.SERIALIZABLE_FIELDS[which]
+        args = getattr(proto, which)
+        args_dict = {}
+        for field in proto.DESCRIPTOR.fields_by_name[which].message_type.fields:
+            args_dict[field.name] = getattr(args, field.name)
+        return class_name(**args_dict)

src/dailyai/services/ai_services.py

@@ -10,6 +10,8 @@ from dailyai.pipeline.frames import (
     EndPipeFrame,
     ImageFrame,
     Frame,
+    TTSEndFrame,
+    TTSStartFrame,
     TextFrame,
     TranscriptionQueueFrame,
 )
@@ -47,12 +49,18 @@ class TTSService(AIService):
         # yield empty bytes here, so linting can infer what this method does
         yield bytes()
 
+    async def wrap_tts(self, text) -> AsyncGenerator[Frame, None]:
+        yield TTSStartFrame()
+        async for audio_chunk in self.run_tts(text):
+            yield AudioFrame(audio_chunk)
+        yield TTSEndFrame()
+        yield TextFrame(text)
+
     async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
         if isinstance(frame, EndFrame) or isinstance(frame, EndPipeFrame):
             if self.current_sentence:
-                async for audio_chunk in self.run_tts(self.current_sentence):
-                    yield AudioFrame(audio_chunk)
-                yield TextFrame(self.current_sentence)
+                async for cleanup_frame in self.wrap_tts(self.current_sentence):
+                    yield cleanup_frame
 
         if not isinstance(frame, TextFrame):
             yield frame
@@ -68,12 +76,8 @@ class TTSService(AIService):
                 self.current_sentence = ""
 
         if text:
-            async for audio_chunk in self.run_tts(text):
-                yield AudioFrame(audio_chunk)
-
-            # note we pass along the text frame *after* the audio, so the text
-            # frame is completed after the audio is processed.
-            yield TextFrame(text)
+            async for frame in self.wrap_tts(text):
+                yield frame
 
 
 class ImageGenService(AIService):

src/dailyai/services/local_stt_service.py

@@ -42,6 +42,7 @@ class LocalSTTService(STTService):
     async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
         """Processes a frame of audio data, either buffering or transcribing it."""
         if not isinstance(frame, AudioFrame):
+            yield frame
             return
 
         data = frame.data

src/dailyai/services/websocket_transport_service.py

@@ -0,0 +1,121 @@
+import asyncio
+import time
+from typing import AsyncGenerator, List
+import websockets
+
+from dailyai.pipeline.frame_processor import FrameProcessor
+from dailyai.pipeline.frames import AudioFrame, ControlFrame, EndFrame, Frame, TTSEndFrame, TTSStartFrame, TextFrame
+from dailyai.pipeline.pipeline import Pipeline
+from dailyai.serializers.protobuf_serializer import ProtobufFrameSerializer
+from dailyai.services.base_transport_service import BaseTransportService
+
+
+class WebSocketFrameProcessor(FrameProcessor):
+    """This FrameProcessor filters and mutates frames before they're sent over the websocket.
+    This is necessary to aggregate audio frames into sizes that are cleanly playable by the client"""
+
+    def __init__(
+            self,
+            audio_frame_size=16000,
+            sendable_frames: List[Frame] | None = None):
+        super().__init__()
+        self._audio_frame_size = audio_frame_size
+        self._sendable_frames = sendable_frames or [TextFrame, AudioFrame]
+        self._audio_buffer = bytes()
+        self._in_tts_audio = False
+
+    async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
+        if isinstance(frame, TTSStartFrame):
+            self._in_tts_audio = True
+        elif isinstance(frame, AudioFrame):
+            if self._in_tts_audio:
+                self._audio_buffer += frame.data
+                while len(self._audio_buffer) >= self._audio_frame_size:
+                    yield AudioFrame(self._audio_buffer[:self._audio_frame_size])
+                    self._audio_buffer = self._audio_buffer[self._audio_frame_size:]
+        elif isinstance(frame, TTSEndFrame):
+            self._in_tts_audio = False
+            if self._audio_buffer:
+                yield AudioFrame(self._audio_buffer)
+                self._audio_buffer = bytes()
+        elif type(frame) in self._sendable_frames or isinstance(frame, ControlFrame):
+            yield frame
+
+
+class WebsocketTransport(BaseTransportService):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._sample_width = kwargs.get("sample_width", 2)
+        self._n_channels = kwargs.get("n_channels", 1)
+        self._port = kwargs.get("port", 8765)
+        self._host = kwargs.get("host", "localhost")
+        self._audio_frame_size = kwargs.get("audio_frame_size", 16000)
+        self._sendable_frames = kwargs.get(
+            "sendable_frames", [
+                TextFrame, AudioFrame, TTSEndFrame, TTSStartFrame])
+        self._serializer = kwargs.get("serializer", ProtobufFrameSerializer())
+
+        if self._camera_enabled:
+            raise ValueError(
+                "Camera is not supported in WebsocketTransportService")
+
+        if self._speaker_enabled:
+            self._speaker_buffer_pending = bytearray()
+
+        self._server: websockets.WebSocketServer | None = None
+        self._websocket: websockets.WebSocketServerProtocol | None = None
+
+        self._connection_handlers = []
+
+    async def run(self, pipeline: Pipeline, override_pipeline_source_queue=True):
+        self._stop_server_event = asyncio.Event()
+        pipeline.set_sink(self.send_queue)
+        if override_pipeline_source_queue:
+            pipeline.set_source(self.receive_queue)
+
+        pipeline.add_processor(WebSocketFrameProcessor(
+            audio_frame_size=self._audio_frame_size,
+            sendable_frames=self._sendable_frames))
+
+        async def timeout():
+            sleep_time = self._expiration - time.time()
+            await asyncio.sleep(sleep_time)
+            self._stop_server_event.set()
+
+        async def send_task():
+            while not self._stop_server_event.is_set():
+                frame = await self.send_queue.get()
+                if isinstance(frame, EndFrame):
+                    self._stop_server_event.set()
+                    break
+                if self._websocket and frame:
+                    proto = self._serializer.serialize(frame)
+                    await self._websocket.send(proto)
+
+        async def start_server():
+            async with websockets.serve(self._websocket_handler, self._host, self._port) as server:
+                self._logger.debug("Websocket server started.")
+                await self._stop_server_event.wait()
+                self._logger.debug("Websocket server stopped.")
+            await self.receive_queue.put(EndFrame())
+
+        timeout_task = asyncio.create_task(timeout())
+        await asyncio.gather(start_server(), send_task(), pipeline.run_pipeline())
+        timeout_task.cancel()
+
+    def on_connection(self, handler):
+        self._connection_handlers.append(handler)
+
+    async def _websocket_handler(self, websocket: websockets.WebSocketServerProtocol, path):
+        if self._websocket:
+            await self._websocket.close()
+            self._logger.warning(
+                "Got another websocket connection; closing first.")
+
+        for handler in self._connection_handlers:
+            await handler()
+
+        self._websocket = websocket
+        async for message in websocket:
+            frame = self._serializer.deserialize(message)
+            await self.receive_queue.put(frame)