From ef380321cfd5f0daa6289fe9d7c253be8d6d8824 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= <aleix@daily.co>
Date: Mon, 10 Jun 2024 20:33:27 -0700
Subject: [PATCH] services: added new DeepgramSTTService

---
 CHANGELOG.md                     | 10 ++++
 src/pipecat/services/deepgram.py | 98 +++++++++++++++++++++++++++++++-
 2 files changed, 106 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 62d4bfc81..33cc37438 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,16 @@ All notable changes to **pipecat** will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Changed
+
+- Added `DeepgramSTTService`. This service has an ongoing websocket
+  connection. To handle this, it subclasses `AIService` instead of
+  `STTService`. The output of this service will be pushed from the same task,
+  except system frames like `StartFrame`, `CancelFrame` or
+  `StartInterruptionFrame`.
+
 ## [0.0.29] - 2024-06-07
 
 ### Added
diff --git a/src/pipecat/services/deepgram.py b/src/pipecat/services/deepgram.py
index ce418916b..8afc4f0f3 100644
--- a/src/pipecat/services/deepgram.py
+++ b/src/pipecat/services/deepgram.py
@@ -5,11 +5,29 @@
 #
 
 import aiohttp
+import asyncio
+import time
 
 from typing import AsyncGenerator
 
-from pipecat.frames.frames import AudioRawFrame, ErrorFrame, Frame
-from pipecat.services.ai_services import TTSService
+from pipecat.frames.frames import (
+    AudioRawFrame,
+    CancelFrame,
+    EndFrame,
+    ErrorFrame,
+    Frame,
+    InterimTranscriptionFrame,
+    StartFrame,
+    SystemFrame,
+    TranscriptionFrame)
+from pipecat.processors.frame_processor import FrameDirection
+from pipecat.services.ai_services import AIService, TTSService
+
+from deepgram import (
+    DeepgramClient,
+    LiveTranscriptionEvents,
+    LiveOptions,
+)
 
 from loguru import logger
 
@@ -55,3 +73,79 @@ class DeepgramTTSService(TTSService):
                     yield frame
         except Exception as e:
             logger.error(f"Deepgram exception: {e}")
+
+
+class DeepgramSTTService(AIService):
+    def __init__(self,
+                 api_key: str,
+                 live_options: LiveOptions = LiveOptions(
+                     encoding="linear16",
+                     language="en-US",
+                     model="nova-2-conversationalai",
+                     sample_rate=16000,
+                     channels=1,
+                     interim_results=True,
+                     smart_format=True,
+                 ),
+                 **kwargs):
+        super().__init__(**kwargs)
+
+        self._live_options = live_options
+
+        self._client = DeepgramClient(api_key)
+        self._connection = self._client.listen.asynclive.v("1")
+        self._connection.on(LiveTranscriptionEvents.Transcript, self._on_message)
+
+        self._create_push_task()
+
+    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
+        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 _stop(self):
+        await self._connection.finish()
+
+    def _create_push_task(self):
+        self._push_frame_task = self.get_event_loop().create_task(self._push_frame_task_handler())
+        self._push_queue = asyncio.Queue()
+
+    async def _push_frame_task_handler(self):
+        running = True
+        while running:
+            try:
+                (frame, direction) = await self._push_queue.get()
+                await self.push_frame(frame, direction)
+                running = not isinstance(frame, EndFrame)
+            except asyncio.CancelledError:
+                break
+
+    async def _on_message(self, *args, **kwargs):
+        result = kwargs["result"]
+        is_final = result.is_final
+        transcript = result.channel.alternatives[0].transcript
+        if len(transcript) > 0:
+            if is_final:
+                await self._push_queue.put((TranscriptionFrame(transcript, "", int(time.time_ns() / 1000000)), FrameDirection.DOWNSTREAM))
+            else:
+                await self._push_queue.put((InterimTranscriptionFrame(transcript, "", int(time.time_ns() / 1000000)), FrameDirection.DOWNSTREAM))