wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

services(deepgram): abstract StartFrame/EndFrame/CancelFrame

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

20
src/pipecat/services/ai_services.py
View File

@@ -16,6 +16,7 @@ from pipecat.frames.frames import (
EndFrame, EndFrame,
ErrorFrame, ErrorFrame,
Frame, Frame,
StartFrame,
TTSStartedFrame, TTSStartedFrame,
TTSStoppedFrame, TTSStoppedFrame,
TextFrame, TextFrame,
@@ -30,6 +31,25 @@ class AIService(FrameProcessor):
def __init__(self): def __init__(self):
super().__init__() 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 def process_generator(self, generator: AsyncGenerator[Frame, None]):
async for f in generator: async for f in generator:
if isinstance(f, ErrorFrame): if isinstance(f, ErrorFrame):

31
src/pipecat/services/deepgram.py
View File

@@ -101,30 +101,27 @@ class DeepgramSTTService(AIService):
async def process_frame(self, frame: Frame, direction: FrameDirection): async def process_frame(self, frame: Frame, direction: FrameDirection):
await super().process_frame(frame, direction) await super().process_frame(frame, direction)
if isinstance(frame, StartFrame): if isinstance(frame, SystemFrame):
await self._start() await self.push_frame(frame, direction)
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
elif isinstance(frame, AudioRawFrame): elif isinstance(frame, AudioRawFrame):
await self._connection.send(frame.audio) await self._connection.send(frame.audio)
else: else:
await self._push_queue.put((frame, direction)) await self._push_queue.put((frame, direction))
async def _start(self): async def start(self, frame: StartFrame):
if not await self._connection.start(self._live_options): if await self._connection.start(self._live_options):
logger.error("Unable to connect to Deepgram") 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._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): def _create_push_task(self):
self._push_frame_task = self.get_event_loop().create_task(self._push_frame_task_handler()) self._push_frame_task = self.get_event_loop().create_task(self._push_frame_task_handler())