PipelineTask: add add_observer() and remove_observer()

tests/test_pipeline.py

@@ -8,13 +8,8 @@ import asyncio
 import time
 import unittest
 
-from pipecat.frames.frames import (
-    EndFrame,
-    HeartbeatFrame,
-    StartFrame,
-    StopFrame,
-    TextFrame,
-)
+from pipecat.frames.frames import EndFrame, HeartbeatFrame, StartFrame, StopFrame, TextFrame
+from pipecat.observers.base_observer import BaseObserver, FramePushed
 from pipecat.pipeline.parallel_pipeline import ParallelPipeline
 from pipecat.pipeline.pipeline import Pipeline
 from pipecat.pipeline.task import PipelineParams, PipelineTask
@@ -101,6 +96,70 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
         await task.run()
         assert task.has_finished()
 
+    async def test_task_observers(self):
+        frame_received = False
+
+        class CustomObserver(BaseObserver):
+            async def on_push_frame(self, data: FramePushed):
+                nonlocal frame_received
+
+                if isinstance(data.frame, TextFrame):
+                    frame_received = True
+
+        identity = IdentityFilter()
+        pipeline = Pipeline([identity])
+        task = PipelineTask(pipeline, observers=[CustomObserver()])
+        task.set_event_loop(asyncio.get_event_loop())
+
+        await task.queue_frames([TextFrame(text="Hello Downstream!"), EndFrame()])
+        await task.run()
+        assert frame_received
+
+    async def test_task_add_observer(self):
+        frame_received = False
+        frame_add_count = 0
+
+        class CustomObserver(BaseObserver):
+            async def on_push_frame(self, data: FramePushed):
+                nonlocal frame_received
+
+                if isinstance(data.frame, TextFrame):
+                    frame_received = True
+
+        class CustomAddObserver(BaseObserver):
+            async def on_push_frame(self, data: FramePushed):
+                nonlocal frame_add_count
+
+                if isinstance(data.source, IdentityFilter) and isinstance(data.frame, TextFrame):
+                    frame_add_count += 1
+
+        identity = IdentityFilter()
+        pipeline = Pipeline([identity])
+        task = PipelineTask(pipeline, observers=[CustomObserver()])
+        task.set_event_loop(asyncio.get_event_loop())
+
+        async def delayed_add_observer():
+            observer = CustomAddObserver()
+            # Wait after the pipeline is started and add an observer.
+            await asyncio.sleep(0.1)
+            await task.add_observer(observer)
+            # Push a TextFrame and wait for the observer to pick it up.
+            await task.queue_frame(TextFrame(text="Hello Downstream!"))
+            await asyncio.sleep(0.1)
+            # Remove the observer
+            await task.remove_observer(observer)
+            # Push another TextFrame. This time the counter should not
+            # increments since we have removed the observer.
+            await task.queue_frame(TextFrame(text="Hello Downstream!"))
+            await asyncio.sleep(0.1)
+            # Finally end the pipeline.
+            await task.queue_frame(EndFrame())
+
+        await asyncio.gather(task.run(), delayed_add_observer())
+
+        assert frame_received
+        assert frame_add_count == 1
+
     async def test_task_started_ended_event_handler(self):
         start_received = False
         end_received = False