diff --git a/src/pipecat/utils/frame_queue.py b/src/pipecat/utils/frame_queue.py index 888bacb8e..f1e14b923 100644 --- a/src/pipecat/utils/frame_queue.py +++ b/src/pipecat/utils/frame_queue.py @@ -7,17 +7,17 @@ """Frame queue utilities for Pipecat pipeline processors.""" import asyncio -from typing import Any, Callable +from typing import Any, Callable, Dict, Type, Union from pipecat.frames.frames import Frame, UninterruptibleFrame class FrameQueue(asyncio.Queue): - """An asyncio.Queue that tracks whether any UninterruptibleFrame is enqueued. + """An asyncio.Queue that tracks the types of frames currently enqueued. - Extends ``asyncio.Queue`` and maintains an O(1) ``has_uninterruptible`` - flag so interrupt-handling code can decide whether to cancel a task or - merely drain non-uninterruptible items without scanning the queue. + Extends ``asyncio.Queue`` and maintains an internal per-type counter so + callers can check whether a frame of a given type (or any subclass) is + present without scanning the queue. Items may be raw ``Frame`` objects or tuples whose first element is a ``Frame`` (e.g. ``(frame, direction, callback)``). Pass a ``frame_getter`` @@ -39,22 +39,41 @@ class FrameQueue(asyncio.Queue): """ super().__init__() self._frame_getter = frame_getter - self._uninterruptible_count: int = 0 + self._frame_counts: Dict[type, int] = {} @property def has_uninterruptible(self) -> bool: """Return True if any UninterruptibleFrame is currently in the queue.""" - return self._uninterruptible_count > 0 + return self.has_frame(UninterruptibleFrame) + + def has_frame(self, frame_type: Union[Type[Frame], Type[UninterruptibleFrame]]) -> bool: + """Return True if any frame of the given type (or a subclass) is in the queue. + + ``frame_type`` may be ``Frame``, ``UninterruptibleFrame`` (a mixin, not a + ``Frame`` subclass), or any concrete frame type. + + Args: + frame_type: The frame class to check for. + + Returns: + True if at least one enqueued frame is an instance of ``frame_type``. + """ + return any( + count > 0 and issubclass(concrete, frame_type) + for concrete, count in self._frame_counts.items() + ) def _put(self, item: Any) -> None: - if isinstance(self._frame_getter(item), UninterruptibleFrame): - self._uninterruptible_count += 1 + frame_type = type(self._frame_getter(item)) + self._frame_counts[frame_type] = self._frame_counts.get(frame_type, 0) + 1 super()._put(item) def _get(self) -> Any: item = super()._get() - if isinstance(self._frame_getter(item), UninterruptibleFrame): - self._uninterruptible_count -= 1 + frame_type = type(self._frame_getter(item)) + self._frame_counts[frame_type] -= 1 + if self._frame_counts[frame_type] == 0: + del self._frame_counts[frame_type] return item def reset(self) -> None: