...?

src/dailyai/queue_aggregators.py

@@ -1,29 +1,90 @@
 import asyncio
 
-from dailyai.queue_frame import LLMMessagesQueueFrame, QueueFrame, TextQueueFrame, TranscriptionQueueFrame
-from dailyai.services.ai_services import AIService
+from dailyai.queue_frame import EndStreamQueueFrame, LLMMessagesQueueFrame, QueueFrame, TextQueueFrame, TranscriptionQueueFrame
+from dailyai.services.ai_services import AIService, PipeService
 
-from typing import AsyncGenerator, List
+from typing import Any, AsyncGenerator, Callable, List, Tuple
+
+@dataclass
+class SinkAndQueuePair:
+    sink: PipeService
+    queue: asyncio.Queue
 
 
-class QueueTee:
-    async def run_to_queue_and_generate(
-        self,
-        output_queue: asyncio.Queue,
-        generator: AsyncGenerator[QueueFrame, None]
-    ) -> AsyncGenerator[QueueFrame, None]:
-        async for frame in generator:
-            await output_queue.put(frame)
-            yield frame
-
-    async def run_to_queues(
-        self,
-        output_queues: List[asyncio.Queue],
-        generator: AsyncGenerator[QueueFrame, None]
+class QueueTee(PipeService):
+    def __init__(
+        self, source: PipeService, sinks: list[PipeService]
     ):
-        async for frame in generator:
-            for queue in output_queues:
-                await queue.put(frame)
+        self.source: PipeService = source
+
+        self.sinks: List[SinkAndQueuePair] = []
+        for sink in sinks:
+            pair = SinkAndQueuePair()
+            pair.sink = sink
+            pair.queue = asyncio.Queue()
+            self.sinks.append(pair)
+
+    async def process_queue(self):
+        while True:
+            frame = await self.source.sink_queue.get()
+            for sink in self.sinks:
+                await sink.queue.put(frame)
+            if isinstance(frame, EndStreamQueueFrame):
+                break
+
+
+class QueueGateOnFrame(PipeService):
+
+    def __init__(
+        self,
+        source: PipeService,
+        aggregator: Callable[[Any, QueueFrame], Tuple[Any, QueueFrame | None]],
+    ):
+        self.source = source
+        self.aggregator = aggregator
+        self.accumulation = None
+
+    async def process_frame(
+        self, frame: QueueFrame
+    ) -> AsyncGenerator[QueueFrame, None]:
+        output_frame: QueueFrame | None = None
+        (self.aggregation, output_frame) = self.aggregator(
+            self.aggregation, frame
+        )
+        if output_frame:
+            yield output_frame
+
+class QueueMergeGateOnFirst(PipeService):
+
+    def __init__(
+        self, sources: List[PipeService], sink: asyncio.Queue[QueueFrame]
+    ):
+        self.sources = sources
+        self.sink = sink
+
+    async def process_queue(self):
+        (frames): list[QueueFrame] = await asyncio.gather(
+            *[
+                source.sink_queue.get() for source in self.sources
+            ]
+        )
+        for idx, frame in enumerate(frames):
+            await self.sink.put(frame)
+
+            # if the first frame we got from a source is an EndStreamQueueFrame, remove that source
+            if isinstance(frame, EndStreamQueueFrame):
+                self.sources.pop(idx)
+
+        async def pass_through(sink, source):
+            while True:
+                frame = await source.get()
+                await sink.put(frame)
+                if isinstance(frame, EndStreamQueueFrame):
+                    break
+
+        await asyncio.gather(
+            *[pass_through(self.sink, source) for source in self.sources]
+        )
 
 
 class LLMContextAggregator(AIService):

src/dailyai/services/ai_services.py

@@ -23,15 +23,8 @@ class AbstractPipeService:
     def __init__(
         self,
     ):
-        self.source_queue: asyncio.Queue[QueueFrame] = asyncio.Queue()
         self.sink_queue: asyncio.Queue[QueueFrame] = asyncio.Queue()
 
-    async def get(self) -> QueueFrame:
-        return await self.sink_queue.get()
-
-    async def put(self, frame: QueueFrame) -> None:
-        await self.source_queue.put(frame)
-
     @abstractmethod
     async def process_queue(self):
         pass
@@ -41,25 +34,23 @@ class PipeService(AbstractPipeService):
     def __init__(
         self,
         source: AbstractPipeService | None = None,
-        sink: AbstractPipeService | None = None,
     ):
         super().__init__()
         self.logger: logging.Logger = logging.getLogger("dailyai")
         self.source = source
-        self.sink = sink
 
     async def process_queue(self):
         if not self.source:
             return
 
         while True:
-            frame: QueueFrame = await self.source.get()
+            frame: QueueFrame = await self.source.sink_queue.get()
             async for output_frame in self.process_frame(frame):
-                await self.sink.put(output_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.put(output_frame)
+                        await self.sink_queue.put(output_frame)
                     return
 
     @abstractmethod

src/dailyai/tests/test_pipe_service.py

@@ -1,10 +1,33 @@
+import asyncio
+from math import pi
 import unittest
 
 from unittest.mock import MagicMock, patch
 
-from dailyai.queue_frame import AudioQueueFrame, ImageQueueFrame
+from dailyai.queue_frame import AudioQueueFrame, EndStreamQueueFrame, ImageQueueFrame, TextQueueFrame
 from dailyai.services.ai_services import PipeService
 
 class TestDailyTransport(unittest.IsolatedAsyncioTestCase):
-    def test_pipe_chain(self):
-        pipe1 = PipeService()
+    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)
+
+        await pipe1.sink_queue.put(TextQueueFrame("test"))
+        await pipe1.sink_queue.put(EndStreamQueueFrame())
+
+        await asyncio.gather(pipe1.process_queue(), pipe2.process_queue(), pipe3.process_queue())
+
+        self.assertEqual(pipe3.sink_queue.qsize(), 2)
+        frame = await pipe3.sink_queue.get()
+        self.assertIsInstance(frame, TextQueueFrame)
+        if isinstance(frame, TextQueueFrame):
+            self.assertEqual(frame.text, "test")
+
+        frame = await pipe3.sink_queue.get()
+        self.assertIsInstance(frame, EndStreamQueueFrame)