merge aggregator & test

src/dailyai/queue_aggregators.py

@@ -1,6 +1,15 @@
 import asyncio
 
-from dailyai.queue_frame import EndStreamQueueFrame, LLMMessagesQueueFrame, QueueFrame, TextQueueFrame, TranscriptionQueueFrame
+from attr import dataclass
+
+from dailyai.queue_frame import (
+    ControlQueueFrame,
+    EndStreamQueueFrame,
+    LLMMessagesQueueFrame,
+    QueueFrame,
+    TextQueueFrame,
+    TranscriptionQueueFrame,
+)
 from dailyai.services.ai_services import AIService, PipeService
 
 from typing import Any, AsyncGenerator, Callable, List, Tuple
@@ -13,40 +22,51 @@ class SinkAndQueuePair:
 
 class QueueTee(PipeService):
     def __init__(
-        self, source: PipeService, sinks: list[PipeService]
+        self, sinks: list[PipeService], *args, **kwargs
     ):
-        self.source: PipeService = source
+        super().__init__(*args, **kwargs)
 
         self.sinks: List[SinkAndQueuePair] = []
         for sink in sinks:
-            pair = SinkAndQueuePair()
-            pair.sink = sink
-            pair.queue = asyncio.Queue()
+            pair = SinkAndQueuePair(sink, asyncio.Queue())
             self.sinks.append(pair)
+            sink.source_queue = pair.queue
 
     async def process_queue(self):
+        if not self.source_queue:
+            return
+
         while True:
-            frame = await self.source.sink_queue.get()
+            frame: QueueFrame = await self.source_queue.get()
+            print("got frame")
             for sink in self.sinks:
+                print("putting frame in sink")
                 await sink.queue.put(frame)
             if isinstance(frame, EndStreamQueueFrame):
                 break
 
 
-class QueueGateOnFrame(PipeService):
+class QueueFrameAggregator(PipeService):
 
     def __init__(
         self,
-        source: PipeService,
         aggregator: Callable[[Any, QueueFrame], Tuple[Any, QueueFrame | None]],
+        finalizer: Callable[[Any], QueueFrame | None],
+        *args,
+        **kwargs
     ):
-        self.source = source
+        super().__init__(*args, **kwargs)
         self.aggregator = aggregator
-        self.accumulation = None
+        self.finalizer = finalizer
+        self.aggregation = None
 
     async def process_frame(
         self, frame: QueueFrame
     ) -> AsyncGenerator[QueueFrame, None]:
+        if isinstance(frame, ControlQueueFrame):
+            yield frame
+            return
+
         output_frame: QueueFrame | None = None
         (self.aggregation, output_frame) = self.aggregator(
             self.aggregation, frame
@@ -54,38 +74,46 @@ class QueueGateOnFrame(PipeService):
         if output_frame:
             yield output_frame
 
+    async def finalize(self) -> AsyncGenerator[QueueFrame, None]:
+        output_frame = self.finalizer(self.aggregation)
+        if output_frame:
+            yield output_frame
+
 class QueueMergeGateOnFirst(PipeService):
 
     def __init__(
-        self, sources: List[PipeService], sink: asyncio.Queue[QueueFrame]
+        self, source_queues: List[asyncio.Queue[QueueFrame]]
     ):
-        self.sources = sources
-        self.sink = sink
+        super().__init__()
+        self.source_queues = source_queues
 
     async def process_queue(self):
         (frames): list[QueueFrame] = await asyncio.gather(
-            *[
-                source.sink_queue.get() for source in self.sources
-            ]
+            *[source_queue.get() for source_queue in self.source_queues]
         )
         for idx, frame in enumerate(frames):
-            await self.sink.put(frame)
+            print("frame", idx, frame)
 
-            # if the first frame we got from a source is an EndStreamQueueFrame, remove that source
+            # if the frame we got from a source is an EndStreamQueueFrame, remove that source
             if isinstance(frame, EndStreamQueueFrame):
-                self.sources.pop(idx)
+                self.source_queues.pop(idx)
+            else:
+                await self.sink_queue.put(frame)
 
         async def pass_through(sink, source):
             while True:
                 frame = await source.get()
-                await sink.put(frame)
                 if isinstance(frame, EndStreamQueueFrame):
                     break
+                else:
+                    await sink.put(frame)
 
         await asyncio.gather(
-            *[pass_through(self.sink, source) for source in self.sources]
+            *[pass_through(self.sink_queue, source) for source in self.source_queues]
         )
 
+        await self.sink_queue.put(EndStreamQueueFrame())
+
 
 class LLMContextAggregator(AIService):
     def __init__(

src/dailyai/services/ai_services.py

@@ -33,26 +33,28 @@ class PipeService(AbstractPipeService):
 
     def __init__(
         self,
-        source: AbstractPipeService | None = None,
+        source_queue: asyncio.Queue[QueueFrame] | None = None,
     ):
         super().__init__()
         self.logger: logging.Logger = logging.getLogger("dailyai")
-        self.source = source
+        self.source_queue = source_queue
 
     async def process_queue(self):
-        if not self.source:
+        if not self.source_queue:
             return
 
         while True:
-            frame: QueueFrame = await self.source.sink_queue.get()
+            frame: QueueFrame = await self.source_queue.get()
             async for output_frame in self.process_frame(frame):
-                await self.sink_queue.put(output_frame)
                 if isinstance(frame, EndStreamQueueFrame):
                     print("end of stream", type(self))
-                    async for output_frame in self.finalize():
-                        await self.sink_queue.put(output_frame)
+                    async for final_frame in self.finalize():
+                        await self.sink_queue.put(final_frame)
+                    await self.sink_queue.put(output_frame)
                     return
 
+                await self.sink_queue.put(output_frame)
+
     @abstractmethod
     async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
         yield frame

src/dailyai/tests/test_pipe_aggregators.py

@@ -0,0 +1,124 @@
+import asyncio
+import unittest
+
+from dailyai.queue_frame import EndStreamQueueFrame, TextQueueFrame
+from dailyai.services.ai_services import PipeService
+from dailyai.queue_aggregators import QueueFrameAggregator, QueueMergeGateOnFirst, QueueTee
+
+class IncomingPipeService(PipeService):
+    def __init__(self):
+        super().__init__()
+        self.sink_queue = asyncio.Queue()
+
+class QueueTeeTest(unittest.IsolatedAsyncioTestCase):
+
+    async def test_queue_tee(self):
+        originpipe = IncomingPipeService()
+        inpipe1 = PipeService(originpipe.sink_queue)
+        outpipe1 = PipeService()
+        outpipe2 = PipeService()
+        teepipe = QueueTee(source_queue=inpipe1.sink_queue, sinks=[outpipe1, outpipe2])
+
+        originpipe.sink_queue.put_nowait(TextQueueFrame("test"))
+        originpipe.sink_queue.put_nowait(EndStreamQueueFrame())
+
+        await asyncio.gather(*[pipe.process_queue() for pipe in [originpipe, inpipe1, outpipe1, outpipe2, teepipe]])
+
+        def validateOutputPipe(pipe: PipeService):
+            self.assertEqual(pipe.sink_queue.qsize(), 2)
+            frame = pipe.sink_queue.get_nowait()
+            self.assertIsInstance(frame, TextQueueFrame)
+            if isinstance(frame, TextQueueFrame):
+                self.assertEqual(frame.text, "test")
+            self.assertIsInstance(pipe.sink_queue.get_nowait(), EndStreamQueueFrame)
+
+        validateOutputPipe(outpipe1)
+        validateOutputPipe(outpipe2)
+
+class QueueFrameAggregatorTest(unittest.IsolatedAsyncioTestCase):
+    async def test_queue_frame_aggregator(self):
+        def aggregate_sentences(accumulation, frame):
+            if not accumulation:
+                accumulation = ""
+            if isinstance(frame, TextQueueFrame):
+                accumulation += frame.text
+            if accumulation.endswith((".", "!", "?")):
+                return ("", TextQueueFrame(accumulation))
+            return (accumulation, None)
+
+        def finalize_sentences(accumulation):
+            return TextQueueFrame(accumulation)
+
+        originpipe = IncomingPipeService()
+        aggregator_pipe = QueueFrameAggregator(
+            source_queue=originpipe.sink_queue,
+            aggregator=aggregate_sentences,
+            finalizer=finalize_sentences,
+        )
+
+        originpipe.sink_queue.put_nowait(TextQueueFrame("testing, "))
+        originpipe.sink_queue.put_nowait(TextQueueFrame("one."))
+        originpipe.sink_queue.put_nowait(TextQueueFrame("two."))
+        originpipe.sink_queue.put_nowait(TextQueueFrame("three."))
+        originpipe.sink_queue.put_nowait(TextQueueFrame("can you "))
+        originpipe.sink_queue.put_nowait(TextQueueFrame("hear me"))
+        originpipe.sink_queue.put_nowait(EndStreamQueueFrame())
+
+        await asyncio.gather(originpipe.process_queue(), aggregator_pipe.process_queue())
+
+        self.assertEqual(aggregator_pipe.sink_queue.qsize(), 5)
+        expected_text = ["testing, one.", "two.", "three.", "can you hear me"]
+        for exepectation in expected_text:
+            frame = aggregator_pipe.sink_queue.get_nowait()
+            print(frame)
+            self.assertIsInstance(frame, TextQueueFrame)
+            if isinstance(frame, TextQueueFrame):
+                self.assertEqual(frame.text, exepectation)
+
+        self.assertIsInstance(aggregator_pipe.sink_queue.get_nowait(), EndStreamQueueFrame)
+
+class QueueMergeGateOnFirstTest(unittest.IsolatedAsyncioTestCase):
+    async def test_queue_merge_gate_on_first(self):
+        pipe1 = IncomingPipeService()
+        pipe2 = IncomingPipeService()
+
+        merge_pipe = QueueMergeGateOnFirst(
+            source_queues=[pipe1.sink_queue, pipe2.sink_queue],
+        )
+
+        evt = asyncio.Event()
+
+        async def add_items_to_first_pipe():
+            await evt.wait()
+            await pipe1.sink_queue.put(TextQueueFrame("pipe1.1"))
+            await pipe1.sink_queue.put(TextQueueFrame("pipe1.2"))
+            await pipe1.sink_queue.put(EndStreamQueueFrame())
+
+        async def add_items_to_second_pipe():
+            await pipe2.sink_queue.put(TextQueueFrame("pipe2.1"))
+            evt.set()
+            await pipe2.sink_queue.put(EndStreamQueueFrame())
+
+        await asyncio.gather(
+            *[pipe.process_queue() for pipe in [pipe1, pipe2, merge_pipe]],
+            add_items_to_first_pipe(),
+            add_items_to_second_pipe())
+
+        self.assertEqual(merge_pipe.sink_queue.qsize(), 4)
+        frame = merge_pipe.sink_queue.get_nowait()
+        assert isinstance(frame, TextQueueFrame)
+        if isinstance(frame, TextQueueFrame):
+            self.assertEqual(frame.text, "pipe1.1")
+
+        frame = merge_pipe.sink_queue.get_nowait()
+        assert isinstance(frame, TextQueueFrame)
+        if isinstance(frame, TextQueueFrame):
+            self.assertEqual(frame.text, "pipe2.1")
+
+        frame = merge_pipe.sink_queue.get_nowait()
+        assert isinstance(frame, TextQueueFrame)
+        if isinstance(frame, TextQueueFrame):
+            self.assertEqual(frame.text, "pipe1.2")
+
+        frame = merge_pipe.sink_queue.get_nowait()
+        assert isinstance(frame, EndStreamQueueFrame)

src/dailyai/tests/test_pipe_service.py

@@ -1,22 +1,19 @@
 import asyncio
-from math import pi
 import unittest
 
-from unittest.mock import MagicMock, patch
-
-from dailyai.queue_frame import AudioQueueFrame, EndStreamQueueFrame, ImageQueueFrame, TextQueueFrame
+from dailyai.queue_frame import EndStreamQueueFrame, TextQueueFrame
 from dailyai.services.ai_services import PipeService
 
-class TestDailyTransport(unittest.IsolatedAsyncioTestCase):
+class TestPipeService(unittest.IsolatedAsyncioTestCase):
     class IncomingPipeService(PipeService):
         def __init__(self):
             super().__init__()
             self.sink_queue = asyncio.Queue()
 
     async def test_pipe_chain(self):
-        pipe1 = TestDailyTransport.IncomingPipeService()
-        pipe2 = PipeService(pipe1)
-        pipe3 = PipeService(pipe2)
+        pipe1 = TestPipeService.IncomingPipeService()
+        pipe2 = PipeService(pipe1.sink_queue)
+        pipe3 = PipeService(pipe2.sink_queue)
 
         await pipe1.sink_queue.put(TextQueueFrame("test"))
         await pipe1.sink_queue.put(EndStreamQueueFrame())