diff --git a/src/pipecat/pipeline/task.py b/src/pipecat/pipeline/task.py index 520998988..6bf5dd688 100644 --- a/src/pipecat/pipeline/task.py +++ b/src/pipecat/pipeline/task.py @@ -310,8 +310,8 @@ class PipelineTask(BaseTask): """Return the turn trace observer if enabled.""" return self._turn_trace_observer - async def add_observer(self, observer: BaseObserver): - await self._observer.add_observer(observer) + def add_observer(self, observer: BaseObserver): + self._observer.add_observer(observer) async def remove_observer(self, observer: BaseObserver): await self._observer.remove_observer(observer) diff --git a/src/pipecat/pipeline/task_observer.py b/src/pipecat/pipeline/task_observer.py index b7a54f58d..950ddc43b 100644 --- a/src/pipecat/pipeline/task_observer.py +++ b/src/pipecat/pipeline/task_observer.py @@ -49,21 +49,31 @@ class TaskObserver(BaseObserver): super().__init__(**kwargs) self._observers = observers or [] self._task_manager = task_manager - self._proxies: Dict[BaseObserver, Proxy] = {} + self._proxies: Optional[Dict[BaseObserver, Proxy]] = ( + None # Becomes a dict after start() is called + ) - async def add_observer(self, observer: BaseObserver): - proxy = self._create_proxy(observer) - self._proxies[observer] = proxy + def add_observer(self, observer: BaseObserver): + # Add the observer to the list. self._observers.append(observer) + # If we already started, create a new proxy for the observer. + # Otherwise, it will be created in start(). + if self._started(): + proxy = self._create_proxy(observer) + self._proxies[observer] = proxy + async def remove_observer(self, observer: BaseObserver): + # If the observer has a proxy, remove it. if observer in self._proxies: proxy = self._proxies[observer] # Remove the proxy so it doesn't get called anymore. del self._proxies[observer] # Cancel the proxy task right away. await self._task_manager.cancel_task(proxy.task) - # Remove the observer. + + # Remove the observer from the list. + if observer in self._observers: self._observers.remove(observer) async def start(self): @@ -79,6 +89,9 @@ class TaskObserver(BaseObserver): for proxy in self._proxies.values(): await proxy.queue.put(data) + def _started(self) -> bool: + return self._proxies is not None + def _create_proxy(self, observer: BaseObserver) -> Proxy: queue = asyncio.Queue() task = self._task_manager.create_task( diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 9bc737b1a..10b0ccaea 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -142,7 +142,7 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase): observer = CustomAddObserver() # Wait after the pipeline is started and add an observer. await asyncio.sleep(0.1) - await task.add_observer(observer) + 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)