RimeTTSService: use AudioContextWordTTSService

This commit is contained in:
Aleix Conchillo Flaqué
2025-02-14 11:39:05 -08:00
parent aeadb40c3f
commit f53ee79ddb

View File

@@ -7,7 +7,7 @@
import base64 import base64
import json import json
import uuid import uuid
from typing import AsyncGenerator, Optional, Union from typing import AsyncGenerator, Optional
import aiohttp import aiohttp
from loguru import logger from loguru import logger
@@ -28,7 +28,7 @@ from pipecat.frames.frames import (
TTSStoppedFrame, TTSStoppedFrame,
) )
from pipecat.processors.frame_processor import FrameDirection from pipecat.processors.frame_processor import FrameDirection
from pipecat.services.ai_services import TTSService, WordTTSService from pipecat.services.ai_services import AudioContextWordTTSService, TTSService
from pipecat.services.websocket_service import WebsocketService from pipecat.services.websocket_service import WebsocketService
from pipecat.transcriptions.language import Language from pipecat.transcriptions.language import Language
@@ -58,7 +58,7 @@ def language_to_rime_language(language: Language) -> str:
return LANGUAGE_MAP.get(language, "eng") return LANGUAGE_MAP.get(language, "eng")
class RimeTTSService(WordTTSService, WebsocketService): class RimeTTSService(AudioContextWordTTSService, WebsocketService):
"""Text-to-Speech service using Rime's websocket API. """Text-to-Speech service using Rime's websocket API.
Uses Rime's websocket JSON API to convert text to speech with word-level timing Uses Rime's websocket JSON API to convert text to speech with word-level timing
@@ -95,7 +95,7 @@ class RimeTTSService(WordTTSService, WebsocketService):
params: Additional configuration parameters. params: Additional configuration parameters.
""" """
# Initialize with parent class settings for proper frame handling # Initialize with parent class settings for proper frame handling
WordTTSService.__init__( AudioContextWordTTSService.__init__(
self, self,
aggregate_sentences=True, aggregate_sentences=True,
push_text_frames=False, push_text_frames=False,
@@ -249,12 +249,18 @@ class RimeTTSService(WordTTSService, WebsocketService):
return word_pairs return word_pairs
async def flush_audio(self):
if not self._context_id or not self._websocket:
return
logger.trace(f"{self}: flushing audio")
self._context_id = None
async def _receive_messages(self): async def _receive_messages(self):
"""Process incoming websocket messages.""" """Process incoming websocket messages."""
async for message in self._get_websocket(): async for message in self._get_websocket():
msg = json.loads(message) msg = json.loads(message)
if not msg or msg["contextId"] != self._context_id: if not msg or not self.audio_context_available(msg["contextId"]):
continue continue
if msg["type"] == "chunk": if msg["type"] == "chunk":
@@ -266,7 +272,7 @@ class RimeTTSService(WordTTSService, WebsocketService):
sample_rate=self.sample_rate, sample_rate=self.sample_rate,
num_channels=1, num_channels=1,
) )
await self.push_frame(frame) await self.append_to_audio_context(msg["contextId"], frame)
elif msg["type"] == "timestamps": elif msg["type"] == "timestamps":
# Process word timing information # Process word timing information
@@ -288,6 +294,7 @@ class RimeTTSService(WordTTSService, WebsocketService):
await self.push_frame(TTSStoppedFrame()) await self.push_frame(TTSStoppedFrame())
await self.stop_all_metrics() await self.stop_all_metrics()
await self.push_error(ErrorFrame(f"{self} error: {msg['message']}")) await self.push_error(ErrorFrame(f"{self} error: {msg['message']}"))
self._context_id = None
async def push_frame(self, frame: Frame, direction: FrameDirection = FrameDirection.DOWNSTREAM): async def push_frame(self, frame: Frame, direction: FrameDirection = FrameDirection.DOWNSTREAM):
"""Push frame and handle end-of-turn conditions.""" """Push frame and handle end-of-turn conditions."""
@@ -329,6 +336,7 @@ class RimeTTSService(WordTTSService, WebsocketService):
self._started = True self._started = True
self._cumulative_time = 0 self._cumulative_time = 0
self._context_id = str(uuid.uuid4()) self._context_id = str(uuid.uuid4())
await self.create_audio_context(self._context_id)
msg = self._build_msg(text=text) msg = self._build_msg(text=text)
await self._get_websocket().send(json.dumps(msg)) await self._get_websocket().send(json.dumps(msg))