From 0aedaa8553e01f152f611b55963c963b1632d915 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= <aleix@daily.co>
Date: Mon, 10 Jun 2024 21:18:42 -0700
Subject: [PATCH] services(deepgram): abstract StartFrame/EndFrame/CancelFrame

---
 src/pipecat/services/ai_services.py | 20 +++++++++++++++++++
 src/pipecat/services/deepgram.py    | 31 +++++++++++++----------------
 2 files changed, 34 insertions(+), 17 deletions(-)

diff --git a/src/pipecat/services/ai_services.py b/src/pipecat/services/ai_services.py
index a46f41f35..82175b797 100644
--- a/src/pipecat/services/ai_services.py
+++ b/src/pipecat/services/ai_services.py
@@ -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):
diff --git a/src/pipecat/services/deepgram.py b/src/pipecat/services/deepgram.py
index 8afc4f0f3..1c3b890ee 100644
--- a/src/pipecat/services/deepgram.py
+++ b/src/pipecat/services/deepgram.py
@@ -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())