diff --git a/src/dailyai/queue_aggregators.py b/src/dailyai/queue_aggregators.py index c6a854582..3935a056f 100644 --- a/src/dailyai/queue_aggregators.py +++ b/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__( diff --git a/src/dailyai/services/ai_services.py b/src/dailyai/services/ai_services.py index 455145d81..3bac0ac21 100644 --- a/src/dailyai/services/ai_services.py +++ b/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 diff --git a/src/dailyai/tests/test_pipe_aggregators.py b/src/dailyai/tests/test_pipe_aggregators.py new file mode 100644 index 000000000..da19fef59 --- /dev/null +++ b/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) diff --git a/src/dailyai/tests/test_pipe_service.py b/src/dailyai/tests/test_pipe_service.py index fae3d11c9..0d1b57786 100644 --- a/src/dailyai/tests/test_pipe_service.py +++ b/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())