diff --git a/src/pipecat/pipeline/parallel_pipeline.py b/src/pipecat/pipeline/parallel_pipeline.py index 48ec1a172..40bfea90d 100644 --- a/src/pipecat/pipeline/parallel_pipeline.py +++ b/src/pipecat/pipeline/parallel_pipeline.py @@ -7,7 +7,7 @@ import asyncio from itertools import chain -from typing import List +from typing import Awaitable, Callable, List from pipecat.pipeline.base_pipeline import BasePipeline from pipecat.pipeline.pipeline import Pipeline @@ -18,26 +18,37 @@ from loguru import logger class Source(FrameProcessor): - def __init__(self, upstream_queue: asyncio.Queue): + def __init__( + self, + upstream_queue: asyncio.Queue, + push_frame_func: Callable[[Frame, FrameDirection], Awaitable[None]], + ): super().__init__() self._up_queue = upstream_queue + self._push_frame_func = push_frame_func async def process_frame(self, frame: Frame, direction: FrameDirection): await super().process_frame(frame, direction) match direction: case FrameDirection.UPSTREAM: - # SystemFrames are pushed directly from ParallelPipeline. - if not isinstance(frame, SystemFrame): + if isinstance(frame, SystemFrame): + await self._push_frame_func(frame, direction) + else: await self._up_queue.put(frame) case FrameDirection.DOWNSTREAM: await self.push_frame(frame, direction) class Sink(FrameProcessor): - def __init__(self, downstream_queue: asyncio.Queue): + def __init__( + self, + downstream_queue: asyncio.Queue, + push_frame_func: Callable[[Frame, FrameDirection], Awaitable[None]], + ): super().__init__() self._down_queue = downstream_queue + self._push_frame_func = push_frame_func async def process_frame(self, frame: Frame, direction: FrameDirection): await super().process_frame(frame, direction) @@ -46,8 +57,9 @@ class Sink(FrameProcessor): case FrameDirection.UPSTREAM: await self.push_frame(frame, direction) case FrameDirection.DOWNSTREAM: - # SystemFrames are pushed directly from ParallelPipeline. - if not isinstance(frame, SystemFrame): + if isinstance(frame, SystemFrame): + await self._push_frame_func(frame, direction) + else: await self._down_queue.put(frame) @@ -60,6 +72,7 @@ class ParallelPipeline(BasePipeline): self._sources = [] self._sinks = [] + self._seen_ids = set() self._up_queue = asyncio.Queue() self._down_queue = asyncio.Queue() @@ -74,8 +87,8 @@ class ParallelPipeline(BasePipeline): raise TypeError(f"ParallelPipeline argument {processors} is not a list") # We will add a source before the pipeline and a sink after. - source = Source(self._up_queue) - sink = Sink(self._down_queue) + source = Source(self._up_queue, self._parallel_push_frame) + sink = Sink(self._down_queue, self._parallel_push_frame) self._sources.append(source) self._sinks.append(sink) @@ -119,11 +132,6 @@ class ParallelPipeline(BasePipeline): # If we get a downstream frame we process it in each source. await asyncio.gather(*[s.queue_frame(frame, direction) for s in self._sources]) - # If we have a SystemFrame we will push it from this task. Note that the - # connected sinks and sources ignore SystemFrames. - if isinstance(frame, SystemFrame): - await self.push_frame(frame, direction) - # If we get an EndFrame we stop our queue processing tasks and wait on # all the pipelines to finish. if isinstance(frame, (CancelFrame, EndFrame)): @@ -135,24 +143,25 @@ class ParallelPipeline(BasePipeline): if self._down_task: await self._down_task + async def _parallel_push_frame(self, frame: Frame, direction: FrameDirection): + if frame.id not in self._seen_ids: + self._seen_ids.add(frame.id) + await self.push_frame(frame, direction) + async def _process_up_queue(self): running = True - seen_ids = set() while running: frame = await self._up_queue.get() - if frame and frame.id not in seen_ids: - await self.push_frame(frame, FrameDirection.UPSTREAM) - seen_ids.add(frame.id) + if frame: + await self._parallel_push_frame(frame, FrameDirection.UPSTREAM) running = frame is not None self._up_queue.task_done() async def _process_down_queue(self): running = True - seen_ids = set() while running: frame = await self._down_queue.get() - if frame and frame.id not in seen_ids: - await self.push_frame(frame, FrameDirection.DOWNSTREAM) - seen_ids.add(frame.id) + if frame: + await self._parallel_push_frame(frame, FrameDirection.DOWNSTREAM) running = frame is not None self._down_queue.task_done()