getting started

src/dailyai/pipeline/aggregators.py

@@ -0,0 +1,182 @@
+import asyncio
+import re
+
+from tblib import Frame
+from dailyai.pipeline.frame_processor import FrameProcessor
+
+from dailyai.pipeline.frames import (
+    ControlQueueFrame,
+    EndStreamQueueFrame,
+    LLMMessagesQueueFrame,
+    QueueFrame,
+    TextQueueFrame,
+    TranscriptionQueueFrame,
+)
+from dailyai.pipeline.pipeline import Pipeline
+from dailyai.services.ai_services import AIService
+
+from typing import AsyncGenerator, List
+
+
+class LLMContextAggregator(AIService):
+    def __init__(
+        self,
+        messages: list[dict],
+        role: str,
+        bot_participant_id=None,
+        complete_sentences=True,
+        pass_through=True,
+    ):
+        super().__init__()
+        self.messages = messages
+        self.bot_participant_id = bot_participant_id
+        self.role = role
+        self.sentence = ""
+        self.complete_sentences = complete_sentences
+        self.pass_through = pass_through
+
+    async def process_frame(
+        self, frame: QueueFrame
+    ) -> AsyncGenerator[QueueFrame, None]:
+        # We don't do anything with non-text frames, pass it along to next in the pipeline.
+        if not isinstance(frame, TextQueueFrame):
+            yield frame
+            return
+
+        # Ignore transcription frames from the bot
+        if isinstance(frame, TranscriptionQueueFrame):
+            if frame.participantId == self.bot_participant_id:
+                return
+
+        # The common case for "pass through" is receiving frames from the LLM that we'll
+        # use to update the "assistant" LLM messages, but also passing the text frames
+        # along to a TTS service to be spoken to the user.
+        if self.pass_through:
+            yield frame
+
+        # TODO: split up transcription by participant
+        if self.complete_sentences:
+            # type: ignore -- the linter thinks this isn't a TextQueueFrame, even
+            # though we check it above
+            self.sentence += frame.text
+            if self.sentence.endswith((".", "?", "!")):
+                self.messages.append({"role": self.role, "content": self.sentence})
+                self.sentence = ""
+                yield LLMMessagesQueueFrame(self.messages)
+        else:
+            # type: ignore -- the linter thinks this isn't a TextQueueFrame, even
+            # though we check it above
+            self.messages.append({"role": self.role, "content": frame.text})
+            yield LLMMessagesQueueFrame(self.messages)
+
+    async def finalize(self) -> AsyncGenerator[QueueFrame, None]:
+        # Send any dangling words that weren't finished with punctuation.
+        if self.complete_sentences and self.sentence:
+            self.messages.append({"role": self.role, "content": self.sentence})
+            yield LLMMessagesQueueFrame(self.messages)
+
+
+class LLMUserContextAggregator(LLMContextAggregator):
+    def __init__(
+        self, messages: list[dict], bot_participant_id=None, complete_sentences=True
+    ):
+        super().__init__(
+            messages, "user", bot_participant_id, complete_sentences, pass_through=False
+        )
+
+
+class LLMAssistantContextAggregator(LLMContextAggregator):
+    def __init__(
+        self, messages: list[dict], bot_participant_id=None, complete_sentences=True
+    ):
+        super().__init__(
+            messages,
+            "assistant",
+            bot_participant_id,
+            complete_sentences,
+            pass_through=True,
+        )
+
+
+class SentenceAggregator(FrameProcessor):
+
+    def __init__(self):
+        self.aggregation = ""
+
+    async def process_frame(
+        self, frame: QueueFrame
+    ) -> AsyncGenerator[QueueFrame, None]:
+        if isinstance(frame, TextQueueFrame):
+            m = re.search("(.*[?.!])(.*)", frame.text)
+            if m:
+                yield TextQueueFrame(self.aggregation + m.group(1))
+                self.aggregation = m.group(2)
+            else:
+                self.aggregation += frame.text
+        elif isinstance(frame, EndStreamQueueFrame):
+            if self.aggregation:
+                yield TextQueueFrame(self.aggregation)
+            yield frame
+        else:
+            yield frame
+
+
+class StatelessTextTransformer(FrameProcessor):
+    def __init__(self, transform_fn):
+        self.transform_fn = transform_fn
+
+    async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
+        if isinstance(frame, TextQueueFrame):
+            yield TextQueueFrame(self.transform_fn(frame.text))
+        else:
+            yield frame
+
+class ParallelPipeline(FrameProcessor):
+    def __init__(self, pipeline_definitions: List[List[FrameProcessor]]):
+        self.sources = [asyncio.Queue() for _ in pipeline_definitions]
+        self.sink: asyncio.Queue[QueueFrame] = asyncio.Queue()
+        self.pipelines: list[Pipeline] = [
+            Pipeline(source, self.sink, pipeline_definition)
+            for source, pipeline_definition in zip(self.sources, pipeline_definitions)
+        ]
+
+    async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
+        # Short circuit, because we use EndStreamQueueFrame for our own internal process control.
+        if isinstance(frame, EndStreamQueueFrame):
+            yield frame
+
+        for source in self.sources:
+            await source.put(frame)
+            await source.put(EndStreamQueueFrame())
+
+        await asyncio.gather(*[pipeline.run_pipeline() for pipeline in self.pipelines])
+
+        while not self.sink.empty():
+            frame = await self.sink.get()
+            # Skip passing along EndStreamQueueFrames, because we use them for our own flow control.
+            if not isinstance(frame, EndStreamQueueFrame):
+                yield frame
+
+class GatedAccumulator(FrameProcessor):
+    def __init__(self, gate_open_fn, gate_close_fn, start_open):
+        self.gate_open_fn = gate_open_fn
+        self.gate_close_fn = gate_close_fn
+        self.gate_open = start_open
+        self.accumulator: List[QueueFrame] = []
+
+    async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
+        if self.gate_open:
+            if self.gate_close_fn(frame):
+                self.gate_open = False
+        else:
+            if self.gate_open_fn(frame):
+                self.gate_open = True
+
+        if self.gate_open:
+            yield frame
+            if self.accumulator:
+                for frame in self.accumulator:
+                    yield frame
+            self.accumulator = []
+        else:
+            self.accumulator.append(frame)

src/dailyai/pipeline/frame_processor.py

@@ -0,0 +1,24 @@
+from abc import abstractmethod
+from typing import AsyncGenerator
+
+from dailyai.pipeline.frames import ControlQueueFrame, QueueFrame
+
+
+class FrameProcessor:
+    @abstractmethod
+    async def process_frame(
+        self, frame: QueueFrame
+    ) -> AsyncGenerator[QueueFrame, None]:
+        if isinstance(frame, ControlQueueFrame):
+            yield frame
+
+    @abstractmethod
+    async def finalize(self) -> AsyncGenerator[QueueFrame, None]:
+        # This is a trick for the interpreter (and linter) to know that this is a generator.
+        if False:
+            yield QueueFrame()
+
+    @abstractmethod
+    async def interrupted(self) -> None:
+        pass
+

src/dailyai/pipeline/frames.py

@@ -0,0 +1,70 @@
+from dataclasses import dataclass
+from typing import Any
+
+
+class QueueFrame:
+    def __eq__(self, other):
+        return isinstance(other, self.__class__)
+
+
+class ControlQueueFrame(QueueFrame):
+    pass
+
+
+class StartStreamQueueFrame(ControlQueueFrame):
+    pass
+
+
+class EndStreamQueueFrame(ControlQueueFrame):
+    pass
+
+
+class LLMResponseStartQueueFrame(QueueFrame):
+    pass
+
+
+class LLMResponseEndQueueFrame(QueueFrame):
+    pass
+
+
+@dataclass()
+class AudioQueueFrame(QueueFrame):
+    data: bytes
+
+
+@dataclass()
+class ImageQueueFrame(QueueFrame):
+    url: str | None
+    image: bytes
+
+
+@dataclass()
+class SpriteQueueFrame(QueueFrame):
+    images: list[bytes]
+
+
+@dataclass()
+class TextQueueFrame(QueueFrame):
+    text: str
+
+
+@dataclass()
+class TranscriptionQueueFrame(TextQueueFrame):
+    participantId: str
+    timestamp: str
+
+
+@dataclass()
+class LLMMessagesQueueFrame(QueueFrame):
+    messages: list[dict[str, str]]  # TODO: define this more concretely!
+
+
+class AppMessageQueueFrame(QueueFrame):
+    message: Any
+    participantId: str
+
+class UserStartedSpeakingFrame(QueueFrame):
+    pass
+
+class UserStoppedSpeakingFrame(QueueFrame):
+    pass

src/dailyai/pipeline/pipeline.py

@@ -0,0 +1,42 @@
+import asyncio
+from typing import AsyncGenerator, List
+from dailyai.pipeline.frame_processor import FrameProcessor
+
+from dailyai.pipeline.frames import EndStreamQueueFrame, QueueFrame
+
+
+class Pipeline:
+    def __init__(
+        self,
+        source: asyncio.Queue,
+        sink: asyncio.Queue[QueueFrame],
+        processors: List[FrameProcessor],
+    ):
+        self.source: asyncio.Queue[QueueFrame] = source
+        self.sink: asyncio.Queue[QueueFrame] = sink
+        self.processors = processors
+
+    async def get_next_source_frame(self) -> AsyncGenerator[QueueFrame, None]:
+        yield await self.source.get()
+
+    async def run_pipeline(self):
+        try:
+            while True:
+                frame_generators = [self.get_next_source_frame()]
+                for processor in self.processors:
+                    next_frame_generators = []
+                    for frame_generator in frame_generators:
+                        async for frame in frame_generator:
+                            next_frame_generators.append(processor.process_frame(frame))
+                    frame_generators = next_frame_generators
+
+                for frame_generator in frame_generators:
+                    async for frame in frame_generator:
+                        await self.sink.put(frame)
+                    if isinstance(frame, EndStreamQueueFrame):
+                        return
+        except asyncio.CancelledError:
+            # this means there's been an interruption, do any cleanup necessary here.
+            for processor in self.processors:
+                await processor.interrupted()
+            pass