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)])