diff --git a/src/dailyai/pipeline/aggregators.py b/src/dailyai/pipeline/aggregators.py index ef9d1e7ae..742555006 100644 --- a/src/dailyai/pipeline/aggregators.py +++ b/src/dailyai/pipeline/aggregators.py @@ -6,6 +6,7 @@ from dailyai.pipeline.frame_processor import FrameProcessor from dailyai.pipeline.frames import ( ControlQueueFrame, + EndParallelPipeQueueFrame, EndStreamQueueFrame, LLMMessagesQueueFrame, QueueFrame, @@ -15,7 +16,7 @@ from dailyai.pipeline.frames import ( from dailyai.pipeline.pipeline import Pipeline from dailyai.services.ai_services import AIService -from typing import AsyncGenerator, List +from typing import AsyncGenerator, Coroutine, List class LLMContextAggregator(AIService): @@ -127,7 +128,11 @@ class StatelessTextTransformer(FrameProcessor): async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]: if isinstance(frame, TextQueueFrame): - yield TextQueueFrame(self.transform_fn(frame.text)) + result = self.transform_fn(frame.text) + if isinstance(result, Coroutine): + result = await result + + yield TextQueueFrame(result) else: yield frame @@ -141,20 +146,16 @@ class ParallelPipeline(FrameProcessor): ] 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 source.put(EndParallelPipeQueueFrame()) 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): + # Skip passing along EndParallelPipeQueueFrame, because we use them for our own flow control. + if not isinstance(frame, EndParallelPipeQueueFrame): yield frame class GatedAccumulator(FrameProcessor): diff --git a/src/dailyai/pipeline/frames.py b/src/dailyai/pipeline/frames.py index f0542e561..e8a31d0eb 100644 --- a/src/dailyai/pipeline/frames.py +++ b/src/dailyai/pipeline/frames.py @@ -18,6 +18,9 @@ class StartStreamQueueFrame(ControlQueueFrame): class EndStreamQueueFrame(ControlQueueFrame): pass +class EndParallelPipeQueueFrame(ControlQueueFrame): + pass + class LLMResponseStartQueueFrame(QueueFrame): pass diff --git a/src/dailyai/pipeline/pipeline.py b/src/dailyai/pipeline/pipeline.py index 0ecc594bb..780d31750 100644 --- a/src/dailyai/pipeline/pipeline.py +++ b/src/dailyai/pipeline/pipeline.py @@ -2,7 +2,7 @@ import asyncio from typing import AsyncGenerator, List from dailyai.pipeline.frame_processor import FrameProcessor -from dailyai.pipeline.frames import EndStreamQueueFrame, QueueFrame +from dailyai.pipeline.frames import EndParallelPipeQueueFrame, EndStreamQueueFrame, QueueFrame class Pipeline: @@ -33,7 +33,11 @@ class Pipeline: for frame_generator in frame_generators: async for frame in frame_generator: await self.sink.put(frame) - if isinstance(frame, EndStreamQueueFrame): + if isinstance( + frame, EndStreamQueueFrame + ) or isinstance( + frame, EndParallelPipeQueueFrame + ): return except asyncio.CancelledError: # this means there's been an interruption, do any cleanup necessary here. diff --git a/src/dailyai/tests/test_aggregators.py b/src/dailyai/tests/test_aggregators.py index 4e4a5b304..986b28226 100644 --- a/src/dailyai/tests/test_aggregators.py +++ b/src/dailyai/tests/test_aggregators.py @@ -1,11 +1,22 @@ import asyncio -from doctest import OutputChecker -from typing import Text +import functools import unittest -import llm -from dailyai.pipeline.aggregators import GatedAccumulator, SentenceAggregator, StatelessTextTransformer -from dailyai.pipeline.frames import AudioQueueFrame, EndStreamQueueFrame, ImageQueueFrame, LLMResponseEndQueueFrame, LLMResponseStartQueueFrame, TextQueueFrame +from dailyai.pipeline.aggregators import ( + GatedAccumulator, + ParallelPipeline, + SentenceAggregator, + StatelessTextTransformer, +) +from dailyai.pipeline.frames import ( + AudioQueueFrame, + EndStreamQueueFrame, + ImageQueueFrame, + LLMResponseEndQueueFrame, + LLMResponseStartQueueFrame, + QueueFrame, + TextQueueFrame, +) from dailyai.pipeline.pipeline import Pipeline @@ -63,4 +74,42 @@ class TestDailyFrameAggregators(unittest.IsolatedAsyncioTestCase): self.assertEqual(expected_output_frames, []) async def test_parallel_pipeline(self): - pass + + async def slow_add(sleep_time:float, name:str, x: str): + await asyncio.sleep(sleep_time) + return ":".join([x, name]) + + pipe1_annotation = StatelessTextTransformer(functools.partial(slow_add, 0.1, 'pipe1')) + pipe2_annotation = StatelessTextTransformer(functools.partial(slow_add, 0.2, 'pipe2')) + sentence_aggregator = SentenceAggregator() + add_dots = StatelessTextTransformer(lambda x: x + ".") + + source = asyncio.Queue() + sink = asyncio.Queue() + pipeline = Pipeline( + source, + sink, + [ParallelPipeline([[pipe1_annotation], [sentence_aggregator, pipe2_annotation]]), add_dots], + ) + + frames = [ + TextQueueFrame("Hello, "), + TextQueueFrame("world."), + EndStreamQueueFrame() + ] + + expected_output_frames: list[QueueFrame] = [ + TextQueueFrame(text='Hello, :pipe1.'), + TextQueueFrame(text='world.:pipe1.'), + TextQueueFrame(text='Hello, world.:pipe2.'), + EndStreamQueueFrame() + ] + + for frame in frames: + await source.put(frame) + + await pipeline.run_pipeline() + + while not sink.empty(): + frame = await sink.get() + self.assertEqual(frame, expected_output_frames.pop(0)) diff --git a/src/examples/foundational/04-utterance-and-speech.py b/src/examples/foundational/04-utterance-and-speech.py index b46ef8597..57ba3dc91 100644 --- a/src/examples/foundational/04-utterance-and-speech.py +++ b/src/examples/foundational/04-utterance-and-speech.py @@ -2,6 +2,7 @@ import asyncio import os import aiohttp +from dailyai.pipeline.pipeline import Pipeline from dailyai.services.daily_transport_service import DailyTransportService from dailyai.services.azure_ai_services import AzureLLMService, AzureTTSService @@ -41,13 +42,10 @@ async def main(room_url: str): # will run in parallel with generating and speaking the audio for static text, so there's no delay to # speak the LLM response. buffer_queue = asyncio.Queue() - llm_response_task = asyncio.create_task( - elevenlabs_tts.run_to_queue( - buffer_queue, - llm.run([LLMMessagesQueueFrame(messages)]), - True, - ) - ) + source_queue = asyncio.Queue() + pipeline = Pipeline(source = source_queue, sink=buffer_queue, processors=[llm, elevenlabs_tts]) + source_queue.put_nowait(LLMMessagesQueueFrame(messages)) + pipeline_run_task = pipeline.run_pipeline() @transport.event_handler("on_first_other_participant_joined") async def on_first_other_participant_joined(transport): @@ -61,7 +59,7 @@ async def main(room_url: str): if isinstance(frame, EndStreamQueueFrame): break - await asyncio.gather(llm_response_task, buffer_to_send_queue()) + await asyncio.gather(pipeline_run_task, buffer_to_send_queue()) await transport.stop_when_done()