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):
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
#

View File

@@ -8,7 +8,7 @@ import asyncio
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.utils.utils import obj_count, obj_id
@@ -74,7 +74,9 @@ class PipelineTask:
frame = await self._task_queue.get()
await self._source.process_frame(frame, FrameDirection.DOWNSTREAM)
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.
await self._up_queue.put(None)