diff --git a/changelog/4464.added.2.md b/changelog/4464.added.2.md new file mode 100644 index 000000000..094948ff7 --- /dev/null +++ b/changelog/4464.added.2.md @@ -0,0 +1 @@ +- Added NVIDIA Magpie TTS services via AWS SageMaker: `NvidiaSageMakerHTTPTTSService` (single HTTP invocation, streams raw PCM back) and `NvidiaSageMakerWebsocketTTSService` (persistent HTTP/2 bidi-stream with full interruption support via `InterruptibleTTSService`). diff --git a/changelog/4464.added.md b/changelog/4464.added.md new file mode 100644 index 000000000..9935c47f0 --- /dev/null +++ b/changelog/4464.added.md @@ -0,0 +1 @@ +- Added `NvidiaSageMakerWebsocketSTTService` for streaming speech recognition using NVIDIA Nemotron ASR via an AWS SageMaker bidirectional-stream endpoint. Produces `InterimTranscriptionFrame` and `TranscriptionFrame` frames, is VAD-aware, and automatically reconnects on error. diff --git a/env.example b/env.example index d449db6b2..6d69cc0e9 100644 --- a/env.example +++ b/env.example @@ -132,6 +132,10 @@ NOVITA_API_KEY=... # NVIDIA NVIDIA_API_KEY=... +# For a full example of how to deploy to SageMaker, see: +# https://github.com/pipecat-ai/pipecat-examples/tree/main/nvidia_sagemaker_example/deployment/aws-sagemaker-nvidia +SAGEMAKER_ASR_ENDPOINT_NAME=... +SAGEMAKER_MAGPIE_ENDPOINT_NAME=... # OpenAI OPENAI_API_KEY=... diff --git a/examples/voice/voice-nvidia-sagemaker.py b/examples/voice/voice-nvidia-sagemaker.py new file mode 100644 index 000000000..ac0c6a365 --- /dev/null +++ b/examples/voice/voice-nvidia-sagemaker.py @@ -0,0 +1,129 @@ +# +# Copyright (c) 2024-2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# +# For a full example of how to deploy to SageMaker, see: +# https://github.com/pipecat-ai/pipecat-examples/tree/main/nvidia_sagemaker_example/deployment/aws-sagemaker-nvidia + +import os + +from dotenv import load_dotenv +from loguru import logger + +from pipecat.audio.vad.silero import SileroVADAnalyzer +from pipecat.frames.frames import LLMRunFrame +from pipecat.pipeline.pipeline import Pipeline +from pipecat.pipeline.runner import PipelineRunner +from pipecat.pipeline.task import PipelineParams, PipelineTask +from pipecat.processors.aggregators.llm_context import LLMContext +from pipecat.processors.aggregators.llm_response_universal import ( + LLMContextAggregatorPair, + LLMUserAggregatorParams, +) +from pipecat.runner.types import RunnerArguments +from pipecat.runner.utils import create_transport +from pipecat.services.nvidia.llm import NvidiaLLMService +from pipecat.services.nvidia.sagemaker.stt import NvidiaSageMakerSTTService +from pipecat.services.nvidia.sagemaker.tts import NvidiaSageMakerTTSService +from pipecat.transports.base_transport import BaseTransport, TransportParams +from pipecat.transports.daily.transport import DailyParams +from pipecat.transports.websocket.fastapi import FastAPIWebsocketParams + +load_dotenv(override=True) + +# We use lambdas to defer transport parameter creation until the transport +# type is selected at runtime. +transport_params = { + "daily": lambda: DailyParams( + audio_in_enabled=True, + audio_out_enabled=True, + ), + "twilio": lambda: FastAPIWebsocketParams( + audio_in_enabled=True, + audio_out_enabled=True, + ), + "webrtc": lambda: TransportParams( + audio_in_enabled=True, + audio_out_enabled=True, + ), +} + + +async def run_bot(transport: BaseTransport, runner_args: RunnerArguments): + logger.info(f"Starting bot") + + stt = NvidiaSageMakerSTTService( + endpoint_name=os.environ["SAGEMAKER_ASR_ENDPOINT_NAME"], + region=os.getenv("AWS_REGION", "us-west-2"), + ) + + llm = NvidiaLLMService( + api_key=os.environ["NVIDIA_API_KEY"], + settings=NvidiaLLMService.Settings( + model="meta/llama-3.3-70b-instruct", + system_instruction="You are a helpful assistant in a voice conversation. Your responses will be spoken aloud, so avoid emojis, bullet points, or other formatting that can't be spoken. Respond to what the user said in a creative, helpful, and brief way.", + ), + ) + + tts = NvidiaSageMakerTTSService( + endpoint_name=os.environ["SAGEMAKER_MAGPIE_ENDPOINT_NAME"], + region=os.getenv("AWS_REGION", "us-west-2"), + ) + + context = LLMContext() + user_aggregator, assistant_aggregator = LLMContextAggregatorPair( + context, + user_params=LLMUserAggregatorParams(vad_analyzer=SileroVADAnalyzer()), + ) + + pipeline = Pipeline( + [ + transport.input(), # Transport user input + stt, # STT + user_aggregator, # User responses + llm, # LLM + tts, # TTS + transport.output(), # Transport bot output + assistant_aggregator, # Assistant spoken responses + ] + ) + + task = PipelineTask( + pipeline, + params=PipelineParams( + enable_metrics=True, + enable_usage_metrics=True, + ), + idle_timeout_secs=runner_args.pipeline_idle_timeout_secs, + ) + + @transport.event_handler("on_client_connected") + async def on_client_connected(transport, client): + logger.info(f"Client connected") + # Kick off the conversation. + context.add_message( + {"role": "developer", "content": "Please introduce yourself to the user."} + ) + await task.queue_frames([LLMRunFrame()]) + + @transport.event_handler("on_client_disconnected") + async def on_client_disconnected(transport, client): + logger.info(f"Client disconnected") + await task.cancel() + + runner = PipelineRunner(handle_sigint=runner_args.handle_sigint) + + await runner.run(task) + + +async def bot(runner_args: RunnerArguments): + """Main bot entry point compatible with Pipecat Cloud.""" + transport = await create_transport(runner_args, transport_params) + await run_bot(transport, runner_args) + + +if __name__ == "__main__": + from pipecat.runner.run import main + + main() diff --git a/src/pipecat/services/aws/sagemaker/bidi_client.py b/src/pipecat/services/aws/sagemaker/bidi_client.py index 8d7bdeaa1..1165bb89b 100644 --- a/src/pipecat/services/aws/sagemaker/bidi_client.py +++ b/src/pipecat/services/aws/sagemaker/bidi_client.py @@ -63,8 +63,8 @@ class SageMakerBidiClient: self, endpoint_name: str, region: str, - model_invocation_path: str = "", - model_query_string: str = "", + model_invocation_path: str | None = "", + model_query_string: str | None = "", ): """Initialize the SageMaker BiDi client. diff --git a/src/pipecat/services/nvidia/sagemaker/__init__.py b/src/pipecat/services/nvidia/sagemaker/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/pipecat/services/nvidia/sagemaker/stt.py b/src/pipecat/services/nvidia/sagemaker/stt.py new file mode 100644 index 000000000..cf6c8c6d8 --- /dev/null +++ b/src/pipecat/services/nvidia/sagemaker/stt.py @@ -0,0 +1,353 @@ +# +# Copyright (c) 2024-2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""NVIDIA Nemotron ASR STT service backed by an AWS SageMaker bidirectional-stream endpoint. + +Uses SageMaker's HTTP/2 bidi-stream API to maintain a persistent connection to +the wrapper's /invocations-bidirectional-stream endpoint, which proxies to NIM's +realtime WebSocket. + +Audio is streamed as base64-encoded PCM16 chunks via input_audio_buffer.append +events. Transcription deltas arrive as InterimTranscriptionFrames and final +results as TranscriptionFrames. + +When the VAD detects the user has stopped speaking, input_audio_buffer.commit +is sent to trigger NIM to finalise the current utterance. +""" + +import asyncio +import base64 +import json +from collections.abc import AsyncGenerator +from dataclasses import dataclass + +from loguru import logger + +from pipecat.frames.frames import ( + CancelFrame, + EndFrame, + ErrorFrame, + Frame, + InterimTranscriptionFrame, + StartFrame, + TranscriptionFrame, + VADUserStartedSpeakingFrame, + VADUserStoppedSpeakingFrame, +) +from pipecat.processors.frame_processor import FrameDirection +from pipecat.services.aws.sagemaker.bidi_client import SageMakerBidiClient +from pipecat.services.settings import STTSettings, assert_given +from pipecat.services.stt_service import STTService +from pipecat.transcriptions.language import Language +from pipecat.utils.time import time_now_iso8601 +from pipecat.utils.tracing.service_decorators import traced_stt + + +@dataclass +class NvidiaSageMakerSTTSettings(STTSettings): + """Settings for NvidiaSageMakerSTTService. + + Parameters: + language: ISO-639-1 language code passed to NIM (e.g. ``en-US``). + """ + + +class NvidiaSageMakerSTTService(STTService): + """NVIDIA Nemotron ASR STT service using SageMaker bidirectional streaming. + + Maintains a persistent HTTP/2 bidi-stream session to the SageMaker endpoint + for the lifetime of the pipeline. Audio chunks are forwarded as base64-encoded + PCM16 via NIM realtime events; transcription results arrive asynchronously and + are pushed as :class:`InterimTranscriptionFrame` and :class:`TranscriptionFrame` + frames. + + Example:: + + stt = NvidiaSageMakerSTTService( + endpoint_name=os.getenv("SAGEMAKER_ASR_ENDPOINT_NAME"), + region=os.getenv("AWS_REGION", "us-west-2"), + settings=NvidiaSageMakerSTTService.Settings( + language="en-US", + ), + ) + """ + + Settings = NvidiaSageMakerSTTSettings + + def __init__( + self, + *, + endpoint_name: str, + region: str = "us-west-2", + sample_rate: int | None = None, + settings: NvidiaSageMakerSTTSettings | None = None, + ttfs_p99_latency: float | None = 1.5, + **kwargs, + ): + """Initialize the SageMaker WebSocket STT service. + + Args: + endpoint_name: Name of the deployed SageMaker endpoint. + region: AWS region where the endpoint lives. + sample_rate: Input sample rate in Hz. Defaults to pipeline rate. + settings: Runtime-updatable settings (language, model). + ttfs_p99_latency: Expected p99 time-to-first-segment latency in seconds. + **kwargs: Forwarded to :class:`STTService`. + """ + default_settings = self.Settings( + model="cache-aware-parakeet-rnnt-en-US-asr-streaming-sortformer", + language="en-US", + ) + + if settings is not None: + default_settings.apply_update(settings) + + super().__init__( + sample_rate=sample_rate, + settings=default_settings, + ttfs_p99_latency=ttfs_p99_latency, + **kwargs, + ) + + self._endpoint_name = endpoint_name + self._region = region + self._client: SageMakerBidiClient | None = None + self._response_task: asyncio.Task | None = None + + def can_generate_metrics(self) -> bool: + """Check if this service can generate processing metrics. + + Returns: + True, as this service supports metrics generation. + """ + return True + + # ── Lifecycle ───────────────────────────────────────────────────────────── + + async def start(self, frame: StartFrame): + """Start the STT service and connect to the SageMaker endpoint. + + Args: + frame: The start frame containing initialization parameters. + """ + await super().start(frame) + await self._connect() + + async def stop(self, frame: EndFrame): + """Stop the STT service and disconnect from the SageMaker endpoint. + + Args: + frame: The end frame. + """ + await super().stop(frame) + await self._disconnect() + + async def cancel(self, frame: CancelFrame): + """Cancel the STT service and disconnect from the SageMaker endpoint. + + Args: + frame: The cancel frame. + """ + await super().cancel(frame) + await self._disconnect() + + # ── Audio input ─────────────────────────────────────────────────────────── + + async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame | None, None]: + """Send an audio chunk to NIM; transcription results arrive asynchronously. + + Each chunk is appended and immediately committed, matching the NVIDIA + reference client pattern for continuous streaming transcription. + """ + if self._client and self._client.is_active: + try: + await self._client.send_json( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(audio).decode(), + } + ) + await self._client.send_json({"type": "input_audio_buffer.commit"}) + except Exception as e: + yield ErrorFrame(error=f"Unknown error occurred: {e}") + yield None + + # ── VAD integration ─────────────────────────────────────────────────────── + + async def process_frame(self, frame: Frame, direction: FrameDirection): + """Process frames with VAD-specific handling for metrics lifecycle. + + Args: + frame: The frame to process. + direction: The direction of frame processing. + """ + await super().process_frame(frame, direction) + + if isinstance(frame, VADUserStartedSpeakingFrame): + logger.debug(f"{self}: VAD user started speaking") + await self.start_processing_metrics() + if isinstance(frame, VADUserStoppedSpeakingFrame): + logger.debug(f"{self}: VAD user stopped speaking") + + # ── Connection management ───────────────────────────────────────────────── + + async def _open_client_session(self): + self._client = SageMakerBidiClient( + endpoint_name=self._endpoint_name, + region=self._region, + model_query_string=None, + model_invocation_path=None, + ) + await self._client.start_session() + await self._send_session_config() + + async def _close_client_session(self): + if self._client and self._client.is_active: + try: + await self._client.send_json({"type": "session.end"}) + except Exception as e: + logger.warning(f"{self}: error sending session.end: {e}") + await self._client.close_session() + self._client = None + + async def _connect(self): + logger.debug( + f"{self}: connecting to SageMaker bidi-stream endpoint '{self._endpoint_name}'" + ) + try: + await self._open_client_session() + self._response_task = self.create_task(self._process_responses()) + logger.debug(f"{self}: connected") + await self._call_event_handler("on_connected") + except Exception as e: + logger.error(f"{self}: connection error: {e}") + self._client = None + await self._call_event_handler("on_connection_error", f"{e}") + + async def _disconnect(self): + if self._response_task and not self._response_task.done(): + await self.cancel_task(self._response_task) + self._response_task = None + await self._close_client_session() + await self._call_event_handler("on_disconnected") + + async def _do_reconnect(self): + await self._close_client_session() + await self._open_client_session() + + async def _send_session_config(self): + """Send transcription_session.update to configure audio format and params. + + Specifies ``"model": "nemotron-asr-streaming"`` in ``input_audio_transcription`` so + NIM selects the correct Nemotron ASR Streaming model. + """ + logger.debug( + f"{self}: sending session config," + f" sample_rate={self.sample_rate} language={self._settings.language}" + ) + assert self._client is not None + await self._client.send_json( + { + "type": "transcription_session.update", + "session": { + "input_audio_format": "pcm16", + "input_audio_params": { + "sample_rate_hz": self.sample_rate, + "num_channels": 1, + }, + "input_audio_transcription": { + "language": self._settings.language, + "model": self._settings.model, + }, + "recognition_config": { + "enable_automatic_punctuation": True, + }, + }, + } + ) + + # ── Response processing ─────────────────────────────────────────────────── + + async def _process_responses(self): + """Receive NIM JSON events and push transcription frames.""" + try: + while self._client and self._client.is_active: + result = await self._client.receive_response() + + if result is None or not ( + hasattr(result, "value") and hasattr(result.value, "bytes_") # type: ignore[union-attr] + ): + continue + + payload = result.value.bytes_ # type: ignore[union-attr] + if not payload: + continue + + try: + msg = json.loads(payload.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError): + continue + + event_type = msg.get("type", "") + + if event_type not in ( + "conversation.item.input_audio_transcription.delta", + "input_audio_buffer.committed", + ): + logger.debug(f"{self}: received event: {event_type}") + + _lang = assert_given(self._settings.language) + language: Language | None = Language(_lang) if _lang is not None else None + + if event_type == "conversation.item.input_audio_transcription.delta": + delta = msg.get("delta", "") + if delta: + logger.debug(f"{self}: received transcription delta: {delta}") + await self.push_frame( + InterimTranscriptionFrame( + delta, + self._user_id, + time_now_iso8601(), + language=language, + result=msg, + ) + ) + + elif event_type == "conversation.item.input_audio_transcription.completed": + transcript = msg.get("transcript", "") + if transcript.strip(): + logger.debug(f"{self}: received final transcription: {transcript}") + await self.push_frame( + TranscriptionFrame( + transcript, + self._user_id, + time_now_iso8601(), + language=language, + result=msg, + finalized=True, + ) + ) + await self._handle_transcription(transcript, True) + await self.stop_processing_metrics() + + elif event_type in ( + "conversation.item.input_audio_transcription.failed", + "error", + ): + await self.push_error(error_msg=f"NIM ASR error: {msg}") + # In case of error we need to reconnect, otherwise we are not going to receive from the STT service anymore + await self._request_reconnect() + + except asyncio.CancelledError: + logger.debug(f"{self}: response processor cancelled") + except Exception as e: + await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e) + finally: + logger.debug(f"{self}: response processor stopped") + + @traced_stt + async def _handle_transcription(self, transcript: str, is_final: bool, language=None): + pass diff --git a/src/pipecat/services/nvidia/sagemaker/tts.py b/src/pipecat/services/nvidia/sagemaker/tts.py new file mode 100644 index 000000000..e6737a5f9 --- /dev/null +++ b/src/pipecat/services/nvidia/sagemaker/tts.py @@ -0,0 +1,496 @@ +# +# Copyright (c) 2024-2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""NVIDIA Magpie TTS service backed by an AWS SageMaker endpoint.""" + +import asyncio +import base64 +import json +import os +from collections.abc import AsyncGenerator +from dataclasses import dataclass + +import aioboto3 +from loguru import logger + +from pipecat.frames.frames import ( + CancelFrame, + EndFrame, + ErrorFrame, + Frame, + InterruptionFrame, + StartFrame, + TTSAudioRawFrame, +) +from pipecat.processors.frame_processor import FrameDirection +from pipecat.services.aws.sagemaker.bidi_client import SageMakerBidiClient +from pipecat.services.settings import TTSSettings +from pipecat.services.tts_service import InterruptibleTTSService, TTSService +from pipecat.utils.tracing.service_decorators import traced_tts + + +@dataclass +class NvidiaSageMakerTTSSettings(TTSSettings): + """Settings for NVIDIA SageMaker TTS services. + + Parameters: + voice: NIM voice name (e.g. ``Magpie-Multilingual.EN-US.Aria``). + language: BCP-47 language code passed to NIM (e.g. ``en-US``). + """ + + +class NvidiaSageMakerHTTPTTSService(TTSService): + """NVIDIA Magpie TTS service that calls a SageMaker HTTP endpoint. + + Sends each text segment to the wrapper's ``POST /invocations`` endpoint + as a JSON body and streams the raw PCM audio response back to bot + as :class:`TTSAudioRawFrame` frames. + + Example:: + + tts = NvidiaSageMakerHTTPTTSService( + endpoint_name=os.getenv("SAGEMAKER_MAGPIE_ENDPOINT_NAME"), + region=os.getenv("AWS_REGION", "us-west-2"), + settings=NvidiaSageMakerHTTPTTSService.Settings( + voice="Magpie-Multilingual.EN-US.Aria", + language="en-US", + ), + ) + """ + + Settings = NvidiaSageMakerTTSSettings + + def __init__( + self, + *, + endpoint_name: str, + region: str = "us-west-2", + sample_rate: int | None = None, + settings: NvidiaSageMakerTTSSettings | None = None, + **kwargs, + ): + """Initialize the SageMaker HTTP TTS service. + + Args: + endpoint_name: Name of the deployed SageMaker endpoint. + region: AWS region where the endpoint lives. + sample_rate: Output sample rate in Hz. Defaults to bot's pipeline rate. + settings: Runtime-updatable settings (voice, language). + **kwargs: Forwarded to :class:`TTSService`. + """ + default_settings = self.Settings( + model="magpie", + voice="Magpie-Multilingual.EN-US.Aria", + language="en-US", + ) + + if settings is not None: + default_settings.apply_update(settings) + + super().__init__( + sample_rate=sample_rate, + push_start_frame=True, + push_stop_frames=True, + settings=default_settings, + **kwargs, + ) + + self._endpoint_name = endpoint_name + self._region = region + self._client = None + self._client_ctx = None + + def can_generate_metrics(self) -> bool: + """Check if this service can generate processing metrics. + + Returns: + True, as this service supports metrics generation. + """ + return True + + # ── Lifecycle ───────────────────────────────────────────────────────────── + + async def start(self, frame: StartFrame): + """Start the TTS service and create the SageMaker client. + + Args: + frame: The start frame containing initialization parameters. + """ + await super().start(frame) + session = aioboto3.Session( + aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"), + aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"), + region_name=self._region, + ) + self._client_ctx = session.client("sagemaker-runtime") + self._client = await self._client_ctx.__aenter__() + logger.debug(f"{self}: connected to SageMaker endpoint '{self._endpoint_name}'") + + async def _close_client(self): + if self._client_ctx is not None: + try: + await self._client_ctx.__aexit__(None, None, None) + except Exception as e: + logger.warning(f"{self}: error closing SageMaker client: {e}") + self._client_ctx = None + self._client = None + + async def stop(self, frame: EndFrame): + """Stop the TTS service and close the SageMaker client. + + Args: + frame: The end frame. + """ + await super().stop(frame) + await self._close_client() + + async def cancel(self, frame: CancelFrame): + """Cancel the TTS service and close the SageMaker client. + + Args: + frame: The cancel frame. + """ + await super().cancel(frame) + await self._close_client() + + # ── Synthesis ───────────────────────────────────────────────────────────── + + @traced_tts + async def run_tts(self, text: str, context_id: str) -> AsyncGenerator[Frame, None]: + """Synthesize text via SageMaker and yield a single PCM audio frame. + + Args: + text: The text to synthesize. + context_id: Pipecat audio context identifier. + + Yields: + :class:`TTSAudioRawFrame` chunks of signed 16-bit mono PCM. + """ + logger.debug(f"{self}: Generating TTS [{text}]") + + text = text.strip() + if not text or not any(c.isalnum() for c in text): + return + + try: + assert self._client is not None + body = json.dumps( + { + "text": text, + "voice_name": self._settings.voice, + "language_code": self._settings.language, + "sample_rate_hz": self.sample_rate, + } + ) + + response = await self._client.invoke_endpoint( + EndpointName=self._endpoint_name, + ContentType="application/json", + Accept="application/octet-stream", + Body=body, + ) + + if "Body" not in response: + yield ErrorFrame(error="SageMaker TTS returned no audio stream") + return + + first_chunk = True + async for chunk in response["Body"].iter_chunks(chunk_size=self.chunk_size): + if chunk: + if first_chunk: + await self.stop_ttfb_metrics() + first_chunk = False + yield TTSAudioRawFrame( + audio=chunk, + sample_rate=self.sample_rate, + num_channels=1, + context_id=context_id, + ) + except Exception as e: + logger.error(f"{self}: SageMaker TTS error: {e}") + yield ErrorFrame(error=f"SageMaker TTS error: {e}") + + await self.start_tts_usage_metrics(text) + + +class NvidiaSageMakerTTSService(InterruptibleTTSService): + """NVIDIA Magpie TTS service using SageMaker bidirectional streaming. + + Maintains a persistent HTTP/2 bidi-stream session to the SageMaker endpoint + for the lifetime of the pipeline. Each text segment is sent as NIM realtime + events; audio chunks arrive asynchronously and are pushed as + :class:`TTSAudioRawFrame` frames. + + Example:: + + tts = NvidiaSageMakerTTSService( + endpoint_name=os.getenv("SAGEMAKER_MAGPIE_ENDPOINT_NAME"), + region=os.getenv("AWS_REGION", "us-west-2"), + settings=NvidiaSageMakerTTSService.Settings( + voice="Magpie-Multilingual.EN-US.Aria", + language="en-US", + ), + ) + """ + + Settings = NvidiaSageMakerTTSSettings + + def __init__( + self, + *, + endpoint_name: str, + region: str = "us-west-2", + sample_rate: int | None = None, + settings: NvidiaSageMakerTTSSettings | None = None, + **kwargs, + ): + """Initialize the SageMaker WebSocket TTS service. + + Args: + endpoint_name: Name of the deployed SageMaker endpoint. + region: AWS region where the endpoint lives. + sample_rate: Output sample rate in Hz. Defaults to pipeline rate. + settings: Runtime-updatable settings (voice, language). + **kwargs: Forwarded to :class:`InterruptibleTTSService`. + """ + default_settings = self.Settings( + model="magpie", + voice="Magpie-Multilingual.EN-US.Aria", + language="en-US", + ) + + if settings is not None: + default_settings.apply_update(settings) + + super().__init__( + sample_rate=sample_rate, + push_start_frame=True, + push_stop_frames=True, + pause_frame_processing=True, + append_trailing_space=True, + settings=default_settings, + **kwargs, + ) + + self._endpoint_name = endpoint_name + self._region = region + self._client: SageMakerBidiClient | None = None + self._receive_task = None + self._speech_completed_event = asyncio.Event() + + def can_generate_metrics(self) -> bool: + """Check if this service can generate processing metrics. + + Returns: + True, as this service supports metrics generation. + """ + return True + + # ── Lifecycle ───────────────────────────────────────────────────────────── + + async def start(self, frame: StartFrame): + """Start the TTS service and connect to the SageMaker endpoint. + + Args: + frame: The start frame containing initialization parameters. + """ + await super().start(frame) + await self._connect() + + async def stop(self, frame: EndFrame): + """Stop the TTS service and disconnect from the SageMaker endpoint. + + Args: + frame: The end frame. + """ + await super().stop(frame) + await self._disconnect() + + async def cancel(self, frame: CancelFrame): + """Cancel the TTS service and disconnect from the SageMaker endpoint. + + Args: + frame: The cancel frame. + """ + await super().cancel(frame) + await self._disconnect() + + # ── Connection management (WebsocketService abstract interface) ──────────── + + async def _connect(self): + await super()._connect() + await self._connect_websocket() + if self._client and self._client.is_active and not self._receive_task: + self._receive_task = self.create_task(self._receive_task_handler(self._report_error)) + + async def _disconnect(self): + await super()._disconnect() + if self._receive_task: + await self.cancel_task(self._receive_task) + self._receive_task = None + await self._disconnect_websocket() + + async def _connect_websocket(self): + if self._client and self._client.is_active: + return + + logger.debug( + f"{self}: connecting to SageMaker bidi-stream endpoint '{self._endpoint_name}'" + ) + try: + self._client = SageMakerBidiClient( + endpoint_name=self._endpoint_name, + region=self._region, + model_query_string=None, + model_invocation_path=None, + ) + await self._client.start_session() + await self._send_session_config() + logger.debug(f"{self}: connected") + await self._call_event_handler("on_connected") + except Exception as e: + logger.error(f"{self}: connection error: {e}") + self._client = None + await self._call_event_handler("on_connection_error", f"{e}") + + async def _disconnect_websocket(self): + try: + if self._client and self._client.is_active: + logger.debug(f"{self}: disconnecting") + try: + await self._client.send_json({"type": "session.end"}) + except Exception as e: + logger.warning(f"{self}: error sending session.end: {e}") + await self._client.close_session() + logger.debug(f"{self}: disconnected") + except Exception as e: + logger.warning(f"{self}: error during disconnect: {e}") + finally: + self._client = None + await self._call_event_handler("on_disconnected") + + async def _verify_connection(self): + active = self._client and self._client.is_active + logger.info(f"{self}: verifying if websocket connection is active {active}") + return active + + async def _handle_interruption(self, frame: InterruptionFrame, direction: FrameDirection): + if self._bot_speaking and self._client: + logger.debug( + f"{self}: interruption detected, sending input_text.done and waiting for speech.completed" + ) + self._disconnecting = True + self._speech_completed_event.clear() + try: + await self._client.send_json({"type": "input_text.done"}) + await asyncio.wait_for(self._speech_completed_event.wait(), timeout=5.0) + except TimeoutError: + logger.warning(f"{self}: timed out waiting for conversation.item.speech.completed") + await super()._handle_interruption(frame, direction) + + async def _receive_messages(self): + """Receive NIM JSON events and push audio frames.""" + while self._client and self._client.is_active and not self._disconnecting: + result = await self._client.receive_response() + + if self._disconnecting: + self._speech_completed_event.set() + + if result is None: + break + + if not (hasattr(result, "value") and hasattr(result.value, "bytes_")): # type: ignore[union-attr] + continue + + payload = result.value.bytes_ # type: ignore[union-attr] + if not payload: + continue + + context_id = self.get_active_audio_context_id() + + try: + msg = json.loads(payload.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError): + # Unexpected binary frame — treat as raw PCM + await self.push_frame( + TTSAudioRawFrame( + audio=payload, + sample_rate=self.sample_rate, + num_channels=1, + context_id=context_id, + ) + ) + continue + + event_type = msg.get("type", "") + + if event_type != "conversation.item.speech.data": + logger.debug(f"{self}: received event: {event_type}") + + if event_type == "conversation.item.speech.data": + chunk_b64 = msg.get("audio", "") + if chunk_b64: + await self.stop_ttfb_metrics() + await self.push_frame( + TTSAudioRawFrame( + audio=base64.b64decode(chunk_b64), + sample_rate=self.sample_rate, + num_channels=1, + context_id=context_id, + ) + ) + elif event_type == "error": + await self.push_error(error_msg=f"NIM error: {msg.get('message', msg)}") + # In case of error we need to reconnect, otherwise we are not going to receive audio from the TTS service anymore + break + elif event_type == "conversation.item.speech.completed": + # Need to reconnect to reset the synthesis state and be able to synthesize new text + break + + # synthesize_session.updated, input_text.committed, etc. are ignored. + + async def _send_session_config(self): + """Send synthesize_session.update to configure voice and audio params.""" + logger.debug(f"{self}: sending session config, sample_rate={self.sample_rate}") + assert self._client is not None + await self._client.send_json( + { + "type": "synthesize_session.update", + "session": { + "input_text_synthesis": { + "voice_name": self._settings.voice, + "language_code": self._settings.language, + }, + "output_audio_params": { + "sample_rate_hz": self.sample_rate, + }, + }, + } + ) + + # ── Synthesis ───────────────────────────────────────────────────────────── + + @traced_tts + async def run_tts(self, text: str, context_id: str) -> AsyncGenerator[Frame | None, None]: + """Send text to NIM; audio arrives asynchronously via _receive_messages.""" + logger.debug(f"{self}: Generating TTS [{text}]") + + text = text.strip() + if not text or not any(c.isalnum() for c in text): + return + + try: + if not self._client or not self._client.is_active: + await self._connect() + + assert self._client is not None + await self._client.send_json({"type": "input_text.append", "text": text}) + await self._client.send_json({"type": "input_text.commit"}) + await self.start_tts_usage_metrics(text) + yield None + except Exception as e: + logger.error(f"{self}: TTS error: {e}") + yield ErrorFrame(error=f"NvidiaSageMakerTTSService error: {e}")