Nvidia Sagemaker Nemotron ASR STT service

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