From c8da5314020bd525df1727f348a14121efa66728 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Wed, 15 Jan 2025 09:48:10 -0800 Subject: [PATCH] pipeline(task): add support for pipeline frame observers --- src/pipecat/frames/frames.py | 6 +++- src/pipecat/pipeline/task.py | 42 +++++++++++++++++++++-- src/pipecat/processors/frame_processor.py | 6 ++++ 3 files changed, 51 insertions(+), 3 deletions(-) diff --git a/src/pipecat/frames/frames.py b/src/pipecat/frames/frames.py index 6e5fa744f..e2ac25e7a 100644 --- a/src/pipecat/frames/frames.py +++ b/src/pipecat/frames/frames.py @@ -5,7 +5,7 @@ # from dataclasses import dataclass, field -from typing import Any, Awaitable, Callable, List, Literal, Mapping, Optional, Tuple +from typing import TYPE_CHECKING, Any, Awaitable, Callable, List, Literal, Mapping, Optional, Tuple from pipecat.audio.vad.vad_analyzer import VADParams from pipecat.clocks.base_clock import BaseClock @@ -14,6 +14,9 @@ from pipecat.transcriptions.language import Language from pipecat.utils.time import nanoseconds_to_str from pipecat.utils.utils import obj_count, obj_id +if TYPE_CHECKING: + from pipecat.observers.base_observer import BaseObserver + def format_pts(pts: int | None): return nanoseconds_to_str(pts) if pts else None @@ -386,6 +389,7 @@ class StartFrame(SystemFrame): enable_metrics: bool = False enable_usage_metrics: bool = False report_only_initial_ttfb: bool = False + observer: Optional["BaseObserver"] = None @dataclass diff --git a/src/pipecat/pipeline/task.py b/src/pipecat/pipeline/task.py index 79a708c74..db9a00049 100644 --- a/src/pipecat/pipeline/task.py +++ b/src/pipecat/pipeline/task.py @@ -5,10 +5,10 @@ # import asyncio -from typing import AsyncIterable, Iterable +from typing import AsyncIterable, Iterable, List from loguru import logger -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict from pipecat.clocks.base_clock import BaseClock from pipecat.clocks.system_clock import SystemClock @@ -24,20 +24,31 @@ from pipecat.frames.frames import ( StopTaskFrame, ) from pipecat.metrics.metrics import ProcessingMetricsData, TTFBMetricsData +from pipecat.observers.base_observer import BaseObserver from pipecat.pipeline.base_pipeline import BasePipeline from pipecat.processors.frame_processor import FrameDirection, FrameProcessor from pipecat.utils.utils import obj_count, obj_id class PipelineParams(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + allow_interruptions: 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] = [] class Source(FrameProcessor): + """This is the source processor that is linked at the beginning of the + pipeline given to the pipeline task. It allows us to easily push frames + downstream to the pipeline and also receive upstream frames coming from the + pipeline. + + """ + def __init__(self, up_queue: asyncio.Queue): super().__init__() self._up_queue = up_queue @@ -68,6 +79,12 @@ class Source(FrameProcessor): class Sink(FrameProcessor): + """This is the sink processor that is linked at the end of the pipeline + given to the pipeline task. It allows us to receive downstream frames and + act on them, for example, waiting to receive an EndFrame. + + """ + def __init__(self, down_queue: asyncio.Queue): super().__init__() self._down_queue = down_queue @@ -80,6 +97,24 @@ class Sink(FrameProcessor): await self._down_queue.put(frame) +class Observer(BaseObserver): + """This is a pipeline frame observer that is used as a proxy to the user + provided observers. That is, this is the only observer passed to the frame + processors. Then, every time a frame is pushed this observer will call all + the observers registered to the pipeline task. + + """ + + def __init__(self, observers: List[BaseObserver] = []): + self._observers = observers + + async def on_push_frame( + self, src: FrameProcessor, dst: FrameProcessor, frame: Frame, direction: FrameDirection + ): + for observer in self._observers: + await observer.on_push_frame(src, dst, frame, direction) + + class PipelineTask: def __init__( self, @@ -105,6 +140,8 @@ class PipelineTask: self._sink = Sink(self._down_queue) pipeline.link(self._sink) + self._observer = Observer(params.observers) + def has_finished(self): return self._finished @@ -156,6 +193,7 @@ class PipelineTask: 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, clock=self._clock, ) await self._source.queue_frame(start_frame, FrameDirection.DOWNSTREAM) diff --git a/src/pipecat/processors/frame_processor.py b/src/pipecat/processors/frame_processor.py index 653b313b8..ed6f8a4cb 100644 --- a/src/pipecat/processors/frame_processor.py +++ b/src/pipecat/processors/frame_processor.py @@ -58,6 +58,7 @@ class FrameProcessor: self._enable_metrics = False self._enable_usage_metrics = False self._report_only_initial_ttfb = False + self._observer = None # Cancellation is done through CancelFrame (a system frame). This could # cause other events being triggered (e.g. closing a transport) which @@ -194,6 +195,7 @@ class FrameProcessor: self._enable_metrics = frame.enable_metrics self._enable_usage_metrics = frame.enable_usage_metrics self._report_only_initial_ttfb = frame.report_only_initial_ttfb + self._observer = frame.observer elif isinstance(frame, StartInterruptionFrame): await self._start_interruption() await self.stop_all_metrics() @@ -258,9 +260,13 @@ class FrameProcessor: try: if direction == FrameDirection.DOWNSTREAM and self._next: logger.trace(f"Pushing {frame} from {self} to {self._next}") + if self._observer: + await self._observer.on_push_frame(self, self._next, frame, direction) await self._next.queue_frame(frame, direction) elif direction == FrameDirection.UPSTREAM and self._prev: logger.trace(f"Pushing {frame} upstream from {self} to {self._prev}") + if self._observer: + await self._observer.on_push_frame(self, self._prev, frame, direction) await self._prev.queue_frame(frame, direction) except Exception as e: logger.exception(f"Uncaught exception in {self}: {e}")