pipeline: allow stopping tasks with StopTaskFrame

This commit is contained in:
Aleix Conchillo Flaqué
2024-05-14 00:30:32 -07:00
parent 922cdefee5
commit 11aa9dc803
2 changed files with 15 additions and 2 deletions

View File

@@ -205,6 +205,17 @@ class ErrorFrame(SystemFrame):
def __str__(self): def __str__(self):
return f"{self.name}(error: {self.error})" return f"{self.name}(error: {self.error})"
@dataclass
class StopTaskFrame(SystemFrame):
"""Indicates that a pipeline task should be stopped. This should inform the
pipeline processors that they should stop pushing frames but that they
should be kept in a running state.
"""
pass
# #
# Control frames # Control frames
# #

View File

@@ -8,7 +8,7 @@ import asyncio
from typing import AsyncIterable, Iterable from typing import AsyncIterable, Iterable
from pipecat.frames.frames import CancelFrame, EndFrame, ErrorFrame, Frame, StartFrame from pipecat.frames.frames import CancelFrame, EndFrame, ErrorFrame, Frame, StartFrame, StopTaskFrame
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.utils.utils import obj_count, obj_id from pipecat.utils.utils import obj_count, obj_id
@@ -74,7 +74,9 @@ class PipelineTask:
frame = await self._task_queue.get() frame = await self._task_queue.get()
await self._source.process_frame(frame, FrameDirection.DOWNSTREAM) await self._source.process_frame(frame, FrameDirection.DOWNSTREAM)
self._task_queue.task_done() self._task_queue.task_done()
running = not (isinstance(frame, CancelFrame) or isinstance(frame, EndFrame)) running = not (isinstance(frame, StopTaskFrame) or
isinstance(frame, CancelFrame) or
isinstance(frame, EndFrame))
# We just enqueue None to terminate the task. # We just enqueue None to terminate the task.
await self._up_queue.put(None) await self._up_queue.put(None)