From 568eb2ef4c1d90db041f3ee913edec47f48e1573 Mon Sep 17 00:00:00 2001 From: Kwindla Hultman Kramer Date: Tue, 16 Jul 2024 12:46:33 -0700 Subject: [PATCH] cartesia websockets and streaming --- src/pipecat/services/ai_services.py | 7 +- src/pipecat/services/cartesia.py | 145 ++++++++++------------------ 2 files changed, 58 insertions(+), 94 deletions(-) diff --git a/src/pipecat/services/ai_services.py b/src/pipecat/services/ai_services.py index 46bba673e..cc40035c4 100644 --- a/src/pipecat/services/ai_services.py +++ b/src/pipecat/services/ai_services.py @@ -149,6 +149,10 @@ class TTSService(AIService): async def say(self, text: str): await self.process_frame(TextFrame(text=text), FrameDirection.DOWNSTREAM) + async def handle_interruption(self, frame: StartInterruptionFrame, direction: FrameDirection): + self._current_sentence = "" + await self.push_frame(frame, direction) + async def _process_text_frame(self, frame: TextFrame): text: str | None = None if not self._aggregate_sentences: @@ -182,8 +186,7 @@ class TTSService(AIService): if isinstance(frame, TextFrame): await self._process_text_frame(frame) elif isinstance(frame, StartInterruptionFrame): - self._current_sentence = "" - await self.push_frame(frame, direction) + await self.handle_interruption(frame, direction) elif isinstance(frame, LLMFullResponseEndFrame) or isinstance(frame, EndFrame): self._current_sentence = "" await self._push_tts_frames(self._current_sentence) diff --git a/src/pipecat/services/cartesia.py b/src/pipecat/services/cartesia.py index c73f6f832..8047c06ab 100644 --- a/src/pipecat/services/cartesia.py +++ b/src/pipecat/services/cartesia.py @@ -11,7 +11,8 @@ import asyncio from typing import AsyncGenerator -from pipecat.frames.frames import AudioRawFrame, CancelFrame, EndFrame, Frame, StartFrame +from pipecat.processors.frame_processor import FrameDirection +from pipecat.frames.frames import Frame, AudioRawFrame, StartInterruptionFrame from pipecat.services.ai_services import TTSService from loguru import logger @@ -39,7 +40,7 @@ class CartesiaTTSService(TTSService): encoding: str = "pcm_s16le", sample_rate: int = 16000, **kwargs): - super().__init__(aggregate_sentences=True, **kwargs) + super().__init__(**kwargs) self._api_key = api_key self._cartesia_version = cartesia_version @@ -53,31 +54,49 @@ class CartesiaTTSService(TTSService): } self._language = "en" + self._websocket = None self._context_id = None self._receive_task = None + self._waiting_for_ttfb = False def can_generate_metrics(self) -> bool: return True - async def start(self, frame: StartFrame): + async def connect(self): try: self._websocket = await websockets.connect( f"{self._url}?api_key={self._api_key}&cartesia_version={self._cartesia_version}" ) - # self._receive_task = self.get_event_loop().create_task(self._receive_task_handler()) except Exception as e: logger.exception(f"{self} initialization error: {e}") + async def disconnect(self): + try: + if self._websocket: + ws = self._websocket + self._websocket = None + await ws.close() + except Exception as e: + logger.exception(f"{self} error closing websocket: {e}") + + async def handle_interruption(self, frame: StartInterruptionFrame, direction: FrameDirection): + await super().handle_interruption(frame, direction) + if self._receive_task: + self._receive_task.cancel() + self._receive_task = None + await self.disconnect() + await self.stop_all_metrics() + async def _receive_task_handler(self): - logger.debug("TOP OF RECEIVE TASK ...") async for message in self._websocket: - logger.debug("RECEIVE TASK LOOP") msg = json.loads(message) if not msg: continue - logger.debug(f"Received message: {msg}") + # logger.debug(f"Received message: {msg}") + if self._waiting_for_ttfb: + await self.stop_ttfb_metrics() + self._waiting_for_ttfb = False if msg["done"]: - logger.debug(f"This was a 'done' message, shut down the receive task.") self._context_id = None if self._receive_task: self._receive_task.cancel() @@ -90,96 +109,38 @@ class CartesiaTTSService(TTSService): ) await self.push_frame(frame) - # async for message in self._websocket: - # utterance = json.loads(message) - # if not utterance: - # continue - - # logger.debug(f"Received utterance: {utterance}") - # return - - # # TODO: PORT FROM GLADIA - # if "error" in utterance: - # message = utterance["message"] - # logger.error(f"Gladia error: {message}") - # elif "confidence" in utterance: - # type = utterance["type"] - # confidence = utterance["confidence"] - # transcript = utterance["transcription"] - # if confidence >= self._confidence: - # if type == "final": - # await self.queue_frame(TranscriptionFrame(transcript, "", int(time.time_ns() / 1000000))) - # else: - # await self.queue_frame(InterimTranscriptionFrame(transcript, "", - # int(time.time_ns() / 1000000))) - - async def stop(self, frame: EndFrame): - self._context_id = None - if self._receive_task: - self._receive_task.cancel() - self._receive_task = None - return - - async def cancel(self, frame: CancelFrame): - self._context_id = None - if self._receive_task: - self._receive_task.cancel() - self._receive_task = None - return - async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]: logger.debug(f"Generating TTS: [{text}]") - logger.debug( - f"model_id: {self._model_id}, voice_id: {self._voice_id}, language: {self._language}" - ) try: - await self.start_ttfb_metrics() + if not self._websocket: + await self.connect() + + if not self._waiting_for_ttfb: + await self.start_ttfb_metrics() + self._waiting_for_ttfb = True if not self._context_id: self._context_id = str(uuid.uuid4()) - msg = { - "transcript": text, - "continue": True, - "context_id": self._context_id, - "model_id": self._model_id, - "voice": { - "mode": "id", - "id": self._voice_id - }, - "output_format": self._output_format, - "language": self._language, - } - logger.debug(f"SENDING FIRST MESSAGE {json.dumps(msg)}") - await self._websocket.send(json.dumps(msg)) - logger.debug("AWAITING FIRST RESPONSE MESSAGE") - message = await self._websocket.recv() - msg = json.loads(message) - logger.debug(f"Received message: {msg}") - if (msg["type"] == "error"): - logger.error(f"Error: {msg}") - return - frame = AudioRawFrame( - audio=base64.b64decode(msg["data"]), - sample_rate=self._output_format["sample_rate"], - num_channels=1 - ) - yield frame - if not msg["done"]: - logger.debug("CREATING RECEIVE TASK") - self._receive_task = self.get_event_loop().create_task(self._receive_task_handler()) - # todo: how do we await this task at the app level, so the program doesn't exit? - # we can't await here because we need this function to return - # await self._receive_task - else: - msg = { - "transcript": text, - "continue": True, - "context_id": self._context_id, - } - await asyncio.sleep(0.350) - logger.debug(f"SENDING FOLLOW MESSAGE {json.dumps(msg)}") - await self._websocket.send(json.dumps(msg)) - yield None + + msg = { + "transcript": text, + "continue": True, + "context_id": self._context_id, + "model_id": self._model_id, + "voice": { + "mode": "id", + "id": self._voice_id + }, + "output_format": self._output_format, + "language": self._language, + } + # logger.debug(f"SENDING MESSAGE {json.dumps(msg)}") + await self._websocket.send(json.dumps(msg)) + if not self._receive_task: + # todo: how do we await this task at the app level, so the program doesn't exit? + # we can't await here because we need this function to return + self._receive_task = self.get_event_loop().create_task(self._receive_task_handler()) + yield None except Exception as e: logger.exception(f"{self} exception: {e}")