From 38cd86ad52d5790c23c049a47f8b166d52902134 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= <aleix@daily.co>
Date: Thu, 22 Aug 2024 22:43:30 -0700
Subject: [PATCH] services: added language to transcription frames

---
 src/pipecat/services/deepgram.py         | 26 ++++++++++++++++++++----
 src/pipecat/transports/services/daily.py | 19 +++++++++++++++--
 2 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/src/pipecat/services/deepgram.py b/src/pipecat/services/deepgram.py
index 035fdd25c..30a572d38 100644
--- a/src/pipecat/services/deepgram.py
+++ b/src/pipecat/services/deepgram.py
@@ -22,6 +22,7 @@ from pipecat.frames.frames import (
     TranscriptionFrame)
 from pipecat.processors.frame_processor import FrameDirection
 from pipecat.services.ai_services import AsyncAIService, TTSService
+from pipecat.transcriptions.language import Language
 from pipecat.utils.time import time_now_iso8601
 
 from loguru import logger
@@ -30,10 +31,12 @@ from loguru import logger
 # See .env.example for Deepgram configuration needed
 try:
     from deepgram import (
+        AsyncListenWebSocketClient,
         DeepgramClient,
         DeepgramClientOptions,
         LiveTranscriptionEvents,
         LiveOptions,
+        LiveResultResponse
     )
 except ModuleNotFoundError as e:
     logger.error(f"Exception: {e}")
@@ -42,6 +45,15 @@ except ModuleNotFoundError as e:
     raise Exception(f"Missing module: {e}")
 
 
+def deepgram_language_to_language(language: str) -> Language | None:
+    match language:
+        case "en":
+            return Language.EN
+        case "es":
+            return Language.ES
+    return None
+
+
 class DeepgramTTSService(TTSService):
 
     def __init__(
@@ -128,7 +140,7 @@ class DeepgramSTTService(AsyncAIService):
 
         self._client = DeepgramClient(
             api_key, config=DeepgramClientOptions(url=url, options={"keepalive": "true"}))
-        self._connection = self._client.listen.asynclive.v("1")
+        self._connection: AsyncListenWebSocketClient = self._client.listen.asyncwebsocket.v("1")
         self._connection.on(LiveTranscriptionEvents.Transcript, self._on_message)
 
     async def process_frame(self, frame: Frame, direction: FrameDirection):
@@ -157,11 +169,17 @@ class DeepgramSTTService(AsyncAIService):
         await self._connection.finish()
 
     async def _on_message(self, *args, **kwargs):
-        result = kwargs["result"]
+        result: LiveResultResponse = kwargs["result"]
+        if len(result.channel.alternatives) == 0:
+            return
         is_final = result.is_final
         transcript = result.channel.alternatives[0].transcript
+        language = None
+        if result.channel.alternatives[0].languages:
+            language = result.channel.alternatives[0].languages[0]
+            language = deepgram_language_to_language(language)
         if len(transcript) > 0:
             if is_final:
-                await self.queue_frame(TranscriptionFrame(transcript, "", time_now_iso8601()))
+                await self.queue_frame(TranscriptionFrame(transcript, "", time_now_iso8601(), language))
             else:
-                await self.queue_frame(InterimTranscriptionFrame(transcript, "", time_now_iso8601()))
+                await self.queue_frame(InterimTranscriptionFrame(transcript, "", time_now_iso8601(), language))
diff --git a/src/pipecat/transports/services/daily.py b/src/pipecat/transports/services/daily.py
index 98bc3cc3c..188ae65ce 100644
--- a/src/pipecat/transports/services/daily.py
+++ b/src/pipecat/transports/services/daily.py
@@ -36,6 +36,7 @@ from pipecat.frames.frames import (
     UserImageRawFrame,
     UserImageRequestFrame)
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
+from pipecat.transcriptions.language import Language
 from pipecat.transports.base_input import BaseInputTransport
 from pipecat.transports.base_output import BaseOutputTransport
 from pipecat.transports.base_transport import BaseTransport, TransportParams
@@ -746,6 +747,15 @@ class DailyOutputTransport(BaseOutputTransport):
         await self._client.write_frame_to_camera(frame)
 
 
+def daily_language_to_language(language: str) -> Language | None:
+    match language:
+        case "en":
+            return Language.EN
+        case "es":
+            return Language.ES
+    return None
+
+
 class DailyTransport(BaseTransport):
 
     def __init__(
@@ -950,11 +960,16 @@ class DailyTransport(BaseTransport):
         text = message["text"]
         timestamp = message["timestamp"]
         is_final = message["rawResponse"]["is_final"]
+        try:
+            language = message["rawResponse"]["channel"]["alternatives"][0]["languages"][0]
+            language = daily_language_to_language(language)
+        except KeyError:
+            language = None
         if is_final:
-            frame = TranscriptionFrame(text, participant_id, timestamp)
+            frame = TranscriptionFrame(text, participant_id, timestamp, language)
             logger.debug(f"Transcription (from: {participant_id}): [{text}]")
         else:
-            frame = InterimTranscriptionFrame(text, participant_id, timestamp)
+            frame = InterimTranscriptionFrame(text, participant_id, timestamp, language)
 
         if self._input:
             await self._input.push_transcription_frame(frame)