Handle cancellation, stopping, and restarting

This commit is contained in:
Sharvil Nanavati
2024-08-27 01:24:00 +00:00
parent 1d92fc3199
commit c979762b70
2 changed files with 36 additions and 14 deletions

View File

@@ -238,32 +238,48 @@ class TTSService(AIService):
else: else:
await self.push_frame(frame, direction) await self.push_frame(frame, direction)
async def stop(self, frame: EndFrame):
if self._stop_frame_task:
self._stop_frame_task.cancel()
await self._stop_frame_task
self._stop_frame_task = None
async def cancel(self, frame: CancelFrame):
if self._stop_frame_task:
self._stop_frame_task.cancel()
await self._stop_frame_task
self._stop_frame_task = None
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 self._push_stop_frames and ( if self._push_stop_frames and (
isinstance(frame, StartInterruptionFrame) or isinstance(frame, StartInterruptionFrame) or
isinstance(frame, TTSStartedFrame) or isinstance(frame, TTSStartedFrame) or
isinstance(frame, AudioRawFrame)): isinstance(frame, AudioRawFrame) or
isinstance(frame, TTSStoppedFrame)):
if self._stop_frame_task is None: if self._stop_frame_task is None:
event_loop = self.get_event_loop() event_loop = self.get_event_loop()
self._stop_frame_task = 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) await self._stop_frame_queue.put(frame)
async def _stop_frame_handler(self): async def _stop_frame_handler(self):
has_started = False try:
while True: has_started = False
try: while True:
frame = await asyncio.wait_for(self._stop_frame_queue.get(), try:
self._stop_frame_timeout_s) frame = await asyncio.wait_for(self._stop_frame_queue.get(),
if isinstance(frame, TTSStartedFrame): self._stop_frame_timeout_s)
has_started = True if isinstance(frame, TTSStartedFrame):
elif isinstance(frame, StartInterruptionFrame): has_started = True
has_started = False elif isinstance(frame, (TTSStoppedFrame, StartInterruptionFrame)):
except asyncio.TimeoutError: has_started = False
if has_started: except asyncio.TimeoutError:
await self.push_frame(TTSStoppedFrame()) if has_started:
has_started = False await self.push_frame(TTSStoppedFrame())
has_started = False
except asyncio.CancelledError:
pass
class STTService(AIService): class STTService(AIService):

View File

@@ -19,6 +19,7 @@ from pipecat.frames.frames import (
Frame, Frame,
AudioRawFrame, AudioRawFrame,
StartFrame, StartFrame,
StartInterruptionFrame,
EndFrame, EndFrame,
TTSStartedFrame, TTSStartedFrame,
TTSStoppedFrame, TTSStoppedFrame,
@@ -86,6 +87,11 @@ class LmntTTSService(TTSService):
await super().cancel(frame) await super().cancel(frame)
await self._disconnect() await self._disconnect()
async def push_frame(self, frame: Frame, direction: FrameDirection = FrameDirection.DOWNSTREAM):
await super().push_frame(frame, direction)
if isinstance(frame, (TTSStoppedFrame, StartInterruptionFrame)):
self._started = False
async def _connect(self): async def _connect(self):
try: try:
self._speech = Speech() self._speech = Speech()