wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #835 from pipecat-ai/aleix/even-more-parallel-pipeline-fixes

Browse Source 
parallel_pipeline: fix system frames and parallel pipelines again
This commit is contained in:
Aleix Conchillo Flaqué
2024-12-11 11:15:59 -08:00
committed by GitHub
parent fdefdc9d68 2dd418a38d
commit cff9bb6068
1 changed files with 31 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
src/pipecat/pipeline/parallel_pipeline.py
View File

@@ -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()