services(deepgram): abstract StartFrame/EndFrame/CancelFrame

This commit is contained in:
Aleix Conchillo Flaqué
2024-06-10 21:18:42 -07:00
parent 6554479d39
commit 0aedaa8553
2 changed files with 34 additions and 17 deletions

View File

@@ -16,6 +16,7 @@ from pipecat.frames.frames import (
EndFrame,
ErrorFrame,
Frame,
StartFrame,
TTSStartedFrame,
TTSStoppedFrame,
TextFrame,
@@ -30,6 +31,25 @@ class AIService(FrameProcessor):
def __init__(self):
super().__init__()
async def start(self, frame: StartFrame):
pass
async def stop(self, frame: EndFrame):
pass
async def cancel(self, frame: CancelFrame):
pass
async def process_frame(self, frame: Frame, direction: FrameDirection):
await super().process_frame(frame, direction)
if isinstance(frame, StartFrame):
await self.start(frame)
elif isinstance(frame, CancelFrame):
await self.cancel(frame)
elif isinstance(frame, EndFrame):
await self.stop(frame)
async def process_generator(self, generator: AsyncGenerator[Frame, None]):
async for f in generator:
if isinstance(f, ErrorFrame):

View File

@@ -101,30 +101,27 @@ class DeepgramSTTService(AIService):
async def process_frame(self, frame: Frame, direction: FrameDirection):
await super().process_frame(frame, direction)
if isinstance(frame, StartFrame):
await self._start()
await self.push_frame(frame)
elif isinstance(frame, CancelFrame):
await self._stop()
self._push_frame_task.cancel()
await self.push_frame(frame)
elif isinstance(frame, SystemFrame):
await self.push_frame(frame)
elif isinstance(frame, EndFrame):
await self._stop()
await self._push_queue.put((frame, direction))
await self._push_frame_task
if isinstance(frame, SystemFrame):
await self.push_frame(frame, direction)
elif isinstance(frame, AudioRawFrame):
await self._connection.send(frame.audio)
else:
await self._push_queue.put((frame, direction))
async def _start(self):
if not await self._connection.start(self._live_options):
logger.error("Unable to connect to Deepgram")
async def start(self, frame: StartFrame):
if await self._connection.start(self._live_options):
logger.debug(f"{self}: Connected to Deepgram")
else:
logger.error(f"{self}: Unable to connect to Deepgram")
async def _stop(self):
async def stop(self, frame: EndFrame):
await self._connection.finish()
await self._push_queue.put((frame, FrameDirection.DOWNSTREAM))
await self._push_frame_task
async def cancel(self, frame: CancelFrame):
await self._connection.finish()
self._push_frame_task.cancel()
def _create_push_task(self):
self._push_frame_task = self.get_event_loop().create_task(self._push_frame_task_handler())