wip - using cartesia word timestamps for context management

This commit is contained in:
Kwindla Hultman Kramer
2024-07-17 14:13:01 -07:00
parent 568eb2ef4c
commit 270007b17c
4 changed files with 124 additions and 40 deletions

View File

@@ -13,7 +13,8 @@ from pipecat.frames.frames import EndFrame, LLMMessagesFrame
from pipecat.pipeline.pipeline import Pipeline
from pipecat.pipeline.runner import PipelineRunner
from pipecat.pipeline.task import PipelineTask
from pipecat.services.elevenlabs import ElevenLabsTTSService
# from pipecat.services.elevenlabs import ElevenLabsTTSService
from pipecat.services.cartesia import CartesiaTTSService
from pipecat.services.openai import OpenAILLMService
from pipecat.transports.services.daily import DailyParams, DailyTransport
@@ -36,16 +37,21 @@ async def main(room_url):
"Say One Thing From an LLM",
DailyParams(audio_out_enabled=True))
tts = ElevenLabsTTSService(
aiohttp_session=session,
api_key=os.getenv("ELEVENLABS_API_KEY"),
voice_id=os.getenv("ELEVENLABS_VOICE_ID"),
)
# tts = ElevenLabsTTSService(
# aiohttp_session=session,
# api_key=os.getenv("ELEVENLABS_API_KEY"),
# voice_id=os.getenv("ELEVENLABS_VOICE_ID"),
# )
llm = OpenAILLMService(
api_key=os.getenv("OPENAI_API_KEY"),
model="gpt-4o")
tts = CartesiaTTSService(
api_key=os.getenv("CARTESIA_API_KEY"),
voice_id="a0e99841-438c-4a64-b679-ae501e7d6091", # Barbershop Man
)
messages = [
{
"role": "system",
@@ -58,11 +64,11 @@ async def main(room_url):
@transport.event_handler("on_first_participant_joined")
async def on_first_participant_joined(transport, participant):
await task.queue_frames([LLMMessagesFrame(messages), EndFrame()])
# await task.queue_frames([LLMMessagesFrame(messages), EndFrame()])
await task.queue_frames([LLMMessagesFrame(messages)])
await runner.run(task)
if __name__ == "__main__":
(url, token) = configure()
asyncio.run(main(url))

View File

@@ -72,7 +72,7 @@ async def main(room_url: str, token):
tma_out # Assistant spoken responses
])
task = PipelineTask(pipeline, PipelineParams(allow_interruptions=True))
task = PipelineTask(pipeline, PipelineParams(allow_interruptions=True, enable_metrics=True))
@transport.event_handler("on_first_participant_joined")
async def on_first_participant_joined(transport, participant):

View File

@@ -136,9 +136,15 @@ class LLMService(AIService):
class TTSService(AIService):
def __init__(self, *, aggregate_sentences: bool = True, **kwargs):
def __init__(
self,
*,
aggregate_sentences: bool = True,
push_text_frames: bool = True,
**kwargs):
super().__init__(**kwargs)
self._aggregate_sentences: bool = aggregate_sentences
self._push_text_frames: bool = push_text_frames
self._current_sentence: str = ""
# Converts the text to audio.
@@ -176,9 +182,11 @@ class TTSService(AIService):
await self.process_generator(self.run_tts(text))
await self.stop_processing_metrics()
await self.push_frame(TTSStoppedFrame())
# We send the original text after the audio. This way, if we are
# interrupted, the text is not added to the assistant context.
await self.push_frame(TextFrame(text))
if self._push_text_frames:
print("PUSHING TEXT FRAME")
# We send the original text after the audio. This way, if we are
# interrupted, the text is not added to the assistant context.
await self.push_frame(TextFrame(text))
async def process_frame(self, frame: Frame, direction: FrameDirection):
await super().process_frame(frame, direction)

View File

@@ -8,11 +8,14 @@ import json
import uuid
import base64
import asyncio
import time
from typing import AsyncGenerator
from pipecat.processors.frame_processor import FrameDirection
from pipecat.frames.frames import Frame, AudioRawFrame, StartInterruptionFrame
from pipecat.frames.frames import (
Frame, AudioRawFrame, StartInterruptionFrame, StartFrame, EndFrame, TextFrame
)
from pipecat.services.ai_services import TTSService
from loguru import logger
@@ -39,9 +42,22 @@ class CartesiaTTSService(TTSService):
model_id: str = "sonic-english",
encoding: str = "pcm_s16le",
sample_rate: int = 16000,
language: str = "en",
**kwargs):
super().__init__(**kwargs)
# Aggregating sentences still gives cleaner-sounding results and fewer
# artifacts than streaming one word at a time. On average, waiting for
# a full sentence should only "cost" us 15ms or so with GPT-4o or a Llama 3
# model, and it's worth it for the better audio quality.
self._aggregate_sentences = True
# we don't want to automatically push LLM response text frames, because the
# context aggregators will add them to the LLM context even if we're
# interrupted. cartesia gives us word-by-word timestamps. we can use those
# to generate text frames ourselves aligned with the playout timing of the audio!
self._push_text_frames = False
self._api_key = api_key
self._cartesia_version = cartesia_version
self._url = url
@@ -52,16 +68,32 @@ class CartesiaTTSService(TTSService):
"encoding": encoding,
"sample_rate": sample_rate,
}
self._language = "en"
self._language = language
self._websocket = None
self._context_id = None
self._context_id_start_timestamp = None
self._timestamped_words_buffer = []
self._receive_task = None
self._context_appending_task = None
self._waiting_for_ttfb = False
def can_generate_metrics(self) -> bool:
return True
async def start(self, frame: StartFrame):
await super().start(frame)
await self.connect()
self._context_appending_task = self.get_event_loop().create_task(self._context_appending_task_handler())
async def stop(self, frame: EndFrame):
await super().stop(frame)
await self.disconnect()
if self._context_appending_task:
self._context_appending_task.cancel()
self._context_appending_task = None
pass
async def connect(self):
try:
self._websocket = await websockets.connect(
@@ -69,6 +101,7 @@ class CartesiaTTSService(TTSService):
)
except Exception as e:
logger.exception(f"{self} initialization error: {e}")
self._websocket = None
async def disconnect(self):
try:
@@ -76,38 +109,73 @@ class CartesiaTTSService(TTSService):
ws = self._websocket
self._websocket = None
await ws.close()
self._context_id = None
self._context_id_start_timestamp = None
if self._receive_task:
self._receive_task.cancel()
self._receive_task = None
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()
self._context_id = None
self._context_id_start_timestamp = None
self._timestamped_words_buffer = []
await self.stop_all_metrics()
async def _receive_task_handler(self):
async for message in self._websocket:
msg = json.loads(message)
if not msg:
continue
# logger.debug(f"Received message: {msg}")
if self._waiting_for_ttfb:
await self.stop_ttfb_metrics()
self._waiting_for_ttfb = False
if msg["done"]:
self._context_id = None
if self._receive_task:
self._receive_task.cancel()
self._receive_task = None
return
frame = AudioRawFrame(
audio=base64.b64decode(msg["data"]),
sample_rate=self._output_format["sample_rate"],
num_channels=1
)
await self.push_frame(frame)
try:
async for message in self._websocket:
msg = json.loads(message)
if not msg or msg["context_id"] != self._context_id:
continue
# logger.debug(f"Received message: {msg['type']} {msg['context_id']}")
if msg["type"] == "done":
# unset _context_id but not the _context_id_start_timestamp because we are likely still
# playing out audio and need the timestamp to set send context frames
self._context_id = None
if self._receive_task:
self._receive_task.cancel()
self._receive_task = None
return
if msg["type"] == "timestamps":
# logger.debug(f"TIMESTAMPS: {msg}")
self._timestamped_words_buffer.extend(
list(zip(msg["word_timestamps"]["words"], msg["word_timestamps"]["end"]))
)
continue
if msg["type"] == "chunk":
if not self._context_id_start_timestamp:
self._context_id_start_timestamp = time.time()
if self._waiting_for_ttfb:
await self.stop_ttfb_metrics()
self._waiting_for_ttfb = False
frame = AudioRawFrame(
audio=base64.b64decode(msg["data"]),
sample_rate=self._output_format["sample_rate"],
num_channels=1
)
await self.push_frame(frame)
except Exception as e:
logger.exception(f"{self} exception: {e}")
async def _context_appending_task_handler(self):
try:
while True:
await asyncio.sleep(0.1)
if not self._context_id_start_timestamp:
continue
elapsed_seconds = time.time() - self._context_id_start_timestamp
# pop all words from self._timestamped_words_buffer that are older than the
# elapsed time and print a message about them to the console
# ...
while self._timestamped_words_buffer and self._timestamped_words_buffer[0][1] <= elapsed_seconds:
word, timestamp = self._timestamped_words_buffer.pop(0)
print(f"Word '{word}' with timestamp {timestamp:.2f}s has been spoken.")
await self.push_frame(TextFrame(word))
except Exception as e:
logger.exception(f"{self} exception: {e}")
async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]:
logger.debug(f"Generating TTS: [{text}]")
@@ -124,7 +192,7 @@ class CartesiaTTSService(TTSService):
self._context_id = str(uuid.uuid4())
msg = {
"transcript": text,
"transcript": text + " ",
"continue": True,
"context_id": self._context_id,
"model_id": self._model_id,
@@ -134,8 +202,10 @@ class CartesiaTTSService(TTSService):
},
"output_format": self._output_format,
"language": self._language,
"add_timestamps": True,
}
# logger.debug(f"SENDING MESSAGE {json.dumps(msg)}")
# todo: handle websocket closed
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?