From 79f29e14dddf6f70095ae24d50b2f4cca4b5378c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Wed, 2 Apr 2025 19:28:52 -0700 Subject: [PATCH] processors: add ProducerProcessor and ConsumerProcessor --- CHANGELOG.md | 9 ++ src/pipecat/processors/consumer_processor.py | 65 ++++++++++ src/pipecat/processors/producer_processor.py | 73 +++++++++++ tests/test_producer_consumer.py | 120 +++++++++++++++++++ 4 files changed, 267 insertions(+) create mode 100644 src/pipecat/processors/consumer_processor.py create mode 100644 src/pipecat/processors/producer_processor.py create mode 100644 tests/test_producer_consumer.py diff --git a/CHANGELOG.md b/CHANGELOG.md index aec1ec0fb..8b2cf3958 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Added new processors `ProducerProcessor` and `ConsumerProcessor`. The producer + processor processes frames from the pipeline and decides whether the consumers + should consume it or not. If so, the same frame that is received by the + producer is sent to the consumer. There can be multiple consumers per + producer. These processors can be useful to push frames from one part of a + pipeline to a different one (e.g. when using `ParallelPipeline`). + ### Fixed - Fixed an issue where `LLMAssistantContextAggregator` would prevent a diff --git a/src/pipecat/processors/consumer_processor.py b/src/pipecat/processors/consumer_processor.py new file mode 100644 index 000000000..dbdfc97e5 --- /dev/null +++ b/src/pipecat/processors/consumer_processor.py @@ -0,0 +1,65 @@ +# +# Copyright (c) 2024–2025, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +from typing import Awaitable, Callable, Optional + +from pipecat.frames.frames import CancelFrame, EndFrame, Frame, StartFrame +from pipecat.processors.frame_processor import FrameDirection, FrameProcessor +from pipecat.processors.producer_processor import ProducerProcessor, identity_transformer + + +class ConsumerProcessor(FrameProcessor): + """This class passes-through frames and also consumes frames from a + producer's queue. When a frame from a producer queue is received it will be + pushed to the specified direction. The frames can be transformed into a + different type of frame before being pushed. + + """ + + def __init__( + self, + *, + producer: ProducerProcessor, + transformer: Callable[[Frame], Awaitable[Frame]] = identity_transformer, + direction: FrameDirection = FrameDirection.DOWNSTREAM, + **kwargs, + ): + super().__init__(**kwargs) + self._transformer = transformer + self._direction = direction + self._queue: asyncio.Queue = producer.add_consumer() + self._consumer_task: Optional[asyncio.Task] = None + + async def process_frame(self, frame: Frame, direction: FrameDirection): + await super().process_frame(frame, direction) + + if isinstance(frame, StartFrame): + await self._start(frame) + elif isinstance(frame, EndFrame): + await self._stop(frame) + elif isinstance(frame, CancelFrame): + await self._cancel(frame) + + await self.push_frame(frame, direction) + + async def _start(self, _: StartFrame): + if not self._consumer_task: + self._consumer_task = self.create_task(self._consumer_task_handler()) + + async def _stop(self, _: EndFrame): + if self._consumer_task: + await self.cancel_task(self._consumer_task) + + async def _cancel(self, _: CancelFrame): + if self._consumer_task: + await self.cancel_task(self._consumer_task) + + async def _consumer_task_handler(self): + while True: + frame = await self._queue.get() + new_frame = await self._transformer(frame) + await self.push_frame(new_frame, self._direction) diff --git a/src/pipecat/processors/producer_processor.py b/src/pipecat/processors/producer_processor.py new file mode 100644 index 000000000..6ada2ed83 --- /dev/null +++ b/src/pipecat/processors/producer_processor.py @@ -0,0 +1,73 @@ +# +# Copyright (c) 2024–2025, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +from typing import Awaitable, Callable, List + +from pipecat.frames.frames import Frame +from pipecat.processors.frame_processor import FrameDirection, FrameProcessor + + +async def identity_transformer(frame: Frame): + return frame + + +class ProducerProcessor(FrameProcessor): + """This class optionally passes-through received frames and decides if those + frames should be sent to consumers based on a user-defined filter. The + frames can be transformed into a different type of frame before being + sending them to the consumers. More than one consumer can be added. + + """ + + def __init__( + self, + *, + filter: Callable[[Frame], Awaitable[bool]], + transformer: Callable[[Frame], Awaitable[Frame]] = identity_transformer, + passthrough: bool = True, + ): + super().__init__() + self._filter = filter + self._transformer = transformer + self._passthrough = passthrough + self._consumers: List[asyncio.Queue] = [] + + def add_consumer(self): + """ + Adds a new consumer and returns its associated queue. + + Returns: + asyncio.Queue: The queue for the newly added consumer. + """ + queue = asyncio.Queue() + self._consumers.append(queue) + return queue + + async def process_frame(self, frame: Frame, direction: FrameDirection): + """ + Processes an incoming frame and determines whether to produce it as a ProducerItem. + + If the frame meets the produce criteria, it will be added to the consumer queues. + If passthrough is enabled, the frame will also be sent to consumers. + + Args: + frame (Frame): The frame to process. + direction (FrameDirection): The direction of the frame. + """ + await super().process_frame(frame, direction) + + if await self._filter(frame): + await self._produce(frame) + if self._passthrough: + await self.push_frame(frame, direction) + else: + await self.push_frame(frame, direction) + + async def _produce(self, frame: Frame): + for consumer in self._consumers: + new_frame = await self._transformer(frame) + await consumer.put(new_frame) diff --git a/tests/test_producer_consumer.py b/tests/test_producer_consumer.py new file mode 100644 index 000000000..578d14706 --- /dev/null +++ b/tests/test_producer_consumer.py @@ -0,0 +1,120 @@ +# +# Copyright (c) 2024-2025 Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import unittest + +from pipecat.frames.frames import Frame, InputAudioRawFrame, TextFrame +from pipecat.pipeline.parallel_pipeline import ParallelPipeline +from pipecat.pipeline.pipeline import Pipeline +from pipecat.processors.consumer_processor import ConsumerProcessor +from pipecat.processors.producer_processor import ProducerProcessor +from pipecat.tests.utils import SleepFrame, run_test + + +async def text_frame_filter(frame: Frame): + return isinstance(frame, TextFrame) + + +class TestProducerConsumerProcessor(unittest.IsolatedAsyncioTestCase): + async def test_produce_passthrough(self): + producer = ProducerProcessor(filter=text_frame_filter) + consumer = ConsumerProcessor(producer=producer) + pipeline = Pipeline([producer, consumer]) + frames_to_send = [ + TextFrame("Hello!"), + SleepFrame(), # So we let the consumer go first. + ] + expected_down_frames = [ + TextFrame, # Consumer frame + TextFrame, # Pass-through frame + ] + await run_test( + pipeline, + frames_to_send=frames_to_send, + expected_down_frames=expected_down_frames, + ) + + async def test_produce_no_passthrough(self): + producer = ProducerProcessor(filter=text_frame_filter, passthrough=False) + consumer = ConsumerProcessor(producer=producer) + pipeline = Pipeline([producer, consumer]) + frames_to_send = [TextFrame("Hello!")] + expected_down_frames = [TextFrame] + await run_test( + pipeline, + frames_to_send=frames_to_send, + expected_down_frames=expected_down_frames, + ) + + async def test_produce_multiple_consumer_no_passthrough(self): + producer = ProducerProcessor(filter=text_frame_filter, passthrough=False) + consumer1 = ConsumerProcessor(producer=producer) + consumer2 = ConsumerProcessor(producer=producer) + pipeline = Pipeline([producer, consumer1, consumer2]) + frames_to_send = [TextFrame("Hello!")] + expected_down_frames = [ + TextFrame, # From consumer1 or consumer2 (depending on who runs first) + TextFrame, # From consumer1 or consumer2 (depending on who runs first) + ] + await run_test( + pipeline, + frames_to_send=frames_to_send, + expected_down_frames=expected_down_frames, + ) + + async def test_produce_parallel_pipeline_no_passthrough(self): + producer = ProducerProcessor(filter=text_frame_filter, passthrough=False) + consumer = ConsumerProcessor(producer=producer) + pipeline = Pipeline([ParallelPipeline([producer], [consumer])]) + frames_to_send = [TextFrame("Hello!")] + expected_down_frames = [TextFrame] + await run_test( + pipeline, + frames_to_send=frames_to_send, + expected_down_frames=expected_down_frames, + ) + + async def test_produce_passthrough_transform(self): + async def audio_transformer(_: Frame) -> Frame: + return InputAudioRawFrame(audio=b"", sample_rate=16000, num_channels=1) + + producer = ProducerProcessor(filter=text_frame_filter, transformer=audio_transformer) + consumer = ConsumerProcessor(producer=producer) + pipeline = Pipeline([producer, consumer]) + frames_to_send = [ + TextFrame("Hello!"), + SleepFrame(), # So we let the consumer go first. + ] + expected_down_frames = [ + InputAudioRawFrame, # Consumer frame + TextFrame, # Pass-through frame + ] + await run_test( + pipeline, + frames_to_send=frames_to_send, + expected_down_frames=expected_down_frames, + ) + + async def test_produce_passthrough_consumer_transform(self): + async def audio_transformer(_: Frame) -> Frame: + return InputAudioRawFrame(audio=b"", sample_rate=16000, num_channels=1) + + producer = ProducerProcessor(filter=text_frame_filter) + consumer = ConsumerProcessor(producer=producer, transformer=audio_transformer) + pipeline = Pipeline([producer, consumer]) + frames_to_send = [ + TextFrame("Hello!"), + SleepFrame(), # So we let the consumer go first. + ] + expected_down_frames = [ + InputAudioRawFrame, # Consumer frame + TextFrame, # Pass-through frame + ] + await run_test( + pipeline, + frames_to_send=frames_to_send, + expected_down_frames=expected_down_frames, + )