cleanup, make sample 4 work with new stuff

This commit is contained in:
Moishe Lettvin
2024-03-03 19:37:30 -05:00
parent 643be238f9
commit 15df4a9d58
5 changed files with 80 additions and 25 deletions

View File

@@ -6,6 +6,7 @@ from dailyai.pipeline.frame_processor import FrameProcessor
from dailyai.pipeline.frames import ( from dailyai.pipeline.frames import (
ControlQueueFrame, ControlQueueFrame,
EndParallelPipeQueueFrame,
EndStreamQueueFrame, EndStreamQueueFrame,
LLMMessagesQueueFrame, LLMMessagesQueueFrame,
QueueFrame, QueueFrame,
@@ -15,7 +16,7 @@ from dailyai.pipeline.frames import (
from dailyai.pipeline.pipeline import Pipeline from dailyai.pipeline.pipeline import Pipeline
from dailyai.services.ai_services import AIService from dailyai.services.ai_services import AIService
from typing import AsyncGenerator, List from typing import AsyncGenerator, Coroutine, List
class LLMContextAggregator(AIService): class LLMContextAggregator(AIService):
@@ -127,7 +128,11 @@ class StatelessTextTransformer(FrameProcessor):
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]: async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]:
if isinstance(frame, TextQueueFrame): 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: else:
yield frame yield frame
@@ -141,20 +146,16 @@ class ParallelPipeline(FrameProcessor):
] ]
async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]: 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: for source in self.sources:
await source.put(frame) await source.put(frame)
await source.put(EndStreamQueueFrame()) await source.put(EndParallelPipeQueueFrame())
await asyncio.gather(*[pipeline.run_pipeline() for pipeline in self.pipelines]) await asyncio.gather(*[pipeline.run_pipeline() for pipeline in self.pipelines])
while not self.sink.empty(): while not self.sink.empty():
frame = await self.sink.get() frame = await self.sink.get()
# Skip passing along EndStreamQueueFrames, because we use them for our own flow control. # Skip passing along EndParallelPipeQueueFrame, because we use them for our own flow control.
if not isinstance(frame, EndStreamQueueFrame): if not isinstance(frame, EndParallelPipeQueueFrame):
yield frame yield frame
class GatedAccumulator(FrameProcessor): class GatedAccumulator(FrameProcessor):

View File

@@ -18,6 +18,9 @@ class StartStreamQueueFrame(ControlQueueFrame):
class EndStreamQueueFrame(ControlQueueFrame): class EndStreamQueueFrame(ControlQueueFrame):
pass pass
class EndParallelPipeQueueFrame(ControlQueueFrame):
pass
class LLMResponseStartQueueFrame(QueueFrame): class LLMResponseStartQueueFrame(QueueFrame):
pass pass

View File

@@ -2,7 +2,7 @@ import asyncio
from typing import AsyncGenerator, List from typing import AsyncGenerator, List
from dailyai.pipeline.frame_processor import FrameProcessor from dailyai.pipeline.frame_processor import FrameProcessor
from dailyai.pipeline.frames import EndStreamQueueFrame, QueueFrame from dailyai.pipeline.frames import EndParallelPipeQueueFrame, EndStreamQueueFrame, QueueFrame
class Pipeline: class Pipeline:
@@ -33,7 +33,11 @@ class Pipeline:
for frame_generator in frame_generators: for frame_generator in frame_generators:
async for frame in frame_generator: async for frame in frame_generator:
await self.sink.put(frame) await self.sink.put(frame)
if isinstance(frame, EndStreamQueueFrame): if isinstance(
frame, EndStreamQueueFrame
) or isinstance(
frame, EndParallelPipeQueueFrame
):
return return
except asyncio.CancelledError: except asyncio.CancelledError:
# this means there's been an interruption, do any cleanup necessary here. # this means there's been an interruption, do any cleanup necessary here.

View File

@@ -1,11 +1,22 @@
import asyncio import asyncio
from doctest import OutputChecker import functools
from typing import Text
import unittest import unittest
import llm from dailyai.pipeline.aggregators import (
from dailyai.pipeline.aggregators import GatedAccumulator, SentenceAggregator, StatelessTextTransformer GatedAccumulator,
from dailyai.pipeline.frames import AudioQueueFrame, EndStreamQueueFrame, ImageQueueFrame, LLMResponseEndQueueFrame, LLMResponseStartQueueFrame, TextQueueFrame ParallelPipeline,
SentenceAggregator,
StatelessTextTransformer,
)
from dailyai.pipeline.frames import (
AudioQueueFrame,
EndStreamQueueFrame,
ImageQueueFrame,
LLMResponseEndQueueFrame,
LLMResponseStartQueueFrame,
QueueFrame,
TextQueueFrame,
)
from dailyai.pipeline.pipeline import Pipeline from dailyai.pipeline.pipeline import Pipeline
@@ -63,4 +74,42 @@ class TestDailyFrameAggregators(unittest.IsolatedAsyncioTestCase):
self.assertEqual(expected_output_frames, []) self.assertEqual(expected_output_frames, [])
async def test_parallel_pipeline(self): 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))

View File

@@ -2,6 +2,7 @@ import asyncio
import os import os
import aiohttp import aiohttp
from dailyai.pipeline.pipeline import Pipeline
from dailyai.services.daily_transport_service import DailyTransportService from dailyai.services.daily_transport_service import DailyTransportService
from dailyai.services.azure_ai_services import AzureLLMService, AzureTTSService 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 # will run in parallel with generating and speaking the audio for static text, so there's no delay to
# speak the LLM response. # speak the LLM response.
buffer_queue = asyncio.Queue() buffer_queue = asyncio.Queue()
llm_response_task = asyncio.create_task( source_queue = asyncio.Queue()
elevenlabs_tts.run_to_queue( pipeline = Pipeline(source = source_queue, sink=buffer_queue, processors=[llm, elevenlabs_tts])
buffer_queue, source_queue.put_nowait(LLMMessagesQueueFrame(messages))
llm.run([LLMMessagesQueueFrame(messages)]), pipeline_run_task = pipeline.run_pipeline()
True,
)
)
@transport.event_handler("on_first_other_participant_joined") @transport.event_handler("on_first_other_participant_joined")
async def on_first_other_participant_joined(transport): async def on_first_other_participant_joined(transport):
@@ -61,7 +59,7 @@ async def main(room_url: str):
if isinstance(frame, EndStreamQueueFrame): if isinstance(frame, EndStreamQueueFrame):
break 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() await transport.stop_when_done()