Merge pull request #1139 from pipecat-ai/aleix/task-start-metadata

pipeline task start metadata and unit test improvements

CHANGELOG.md

@@ -9,6 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added a new `start_metadata` field to `PipelineParams`. The provided metadata
+  will be set to the initial `StartFrame` being pushed from the `PipelineTask`.
+
 - Added new fields to `PipelineParams` to control audio input and output sample
   rates for the whole pipeline. This allows controlling sample rates from a
   single place instead of having to specify sample rates in each
@@ -107,6 +110,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Other
 
+- Improved Unit Test `run_test()` to use `PipelineTask` and
+  `PipelineRunner`. There's now also some control around `StartFrame` and
+  `EndFrame`. The `EndTaskFrame` has been removed since it doesn't seem
+  necessary with this new approach.
+
 - Updated `twilio-chatbot` with a few new features: use 8000 sample rate and
   avoid resampling, a new client useful for stress testing and testing locally
   without the need to make phone calls. Also, added audio recording on both the

dev-requirements.txt

@@ -1,11 +1,11 @@
 build~=1.2.2
-grpcio-tools~=1.69.0
+grpcio-tools~=1.67.1
 pip-tools~=7.4.1
 pre-commit~=4.0.1
 pyright~=1.1.392
 pytest~=8.3.4
 pytest-asyncio~=0.25.2
 ruff~=0.9.1
-setuptools~=75.8.0
+setuptools~=70.0.0
 setuptools_scm~=8.1.0
 python-dotenv~=1.0.1

src/pipecat/frames/frames.py

@@ -6,7 +6,18 @@
 
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Awaitable, Callable, List, Literal, Mapping, Optional, Tuple
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Tuple,
+)
 
 from pipecat.audio.vad.vad_analyzer import VADParams
 from pipecat.clocks.base_clock import BaseClock
@@ -48,13 +59,13 @@ class Frame:
     id: int = field(init=False)
     name: str = field(init=False)
     pts: Optional[int] = field(init=False)
-    metadata: dict = field(init=False)
+    metadata: Dict[str, Any] = field(init=False)
 
     def __post_init__(self):
         self.id: int = obj_id()
         self.name: str = f"{self.__class__.__name__}#{obj_count(self)}"
         self.pts: Optional[int] = None
-        self.metadata: dict = {}
+        self.metadata: Dict[str, Any] = {}
 
     def __str__(self):
         return self.name
@@ -433,8 +444,8 @@ class StartFrame(SystemFrame):
     allow_interruptions: bool = False
     enable_metrics: bool = False
     enable_usage_metrics: bool = False
-    report_only_initial_ttfb: bool = False
     observer: Optional["BaseObserver"] = None
+    report_only_initial_ttfb: bool = False
 
 
 @dataclass

src/pipecat/pipeline/task.py

@@ -5,7 +5,7 @@
 #
 
 import asyncio
-from typing import AsyncIterable, Iterable, List
+from typing import Any, AsyncIterable, Dict, Iterable, List
 
 from loguru import logger
 from pydantic import BaseModel, ConfigDict
@@ -40,16 +40,17 @@ HEARTBEAT_MONITOR_SECONDS = HEARTBEAT_SECONDS * 5
 class PipelineParams(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
+    allow_interruptions: bool = False
     audio_in_sample_rate: int = 16000
     audio_out_sample_rate: int = 24000
-    allow_interruptions: bool = False
     enable_heartbeats: bool = False
     enable_metrics: bool = False
     enable_usage_metrics: bool = False
-    send_initial_empty_metrics: bool = True
-    report_only_initial_ttfb: bool = False
-    observers: List[BaseObserver] = []
     heartbeats_period_secs: float = HEARTBEAT_SECONDS
+    observers: List[BaseObserver] = []
+    report_only_initial_ttfb: bool = False
+    send_initial_empty_metrics: bool = True
+    start_metadata: Dict[str, Any] = {}
 
 
 class PipelineTaskSource(FrameProcessor):
@@ -278,13 +279,14 @@ class PipelineTask(BaseTask):
             clock=self._clock,
             task_manager=self._task_manager,
             allow_interruptions=self._params.allow_interruptions,
-            enable_metrics=self._params.enable_metrics,
-            enable_usage_metrics=self._params.enable_usage_metrics,
-            report_only_initial_ttfb=self._params.report_only_initial_ttfb,
-            observer=self._observer,
             audio_in_sample_rate=self._params.audio_in_sample_rate,
             audio_out_sample_rate=self._params.audio_out_sample_rate,
+            enable_metrics=self._params.enable_metrics,
+            enable_usage_metrics=self._params.enable_usage_metrics,
+            observer=self._observer,
+            report_only_initial_ttfb=self._params.report_only_initial_ttfb,
         )
+        start_frame.metadata = self._params.start_metadata
         await self._source.queue_frame(start_frame, FrameDirection.DOWNSTREAM)
 
         if self._params.enable_metrics and self._params.send_initial_empty_metrics:

src/pipecat/tests/utils.py

@@ -5,24 +5,25 @@
 #
 
 import asyncio
-from dataclasses import dataclass
+import sys
-from typing import Awaitable, Callable, Sequence, Tuple
+from typing import Any, Awaitable, Callable, Dict, Sequence, Tuple
+
+from loguru import logger
 
-from pipecat.clocks.system_clock import SystemClock
 from pipecat.frames.frames import (
-    ControlFrame,
+    EndFrame,
     Frame,
     HeartbeatFrame,
     StartFrame,
 )
 from pipecat.observers.base_observer import BaseObserver
+from pipecat.pipeline.pipeline import Pipeline
+from pipecat.pipeline.runner import PipelineRunner
+from pipecat.pipeline.task import PipelineParams, PipelineTask
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
-from pipecat.utils.asyncio import TaskManager
 
-
+logger.remove(0)
-@dataclass
+logger.add(sys.stderr, level="TRACE")
-class EndTestFrame(ControlFrame):
-    pass
 
 
 class HeartbeatsObserver(BaseObserver):
@@ -48,54 +49,58 @@ class HeartbeatsObserver(BaseObserver):
 
 
 class QueuedFrameProcessor(FrameProcessor):
-    def __init__(self, queue: asyncio.Queue, ignore_start: bool = True):
+    def __init__(
+        self, queue: asyncio.Queue, queue_direction: FrameDirection, ignore_start: bool = True
+    ):
         super().__init__()
         self._queue = queue
+        self._queue_direction = queue_direction
         self._ignore_start = ignore_start
 
     async def process_frame(self, frame: Frame, direction: FrameDirection):
         await super().process_frame(frame, direction)
 
-        if self._ignore_start and isinstance(frame, StartFrame):
+        if direction == self._queue_direction:
-            await self.push_frame(frame, direction)
+            if not isinstance(frame, StartFrame) or not self._ignore_start:
-        else:
+                await self._queue.put(frame)
-            await self._queue.put(frame)
+        await self.push_frame(frame, direction)
-            await self.push_frame(frame, direction)
 
 
 async def run_test(
     processor: FrameProcessor,
+    *,
     frames_to_send: Sequence[Frame],
     expected_down_frames: Sequence[type],
     expected_up_frames: Sequence[type] = [],
+    ignore_start: bool = True,
+    start_metadata: Dict[str, Any] = {},
+    send_end_frame: bool = True,
 ) -> Tuple[Sequence[Frame], Sequence[Frame]]:
     received_up = asyncio.Queue()
     received_down = asyncio.Queue()
-    source = QueuedFrameProcessor(received_up)
+    source = QueuedFrameProcessor(received_up, FrameDirection.UPSTREAM, ignore_start)
-    sink = QueuedFrameProcessor(received_down)
+    sink = QueuedFrameProcessor(received_down, FrameDirection.DOWNSTREAM, ignore_start)
 
-    source.link(processor)
+    pipeline = Pipeline([source, processor, sink])
-    processor.link(sink)
 
-    task_manager = TaskManager()
+    task = PipelineTask(pipeline, params=PipelineParams(start_metadata=start_metadata))
-    task_manager.set_event_loop(asyncio.get_event_loop())
-    await source.queue_frame(StartFrame(clock=SystemClock(), task_manager=task_manager))
 
     for frame in frames_to_send:
-        await processor.process_frame(frame, FrameDirection.DOWNSTREAM)
+        await task.queue_frame(frame)
 
-    await processor.queue_frame(EndTestFrame())
+    if send_end_frame:
-    await processor.queue_frame(EndTestFrame(), FrameDirection.UPSTREAM)
+        await task.queue_frame(EndFrame())
+
+    runner = PipelineRunner()
+    await runner.run(task)
 
     #
     # Down frames
     #
     received_down_frames: Sequence[Frame] = []
-    running = True
+    while not received_down.empty():
-    while running:
         frame = await received_down.get()
-        running = not isinstance(frame, EndTestFrame)
+        if not isinstance(frame, EndFrame) or not send_end_frame:
-        if running:
             received_down_frames.append(frame)
 
     print("received DOWN frames =", received_down_frames)
@@ -109,12 +114,9 @@ async def run_test(
     # Up frames
     #
     received_up_frames: Sequence[Frame] = []
-    running = True
+    while not received_up.empty():
-    while running:
         frame = await received_up.get()
-        running = not isinstance(frame, EndTestFrame)
+        received_up_frames.append(frame)
-        if running:
-            received_up_frames.append(frame)
 
     print("received UP frames =", received_up_frames)
 

tests/test_aggregators.py

@@ -31,7 +31,11 @@ class TestSentenceAggregator(unittest.IsolatedAsyncioTestCase):
 
         expected_returned_frames = [TextFrame, TextFrame, TextFrame]
 
-        (received_down, _) = await run_test(aggregator, frames_to_send, expected_returned_frames)
+        (received_down, _) = await run_test(
+            aggregator,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_returned_frames,
+        )
         assert received_down[-3].text == "Hello, world. "
         assert received_down[-2].text == "How are you? "
         assert received_down[-1].text == "I am fine! "
@@ -66,5 +70,7 @@ class TestGatedAggregator(unittest.IsolatedAsyncioTestCase):
         ]
 
         (received_down, _) = await run_test(
-            gated_aggregator, frames_to_send, expected_returned_frames
+            gated_aggregator,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_returned_frames,
         )

tests/test_filters.py

@@ -4,7 +4,6 @@
 # SPDX-License-Identifier: BSD 2-Clause License
 #
 
-import asyncio
 import unittest
 
 from pipecat.frames.frames import (
@@ -19,7 +18,7 @@ from pipecat.processors.filters.frame_filter import FrameFilter
 from pipecat.processors.filters.function_filter import FunctionFilter
 from pipecat.processors.filters.identity_filter import IdentityFilter
 from pipecat.processors.filters.wake_check_filter import WakeCheckFilter
-from pipecat.tests.utils import EndTestFrame, run_test
+from pipecat.tests.utils import run_test
 
 
 class TestIdentifyFilter(unittest.IsolatedAsyncioTestCase):
@@ -27,27 +26,44 @@ class TestIdentifyFilter(unittest.IsolatedAsyncioTestCase):
         filter = IdentityFilter()
         frames_to_send = [UserStartedSpeakingFrame(), UserStoppedSpeakingFrame()]
         expected_returned_frames = [UserStartedSpeakingFrame, UserStoppedSpeakingFrame]
-        await run_test(filter, frames_to_send, expected_returned_frames)
+        await run_test(
+            filter,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_returned_frames,
+        )
 
 
 class TestFrameFilter(unittest.IsolatedAsyncioTestCase):
     async def test_text_frame(self):
-        filter = FrameFilter(types=(TextFrame, EndTestFrame))
+        filter = FrameFilter(types=(TextFrame,))
         frames_to_send = [TextFrame(text="Hello Pipecat!")]
         expected_returned_frames = [TextFrame]
-        await run_test(filter, frames_to_send, expected_returned_frames)
+        await run_test(
+            filter,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_returned_frames,
+        )
 
     async def test_end_frame(self):
-        filter = FrameFilter(types=(EndFrame, EndTestFrame))
+        filter = FrameFilter(types=(EndFrame,))
         frames_to_send = [EndFrame()]
         expected_returned_frames = [EndFrame]
-        await run_test(filter, frames_to_send, expected_returned_frames)
+        await run_test(
+            filter,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_returned_frames,
+            send_end_frame=False,
+        )
 
     async def test_system_frame(self):
-        filter = FrameFilter(types=(EndTestFrame,))
+        filter = FrameFilter(types=())
         frames_to_send = [UserStartedSpeakingFrame()]
         expected_returned_frames = [UserStartedSpeakingFrame]
-        await run_test(filter, frames_to_send, expected_returned_frames)
+        await run_test(
+            filter,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_returned_frames,
+        )
 
 
 class TestFunctionFilter(unittest.IsolatedAsyncioTestCase):
@@ -58,7 +74,11 @@ class TestFunctionFilter(unittest.IsolatedAsyncioTestCase):
         filter = FunctionFilter(filter=passthrough)
         frames_to_send = [TextFrame(text="Hello Pipecat!")]
         expected_returned_frames = [TextFrame]
-        await run_test(filter, frames_to_send, expected_returned_frames)
+        await run_test(
+            filter,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_returned_frames,
+        )
 
     async def test_no_passthrough(self):
         async def no_passthrough(frame: Frame):
@@ -66,14 +86,12 @@ class TestFunctionFilter(unittest.IsolatedAsyncioTestCase):
 
         filter = FunctionFilter(filter=no_passthrough)
         frames_to_send = [TextFrame(text="Hello Pipecat!")]
-        expected_returned_frames = [TextFrame]
+        expected_returned_frames = []
-        try:
+        await run_test(
-            await asyncio.wait_for(
+            filter,
-                run_test(filter, frames_to_send, expected_returned_frames), timeout=0.5
+            frames_to_send=frames_to_send,
-            )
+            expected_down_frames=expected_returned_frames,
-            assert False
+        )
-        except asyncio.TimeoutError:
-            pass
 
 
 class TestWakeCheckFilter(unittest.IsolatedAsyncioTestCase):
@@ -81,7 +99,11 @@ class TestWakeCheckFilter(unittest.IsolatedAsyncioTestCase):
         filter = WakeCheckFilter(wake_phrases=["Hey, Pipecat"])
         frames_to_send = [TranscriptionFrame(user_id="test", text="Phrase 1", timestamp="")]
         expected_returned_frames = []
-        await run_test(filter, frames_to_send, expected_returned_frames)
+        await run_test(
+            filter,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_returned_frames,
+        )
 
     async def test_wake_word(self):
         filter = WakeCheckFilter(wake_phrases=["Hey, Pipecat"])
@@ -90,5 +112,9 @@ class TestWakeCheckFilter(unittest.IsolatedAsyncioTestCase):
             TranscriptionFrame(user_id="test", text="Phrase 1", timestamp=""),
         ]
         expected_returned_frames = [TranscriptionFrame, TranscriptionFrame]
-        (received_down, _) = await run_test(filter, frames_to_send, expected_returned_frames)
+        (received_down, _) = await run_test(
+            filter,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_returned_frames,
+        )
         assert received_down[-1].text == "Phrase 1"

tests/test_pipeline.py

@@ -7,7 +7,7 @@
 import asyncio
 import unittest
 
-from pipecat.frames.frames import EndFrame, HeartbeatFrame, TextFrame
+from pipecat.frames.frames import EndFrame, HeartbeatFrame, StartFrame, TextFrame
 from pipecat.pipeline.parallel_pipeline import ParallelPipeline
 from pipecat.pipeline.pipeline import Pipeline
 from pipecat.pipeline.task import PipelineParams, PipelineTask
@@ -22,7 +22,11 @@ class TestPipeline(unittest.IsolatedAsyncioTestCase):
 
         frames_to_send = [TextFrame(text="Hello from Pipecat!")]
         expected_returned_frames = [TextFrame]
-        await run_test(pipeline, frames_to_send, expected_returned_frames)
+        await run_test(
+            pipeline,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_returned_frames,
+        )
 
     async def test_pipeline_multiple(self):
         identity1 = IdentityFilter()
@@ -33,7 +37,25 @@ class TestPipeline(unittest.IsolatedAsyncioTestCase):
 
         frames_to_send = [TextFrame(text="Hello from Pipecat!")]
         expected_returned_frames = [TextFrame]
-        await run_test(pipeline, frames_to_send, expected_returned_frames)
+        await run_test(
+            pipeline,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_returned_frames,
+        )
+
+    async def test_pipeline_start_metadata(self):
+        pipeline = Pipeline([IdentityFilter()])
+
+        frames_to_send = []
+        expected_returned_frames = [StartFrame]
+        (received_down, _) = await run_test(
+            pipeline,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_returned_frames,
+            ignore_start=False,
+            start_metadata={"foo": "bar"},
+        )
+        assert "foo" in received_down[-1].metadata
 
 
 class TestParallelPipeline(unittest.IsolatedAsyncioTestCase):
@@ -42,7 +64,11 @@ class TestParallelPipeline(unittest.IsolatedAsyncioTestCase):
 
         frames_to_send = [TextFrame(text="Hello from Pipecat!")]
         expected_returned_frames = [TextFrame]
-        await run_test(pipeline, frames_to_send, expected_returned_frames)
+        await run_test(
+            pipeline,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_returned_frames,
+        )
 
     async def test_parallel_multiple(self):
         """Should only passthrough one instance of TextFrame."""
@@ -50,7 +76,11 @@ class TestParallelPipeline(unittest.IsolatedAsyncioTestCase):
 
         frames_to_send = [TextFrame(text="Hello from Pipecat!")]
         expected_returned_frames = [TextFrame]
-        await run_test(pipeline, frames_to_send, expected_returned_frames)
+        await run_test(
+            pipeline,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_returned_frames,
+        )
 
 
 class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
@@ -79,7 +109,9 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
         task = PipelineTask(
             pipeline,
             params=PipelineParams(
-                enable_heartbeats=True, heartbeats_period_secs=0.2, observers=[heartbeats_observer]
+                enable_heartbeats=True,
+                heartbeats_period_secs=0.2,
+                observers=[heartbeats_observer],
             ),
         )
         task.set_event_loop(asyncio.get_event_loop())