From 38cb8ca6a31c4c18b24683be7820e90f5e8329f3 Mon Sep 17 00:00:00 2001 From: Moishe Lettvin Date: Tue, 27 Feb 2024 20:47:08 -0500 Subject: [PATCH] ...? --- src/dailyai/queue_aggregators.py | 101 ++++++++++++++++++++----- src/dailyai/services/ai_services.py | 15 +--- src/dailyai/tests/test_pipe_service.py | 29 ++++++- 3 files changed, 110 insertions(+), 35 deletions(-) diff --git a/src/dailyai/queue_aggregators.py b/src/dailyai/queue_aggregators.py index 182d7d479..c6a854582 100644 --- a/src/dailyai/queue_aggregators.py +++ b/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): diff --git a/src/dailyai/services/ai_services.py b/src/dailyai/services/ai_services.py index ce8b6dfb1..455145d81 100644 --- a/src/dailyai/services/ai_services.py +++ b/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 diff --git a/src/dailyai/tests/test_pipe_service.py b/src/dailyai/tests/test_pipe_service.py index 7a3fbf6db..fae3d11c9 100644 --- a/src/dailyai/tests/test_pipe_service.py +++ b/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)