From 337ca7f5811e58b0e4fe0c9fae1506cb6def7f6e Mon Sep 17 00:00:00 2001 From: Moishe Lettvin Date: Wed, 6 Mar 2024 12:52:44 -0500 Subject: [PATCH 1/3] frame and pipeline docstrings --- src/dailyai/pipeline/frame_processor.py | 30 ++++++++++---------- src/dailyai/pipeline/pipeline.py | 37 +++++++++++++++++++++---- src/dailyai/services/ai_services.py | 15 ++++++++-- 3 files changed, 59 insertions(+), 23 deletions(-) diff --git a/src/dailyai/pipeline/frame_processor.py b/src/dailyai/pipeline/frame_processor.py index 2c5367c1b..d70294ac0 100644 --- a/src/dailyai/pipeline/frame_processor.py +++ b/src/dailyai/pipeline/frame_processor.py @@ -3,27 +3,29 @@ from typing import AsyncGenerator from dailyai.pipeline.frames import ControlFrame, Frame -""" -This is the base class for all frame processors. Frame processors consume a frame -and yield 0 or more frames. Generally frame processors are used as part of a pipeline, -where frames come from a source queue, are processed by a series of frame processors, -then placed on a sink queue. - -By convention, FrameProcessors should immediately yield any frames they don't process. - -Stateful FrameProcessors should watch for the EndStreamQueueFrame and finalize their -output, eg. yielding an unfinished sentence if they're aggregating LLM output to full -sentences. EndStreamQueueFrame is also a chance to clean up any services that need to -be closed, del'd, etc. -""" class FrameProcessor: + """This is the base class for all frame processors. Frame processors consume a frame + and yield 0 or more frames. Generally frame processors are used as part of a pipeline + where frames come from a source queue, are processed by a series of frame processors, + then placed on a sink queue. + + By convention, FrameProcessors should immediately yield any frames they don't process. + + Stateful FrameProcessors should watch for the EndStreamQueueFrame and finalize their + output, eg. yielding an unfinished sentence if they're aggregating LLM output to full + sentences. EndStreamQueueFrame is also a chance to clean up any services that need to + be closed, del'd, etc. + """ + @abstractmethod async def process_frame( self, frame: Frame ) -> AsyncGenerator[Frame, None]: + """Process a single frame and yield 0 or more frames.""" if isinstance(frame, ControlFrame): yield frame + yield frame @abstractmethod async def finalize(self) -> AsyncGenerator[Frame, None]: @@ -33,5 +35,5 @@ class FrameProcessor: @abstractmethod async def interrupted(self) -> None: + """Handle any cleanup if the pipeline was interrupted.""" pass - diff --git a/src/dailyai/pipeline/pipeline.py b/src/dailyai/pipeline/pipeline.py index dad4bc042..8e6c13d2b 100644 --- a/src/dailyai/pipeline/pipeline.py +++ b/src/dailyai/pipeline/pipeline.py @@ -4,14 +4,14 @@ from dailyai.pipeline.frame_processor import FrameProcessor from dailyai.pipeline.frames import EndPipeFrame, EndFrame, Frame -""" -This class manages a pipe of FrameProcessors, and runs them in sequence. The "source" -and "sink" queues are managed by the caller. You can use this class stand-alone to -perform specialized processing, or you can use the Transport's run_pipeline method to -instantiate and run a pipeline with the Transport's sink and source queues. -""" class Pipeline: + """ + This class manages a pipe of FrameProcessors, and runs them in sequence. The "source" + and "sink" queues are managed by the caller. You can use this class stand-alone to + perform specialized processing, or you can use the Transport's run_pipeline method to + instantiate and run a pipeline with the Transport's sink and source queues. + """ def __init__( self, @@ -19,22 +19,47 @@ class Pipeline: source: asyncio.Queue | None = None, sink: asyncio.Queue[Frame] | None = None, ): + """ Create a new pipeline. By default neither the source nor sink + queues are set, so you'll need to pass them to this constructor or + call set_source and set_sink before using the pipeline. Note that + the transport's run_*_pipeline methods will set the source and sink + queues on the pipeline for you. + """ self.processors = processors self.source: asyncio.Queue[Frame] | None = source self.sink: asyncio.Queue[Frame] | None = sink def set_source(self, source: asyncio.Queue[Frame]): + """ Set the source queue for this pipeline. Frames from this queue + will be processed by each frame_processor in the pipeline, or order + from first to last. """ self.source = source def set_sink(self, sink: asyncio.Queue[Frame]): + """ Set the sink queue for this pipeline. After the last frame_processor + has processed a frame, its output will be placed on this queue.""" self.sink = sink async def get_next_source_frame(self) -> AsyncGenerator[Frame, None]: + """ Convenience function to get the next frame from the source queue. This + lets us consistently have an AsyncGenerator yield frames, from either the + source queue or a frame_processor.""" if self.source is None: raise ValueError("Source queue not set") yield await self.source.get() async def run_pipeline(self): + """ Run the pipeline. Take each frame from the source queue, pass it to + the first frame_processor, pass the output of that frame_processor to the + next in the list, etc. until the last frame_processor has processed the + resulting frames, then place those frames in the sink queue. + + The source and sink queues must be set before calling this method. + + This method will exit when an EndStreamQueueFrame is placed on the sink queue. + No more frames will be placed on the sink queue after an EndStreamQueueFrame, even + if it's not the last frame yielded by the last frame_processor in the pipeline..""" + if self.source is None or self.sink is None: raise ValueError("Source or sink queue not set") diff --git a/src/dailyai/services/ai_services.py b/src/dailyai/services/ai_services.py index 61f043bee..0fca44884 100644 --- a/src/dailyai/services/ai_services.py +++ b/src/dailyai/services/ai_services.py @@ -59,9 +59,6 @@ class AIService(FrameProcessor): break else: raise Exception("Frames must be an iterable or async iterable") - - async for output_frame in self.finalize(): - yield output_frame except Exception as e: self.logger.error("Exception occurred while running AI service", e) raise e @@ -103,6 +100,7 @@ class TTSService(AIService): # yield empty bytes here, so linting can infer what this method does yield bytes() +<<<<<<< HEAD async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]: if isinstance(frame, EndFrame): if self.current_sentence: @@ -111,6 +109,14 @@ class TTSService(AIService): yield frame if not isinstance(frame, TextFrame): +======= + async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]: + if not isinstance(frame, TextQueueFrame): + if self.current_sentence: + async for audio_chunk in self.run_tts(self.current_sentence): + yield AudioQueueFrame(audio_chunk) + +>>>>>>> fa5f38c (frame and pipeline docstrings) yield frame return @@ -127,9 +133,12 @@ class TTSService(AIService): async for audio_chunk in self.run_tts(text): yield AudioFrame(audio_chunk) +<<<<<<< HEAD # note we pass along the text frame *after* the audio, so the text frame is completed after the audio is processed. yield TextFrame(text) +======= +>>>>>>> fa5f38c (frame and pipeline docstrings) # Convenience function to send the audio for a sentence to the given queue async def say(self, sentence, queue: asyncio.Queue): await self.run_to_queue(queue, [TextFrame(sentence)]) From 8fb92e3fd796d1abc7a32c9e160be29200ddc260 Mon Sep 17 00:00:00 2001 From: Moishe Lettvin Date: Thu, 7 Mar 2024 12:51:19 -0500 Subject: [PATCH 2/3] Getting started on docstrings --- src/dailyai/pipeline/aggregators.py | 71 +++++++++++++++++++++---- src/dailyai/pipeline/frame_processor.py | 6 --- src/dailyai/services/ai_services.py | 14 +---- 3 files changed, 63 insertions(+), 28 deletions(-) diff --git a/src/dailyai/pipeline/aggregators.py b/src/dailyai/pipeline/aggregators.py index f8816da44..7b7056869 100644 --- a/src/dailyai/pipeline/aggregators.py +++ b/src/dailyai/pipeline/aggregators.py @@ -1,11 +1,9 @@ import asyncio import re -from tblib import Frame from dailyai.pipeline.frame_processor import FrameProcessor from dailyai.pipeline.frames import ( - ControlFrame, EndPipeFrame, EndFrame, LLMMessagesQueueFrame, @@ -94,13 +92,6 @@ class LLMContextAggregator(AIService): self.messages.append({"role": self.role, "content": frame.text}) yield LLMMessagesQueueFrame(self.messages) - async def finalize(self) -> AsyncGenerator[Frame, None]: - # Send any dangling words that weren't finished with punctuation. - if self.complete_sentences and self.sentence: - self.messages.append({"role": self.role, "content": self.sentence}) - yield LLMMessagesQueueFrame(self.messages) - - class LLMUserContextAggregator(LLMContextAggregator): def __init__( self, messages: list[dict], bot_participant_id=None, complete_sentences=True @@ -124,7 +115,22 @@ class LLMAssistantContextAggregator(LLMContextAggregator): class SentenceAggregator(FrameProcessor): + """This frame processor aggregates text frames into complete sentences. + Frame input/output: + TextFrame("Hello,") -> None + TextFrame(" world.") -> TextFrame("Hello world.") + + Doctest: + >>> async def print_frames(aggregator, frame): + ... async for frame in aggregator.process_frame(frame): + ... print(frame.text) + + >>> aggregator = SentenceAggregator() + >>> asyncio.run(print_frames(aggregator, TextFrame("Hello,"))) + >>> asyncio.run(print_frames(aggregator, TextFrame(" world."))) + Hello, world. + """ def __init__(self): self.aggregation = "" @@ -147,6 +153,41 @@ class SentenceAggregator(FrameProcessor): class LLMFullResponseAggregator(FrameProcessor): + """This class aggregates Text frames until it receives a + LLMResponseEndFrame, then emits the concatenated text as + a single text frame. + + given the following frames: + + TextFrame("Hello,") + TextFrame(" world.") + TextFrame(" I am") + TextFrame(" an LLM.") + LLMResponseEndFrame()] + + this processor will yield nothing for the first 4 frames, then + + TextFrame("Hello, world. I am an LLM.") + LLMResponseEndFrame() + + when passed the last frame. + + >>> async def print_frames(aggregator, frame): + ... async for frame in aggregator.process_frame(frame): + ... if isinstance(frame, TextFrame): + ... print(frame.text) + ... else: + ... print(frame.__class__.__name__) + + >>> aggregator = LLMFullResponseAggregator() + >>> asyncio.run(print_frames(aggregator, TextFrame("Hello,"))) + >>> asyncio.run(print_frames(aggregator, TextFrame(" world."))) + >>> asyncio.run(print_frames(aggregator, TextFrame(" I am"))) + >>> asyncio.run(print_frames(aggregator, TextFrame(" an LLM."))) + >>> asyncio.run(print_frames(aggregator, LLMResponseEndFrame())) + Hello, world. I am an LLM. + LLMResponseEndFrame + """ def __init__(self): self.aggregation = "" @@ -157,12 +198,24 @@ class LLMFullResponseAggregator(FrameProcessor): self.aggregation += frame.text elif isinstance(frame, LLMResponseEndFrame): yield TextFrame(self.aggregation) + yield frame self.aggregation = "" else: yield frame class StatelessTextTransformer(FrameProcessor): + """This processor calls the given function on any text in a text frame. + + >>> async def print_frames(aggregator, frame): + ... async for frame in aggregator.process_frame(frame): + ... print(frame.text) + + >>> aggregator = StatelessTextTransformer(lambda x: x.upper()) + >>> asyncio.run(print_frames(aggregator, TextFrame("Hello"))) + HELLO + """ + def __init__(self, transform_fn): self.transform_fn = transform_fn diff --git a/src/dailyai/pipeline/frame_processor.py b/src/dailyai/pipeline/frame_processor.py index d70294ac0..3d361b4d9 100644 --- a/src/dailyai/pipeline/frame_processor.py +++ b/src/dailyai/pipeline/frame_processor.py @@ -27,12 +27,6 @@ class FrameProcessor: yield frame yield frame - @abstractmethod - async def finalize(self) -> AsyncGenerator[Frame, None]: - # This is a trick for the interpreter (and linter) to know that this is a generator. - if False: - yield Frame() - @abstractmethod async def interrupted(self) -> None: """Handle any cleanup if the pipeline was interrupted.""" diff --git a/src/dailyai/services/ai_services.py b/src/dailyai/services/ai_services.py index 0fca44884..6a692fb52 100644 --- a/src/dailyai/services/ai_services.py +++ b/src/dailyai/services/ai_services.py @@ -100,23 +100,14 @@ class TTSService(AIService): # yield empty bytes here, so linting can infer what this method does yield bytes() -<<<<<<< HEAD async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]: if isinstance(frame, EndFrame): if self.current_sentence: async for audio_chunk in self.run_tts(self.current_sentence): yield AudioFrame(audio_chunk) - yield frame + yield TextFrame(self.current_sentence) if not isinstance(frame, TextFrame): -======= - async def process_frame(self, frame: QueueFrame) -> AsyncGenerator[QueueFrame, None]: - if not isinstance(frame, TextQueueFrame): - if self.current_sentence: - async for audio_chunk in self.run_tts(self.current_sentence): - yield AudioQueueFrame(audio_chunk) - ->>>>>>> fa5f38c (frame and pipeline docstrings) yield frame return @@ -133,12 +124,9 @@ class TTSService(AIService): async for audio_chunk in self.run_tts(text): yield AudioFrame(audio_chunk) -<<<<<<< HEAD # note we pass along the text frame *after* the audio, so the text frame is completed after the audio is processed. yield TextFrame(text) -======= ->>>>>>> fa5f38c (frame and pipeline docstrings) # Convenience function to send the audio for a sentence to the given queue async def say(self, sentence, queue: asyncio.Queue): await self.run_to_queue(queue, [TextFrame(sentence)]) From b14f08a7d56c0fe21dcb004ba264a64a9ce20aae Mon Sep 17 00:00:00 2001 From: Moishe Lettvin Date: Thu, 7 Mar 2024 15:16:23 -0500 Subject: [PATCH 3/3] finish docstrings for aggregators --- src/dailyai/pipeline/aggregators.py | 48 +++++++++++++++++++++++++-- src/dailyai/tests/test_aggregators.py | 8 +++++ 2 files changed, 53 insertions(+), 3 deletions(-) diff --git a/src/dailyai/pipeline/aggregators.py b/src/dailyai/pipeline/aggregators.py index 7b7056869..d229b7178 100644 --- a/src/dailyai/pipeline/aggregators.py +++ b/src/dailyai/pipeline/aggregators.py @@ -4,11 +4,12 @@ import re from dailyai.pipeline.frame_processor import FrameProcessor from dailyai.pipeline.frames import ( - EndPipeFrame, EndFrame, + EndPipeFrame, + Frame, + ImageFrame, LLMMessagesQueueFrame, LLMResponseEndFrame, - Frame, LLMResponseStartFrame, TextFrame, TranscriptionQueueFrame, @@ -16,7 +17,7 @@ from dailyai.pipeline.frames import ( from dailyai.pipeline.pipeline import Pipeline from dailyai.services.ai_services import AIService -from typing import AsyncGenerator, Coroutine, List, Text +from typing import AsyncGenerator, Coroutine, List class LLMResponseAggregator(FrameProcessor): def __init__(self, messages: list[dict]): @@ -230,6 +231,23 @@ class StatelessTextTransformer(FrameProcessor): yield frame class ParallelPipeline(FrameProcessor): + """ Run multiple pipelines in parallel. + + This class takes frames from its source queue and sends them to each + sub-pipeline. Each sub-pipeline emits its frames into this class's + sink queue. No guarantees are made about the ordering of frames in + the sink queue (that is, no sub-pipeline has higher priority than + any other, frames are put on the sink in the order they're emitted + by the sub-pipelines). + + After each frame is taken from this class's source queue and placed + in each sub-pipeline's source queue, an EndPipeFrame is put on each + sub-pipeline's source queue. This indicates to the sub-pipe runner + that it should exit. + + Since frame handlers pass through unhandled frames by convention, this + class de-dupes frames in its sink before yielding them. + """ def __init__(self, pipeline_definitions: List[List[FrameProcessor]]): self.sources = [asyncio.Queue() for _ in pipeline_definitions] self.sink: asyncio.Queue[Frame] = asyncio.Queue() @@ -266,6 +284,30 @@ class ParallelPipeline(FrameProcessor): yield frame class GatedAggregator(FrameProcessor): + """Accumulate frames, with custom functions to start and stop accumulation. + Yields gate-opening frame before any accumulated frames, then ensuing frames + until and not including the gate-closed frame. + + >>> async def print_frames(aggregator, frame): + ... async for frame in aggregator.process_frame(frame): + ... if isinstance(frame, TextFrame): + ... print(frame.text) + ... else: + ... print(frame.__class__.__name__) + + >>> aggregator = GatedAggregator( + ... gate_close_fn=lambda x: isinstance(x, LLMResponseStartFrame), + ... gate_open_fn=lambda x: isinstance(x, ImageFrame), + ... start_open=False) + >>> asyncio.run(print_frames(aggregator, TextFrame("Hello"))) + >>> asyncio.run(print_frames(aggregator, TextFrame("Hello again."))) + >>> asyncio.run(print_frames(aggregator, ImageFrame(url='', image=bytes([])))) + ImageFrame + Hello + Hello again. + >>> asyncio.run(print_frames(aggregator, TextFrame("Goodbye."))) + Goodbye. + """ def __init__(self, gate_open_fn, gate_close_fn, start_open): self.gate_open_fn = gate_open_fn self.gate_close_fn = gate_close_fn diff --git a/src/dailyai/tests/test_aggregators.py b/src/dailyai/tests/test_aggregators.py index bad689e73..800066969 100644 --- a/src/dailyai/tests/test_aggregators.py +++ b/src/dailyai/tests/test_aggregators.py @@ -1,4 +1,5 @@ import asyncio +import doctest import functools import unittest @@ -118,3 +119,10 @@ class TestDailyFrameAggregators(unittest.IsolatedAsyncioTestCase): while not sink.empty(): frame = await sink.get() self.assertEqual(frame, expected_output_frames.pop(0)) + + +def load_tests(loader, tests, ignore): + """ Run doctests on the aggregators module. """ + from dailyai.pipeline import aggregators + tests.addTests(doctest.DocTestSuite(aggregators)) + return tests