111 lines
4.6 KiB
Python
111 lines
4.6 KiB
Python
import asyncio
|
|
from typing import AsyncGenerator, AsyncIterable, Iterable, 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 we create the sink and source queues
|
|
if they're not provided, but these can be overridden to point to other
|
|
queues. If this pipeline is run by a transport, its sink and source queues
|
|
will be overridden.
|
|
"""
|
|
self.processors: List[FrameProcessor] = processors
|
|
|
|
self.source: asyncio.Queue[Frame] = source or asyncio.Queue()
|
|
self.sink: asyncio.Queue[Frame] = sink or asyncio.Queue()
|
|
|
|
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."""
|
|
|
|
yield await self.source.get()
|
|
|
|
async def queue_frames(
|
|
self,
|
|
frames: Iterable[Frame] | AsyncIterable[Frame],
|
|
) -> None:
|
|
"""Insert frames directly into a pipeline. This is typically used inside a transport
|
|
participant_joined callback to prompt a bot to start a conversation, for example."""
|
|
|
|
if isinstance(frames, AsyncIterable):
|
|
async for frame in frames:
|
|
await self.source.put(frame)
|
|
elif isinstance(frames, Iterable):
|
|
for frame in frames:
|
|
await self.source.put(frame)
|
|
else:
|
|
raise Exception("Frames must be an iterable or async iterable")
|
|
|
|
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..
|
|
"""
|
|
|
|
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
|
|
|
|
async def _run_pipeline_recursively(
|
|
self, initial_frame: Frame, processors: List[FrameProcessor]
|
|
) -> AsyncGenerator[Frame, None]:
|
|
"""Internal function to add frames to the pipeline as they're yielded
|
|
by each processor."""
|
|
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
|
|
|