diff --git a/src/pipecat/services/nvidia/sagemaker/stt.py b/src/pipecat/services/nvidia/sagemaker/stt.py new file mode 100644 index 000000000..6ab299bf9 --- /dev/null +++ b/src/pipecat/services/nvidia/sagemaker/stt.py @@ -0,0 +1,305 @@ +# +# 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 typing import Optional + +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 +from pipecat.services.stt_service import STTService +from pipecat.utils.time import time_now_iso8601 +from pipecat.utils.tracing.service_decorators import traced_stt + + +@dataclass +class NvidiaSageMakerWSSTTSettings(STTSettings): + """Settings for NvidiaSageMakerWebsocketSTTService. + + Parameters: + language: ISO-639-1 language code passed to NIM (e.g. ``en``). + """ + + language: str = "en-US" + + +class NvidiaSageMakerWebsocketSTTService(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 = NvidiaSageMakerWebsocketSTTService( + endpoint_name=os.getenv("SAGEMAKER_ASR_ENDPOINT_NAME"), + region=os.getenv("AWS_REGION", "us-west-2"), + settings=NvidiaSageMakerWebsocketSTTService.Settings( + language="en-US", + ), + ) + """ + + Settings = NvidiaSageMakerWSSTTSettings + + def __init__( + self, + *, + endpoint_name: str, + region: str = "us-west-2", + sample_rate: int | None = None, + settings: NvidiaSageMakerWSSTTSettings | None = None, + ttfs_p99_latency: float | None = 1.5, + **kwargs, + ): + default_settings = self.Settings( + model="cache-aware-parakeet-rnnt-en-US-asr-streaming-sortformer" + ) + + 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: + return True + + # ── Lifecycle ───────────────────────────────────────────────────────────── + + 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() + + # ── Audio input ─────────────────────────────────────────────────────────── + + async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, 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): + 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 _connect(self): + 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() + 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 + + 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") + + self._client = None + await self._call_event_handler("on_disconnected") + + 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}" + ) + 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_") + ): + continue + + payload = result.value.bytes_ + 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}") + + 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(), + ) + ) + + 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(), + result=msg, + ) + ) + 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._disconnect() + await self._connect() + + 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