Refactoring the FrameQueue to be able to track any Frame.

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``
+    Extends ``asyncio.Queue`` and maintains an internal per-type counter so
-    flag so interrupt-handling code can decide whether to cancel a task or
+    callers can check whether a frame of a given type (or any subclass) is
-    merely drain non-uninterruptible items without scanning the queue.
+    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):
+        frame_type = type(self._frame_getter(item))
-            self._uninterruptible_count += 1
+        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):
+        frame_type = type(self._frame_getter(item))
-            self._uninterruptible_count -= 1
+        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: