Use a single long-lived Task to push TTSStoppedFrame

This commit is contained in:
Sharvil Nanavati
2024-08-24 16:15:42 +00:00
parent 60c3d33def
commit 8ac7fb1a67

View File

@@ -4,11 +4,11 @@
# SPDX-License-Identifier: BSD 2-Clause License # SPDX-License-Identifier: BSD 2-Clause License
# #
import asyncio
import io import io
import wave import wave
from abc import abstractmethod from abc import abstractmethod
from asyncio import Task, sleep
from typing import AsyncGenerator, Optional from typing import AsyncGenerator, Optional
from pipecat.frames.frames import ( from pipecat.frames.frames import (
@@ -169,7 +169,8 @@ class TTSService(AIService):
self._push_text_frames: bool = push_text_frames self._push_text_frames: bool = push_text_frames
self._push_stop_frames: bool = push_stop_frames self._push_stop_frames: bool = push_stop_frames
self._stop_frame_timeout_s: float = stop_frame_timeout_s self._stop_frame_timeout_s: float = stop_frame_timeout_s
self._stop_frame_task: Optional[Task] = None self._stop_frame_task: Optional[asyncio.Task] = None
self._stop_frame_queue: asyncio.Queue = asyncio.Queue()
self._current_sentence: str = "" self._current_sentence: str = ""
@abstractmethod @abstractmethod
@@ -240,18 +241,29 @@ class TTSService(AIService):
async def push_frame(self, frame: Frame, direction: FrameDirection = FrameDirection.DOWNSTREAM): async def push_frame(self, frame: Frame, direction: FrameDirection = FrameDirection.DOWNSTREAM):
await super().push_frame(frame, direction) await super().push_frame(frame, direction)
if isinstance(frame, AudioRawFrame) and self._stop_frame_task is not None: if self._push_stop_frames and (
# Reschedule timeout task if it was already running isinstance(frame, StartInterruptionFrame) or
self._stop_frame_task.cancel() isinstance(frame, TTSStartedFrame) or
self._stop_frame_task = self.get_event_loop().create_task(self._stop_frame_handler()) isinstance(frame, AudioRawFrame)):
elif isinstance(frame, TTSStartedFrame) and self._push_stop_frames: if self._stop_frame_task is None:
# Start timeout task if necessary event_loop = self.get_event_loop()
self._stop_frame_task = self.get_event_loop().create_task(self._stop_frame_handler()) self._stop_frame_task = event_loop.create_task(self._stop_frame_handler())
await self._stop_frame_queue.put(frame)
async def _stop_frame_handler(self): async def _stop_frame_handler(self):
await sleep(self._stop_frame_timeout_s) has_started = False
await self.push_frame(TTSStoppedFrame()) while True:
self._stop_frame_task = None try:
frame = await asyncio.wait_for(self._stop_frame_queue.get(),
self._stop_frame_timeout_s)
if isinstance(frame, TTSStartedFrame):
has_started = True
elif isinstance(frame, StartInterruptionFrame):
has_started = False
except asyncio.TimeoutError:
if has_started:
await self.push_frame(TTSStoppedFrame())
has_started = False
class STTService(AIService): class STTService(AIService):