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

services(deepgram): inherit from STTService instead of AsyncAIService

Browse Source
This commit is contained in:
Aleix Conchillo Flaqué
2024-08-26 11:10:42 -07:00
parent 3931cb3235
commit 6629b853c5
1 changed files with 38 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -16,12 +16,10 @@ from pipecat.frames.frames import (
Frame,
InterimTranscriptionFrame,
StartFrame,
SystemFrame,
TTSStartedFrame,
TTSStoppedFrame,
TranscriptionFrame)
from pipecat.processors.frame_processor import FrameDirection
from pipecat.services.ai_services import AsyncAIService, TTSService
from pipecat.services.ai_services import STTService, TTSService
from pipecat.transcriptions.language import Language
from pipecat.utils.time import time_now_iso8601
@@ -45,15 +43,6 @@ 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__(
@@ -119,7 +108,7 @@ class DeepgramTTSService(TTSService):
logger.exception(f"{self} exception: {e}")
class DeepgramSTTService(AsyncAIService):
class DeepgramSTTService(STTService):
def __init__(self,
*,
api_key: str,
@@ -132,6 +121,8 @@ class DeepgramSTTService(AsyncAIService):
channels=1,
interim_results=True,
smart_format=True,
punctuate=True,
profanity_filter=True,
),
**kwargs):
super().__init__(**kwargs)
@@ -143,30 +134,46 @@ class DeepgramSTTService(AsyncAIService):
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):
await super().process_frame(frame, direction)
async def set_model(self, model: str):
logger.debug(f"Switching STT model to: [{model}]")
self._live_options.model = model
await self._disconnect()
await self._connect()
if isinstance(frame, SystemFrame):
await self.push_frame(frame, direction)
elif isinstance(frame, AudioRawFrame):
await self._connection.send(frame.audio)
else:
await self.queue_frame(frame, direction)
async def set_language(self, language: Language):
logger.debug(f"Switching STT language to: [{language}]")
self._live_options.language = language.value
await self._disconnect()
await self._connect()
async def start(self, frame: StartFrame):
await super().start(frame)
await self._connect()
async def stop(self, frame: EndFrame):
await super().stop(frame)
await self._disconnect()
async def cancel(self, frame: CancelFrame):
await super().cancel(frame)
await self._disconnect()
async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
await self.start_processing_metrics()
await self._connection.send(audio)
yield None
await self.stop_processing_metrics()
async def _connect(self):
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, frame: EndFrame):
await super().stop(frame)
await self._connection.finish()
async def cancel(self, frame: CancelFrame):
await super().cancel(frame)
await self._connection.finish()
async def _disconnect(self):
if self._connection.is_connected:
await self._connection.finish()
logger.debug(f"{self}: Disconnected from Deepgram")
async def _on_message(self, *args, **kwargs):
result: LiveResultResponse = kwargs["result"]
@@ -177,9 +184,9 @@ class DeepgramSTTService(AsyncAIService):
language = None
if result.channel.alternatives[0].languages:
language = result.channel.alternatives[0].languages[0]
language = deepgram_language_to_language(language)
language = Language(language)
if len(transcript) > 0:
if is_final:
await self.queue_frame(TranscriptionFrame(transcript, "", time_now_iso8601(), language))
await self.push_frame(TranscriptionFrame(transcript, "", time_now_iso8601(), language))
else:
await self.queue_frame(InterimTranscriptionFrame(transcript, "", time_now_iso8601(), language))
await self.push_frame(InterimTranscriptionFrame(transcript, "", time_now_iso8601(), language))