89 lines
3.9 KiB
Python
89 lines
3.9 KiB
Python
import asyncio
|
|
from typing import AsyncGenerator, List
|
|
from dailyai.pipeline.frame_processor import FrameProcessor
|
|
|
|
from dailyai.pipeline.frames import EndPipeFrame, EndFrame, Frame
|
|
|
|
|
|
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,
|
|
processors: List[FrameProcessor],
|
|
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_recursively(
|
|
self, initial_frame: Frame, processors: List[FrameProcessor]
|
|
) -> AsyncGenerator[Frame, None]:
|
|
if processors:
|
|
async for frame in processors[0].process_frame(initial_frame):
|
|
async for final_frame in self.run_pipeline_recursively(frame, processors[1:]):
|
|
yield final_frame
|
|
else:
|
|
yield initial_frame
|
|
|
|
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")
|
|
|
|
try:
|
|
while True:
|
|
initial_frame = await self.source.get()
|
|
async for frame in self.run_pipeline_recursively(initial_frame, self.processors):
|
|
await self.sink.put(frame)
|
|
|
|
if isinstance(initial_frame, EndFrame) or isinstance(initial_frame, EndPipeFrame):
|
|
break
|
|
except asyncio.CancelledError:
|
|
# this means there's been an interruption, do any cleanup necessary here.
|
|
for processor in self.processors:
|
|
await processor.interrupted()
|
|
pass
|