430 lines
12 KiB
Python
430 lines
12 KiB
Python
"""
|
|
NVIDIA realtime STT provider for the Voice Assistant.
|
|
|
|
Purpose:
|
|
Implement the generic STTAdapter contract using NVIDIA Speech NIM.
|
|
|
|
Responsibilities:
|
|
- Create and manage NVIDIA realtime transcription sessions.
|
|
- Send canonical PCM audio to the NVIDIA ASR service.
|
|
- Receive provider-specific transcription events.
|
|
- Translate NVIDIA events into provider-independent TranscriptEvent objects.
|
|
- Manage provider connection lifecycle and cleanup.
|
|
|
|
Non-responsibilities:
|
|
- Managing WebSocket client connections to the Audio Gateway.
|
|
- Managing Voice Assistant sessions or turns.
|
|
- Executing Agent or TTS functionality.
|
|
- Defining the canonical Voice Core audio format.
|
|
|
|
Architecture:
|
|
Voice Core
|
|
|
|
|
| canonical PCM audio
|
|
v
|
|
STTAdapter
|
|
|
|
|
v
|
|
NvidiaSTTAdapter
|
|
|
|
|
v
|
|
NVIDIA Speech NIM
|
|
|
|
NVIDIA-specific HTTP, WebSocket, session, and event details remain
|
|
isolated inside this adapter.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
import httpx
|
|
import websockets
|
|
from websockets.asyncio.client import ClientConnection
|
|
|
|
from voice_assistant.audio.models import CANONICAL_AUDIO_FORMAT
|
|
|
|
from .adapter import STTAdapter
|
|
from .models import TranscriptEvent
|
|
|
|
|
|
class NvidiaSTTAdapter(STTAdapter):
|
|
"""Realtime STT adapter for NVIDIA Speech NIM."""
|
|
|
|
def __init__(
|
|
self,
|
|
server: str = "127.0.0.1:9000",
|
|
language_code: str = "pt-BR",
|
|
automatic_punctuation: bool = True,
|
|
) -> None:
|
|
self.server = server
|
|
self.language_code = language_code
|
|
self.automatic_punctuation = automatic_punctuation
|
|
|
|
self._http: httpx.AsyncClient | None = None
|
|
self._websocket: ClientConnection | None = None
|
|
self._receiver_task: asyncio.Task[None] | None = None
|
|
|
|
self._events: asyncio.Queue[TranscriptEvent | None] = asyncio.Queue()
|
|
|
|
self._started = False
|
|
self._stopping = False
|
|
self._aborted = False
|
|
self._abort_lock = asyncio.Lock()
|
|
|
|
|
|
async def start_session(self) -> None:
|
|
if self._started:
|
|
raise RuntimeError("STT session is already started.")
|
|
|
|
self._http = httpx.AsyncClient(timeout=10.0)
|
|
|
|
try:
|
|
session = await self._create_session()
|
|
self._websocket = await self._connect()
|
|
|
|
await self._initialize_session(session)
|
|
|
|
self._started = True
|
|
self._stopping = False
|
|
self._aborted = False
|
|
|
|
self._receiver_task = asyncio.create_task(
|
|
self._receive_events()
|
|
)
|
|
|
|
except Exception:
|
|
await self._cleanup()
|
|
raise
|
|
|
|
async def send_audio(self, audio: bytes) -> None:
|
|
if not self._started or self._websocket is None:
|
|
raise RuntimeError("STT session has not been started.")
|
|
|
|
if not audio:
|
|
raise ValueError("Audio chunk cannot be empty.")
|
|
|
|
await self._send_json(
|
|
{
|
|
"type": "input_audio_buffer.append",
|
|
"audio": base64.b64encode(audio).decode("utf-8"),
|
|
}
|
|
)
|
|
|
|
await self._send_json(
|
|
{
|
|
"type": "input_audio_buffer.commit",
|
|
}
|
|
)
|
|
|
|
async def stop_session(self) -> None:
|
|
if not self._started or self._websocket is None:
|
|
return
|
|
|
|
self._stopping = True
|
|
|
|
try:
|
|
await self._send_json(
|
|
{
|
|
"type": "input_audio_buffer.done",
|
|
}
|
|
)
|
|
|
|
if self._receiver_task is not None:
|
|
await self._receiver_task
|
|
|
|
finally:
|
|
await self._cleanup()
|
|
|
|
|
|
async def abort_session(self) -> None:
|
|
"""
|
|
Abort the NVIDIA transcription session safely.
|
|
|
|
Purpose:
|
|
Stop transcription immediately without waiting for a final
|
|
transcript, including when multiple abort requests overlap.
|
|
|
|
Responsibilities:
|
|
- Serialize concurrent abort requests.
|
|
- Stop accepting new audio.
|
|
- Cancel and await the receiver task.
|
|
- Release WebSocket and HTTP resources.
|
|
- Discard pending transcription events.
|
|
- Publish one end-of-stream marker.
|
|
|
|
Non-responsibilities:
|
|
- Managing Gateway WebSocket connections.
|
|
- Updating VoicePipeline turn state.
|
|
- Cancelling Agent or TTS execution.
|
|
"""
|
|
async with self._abort_lock:
|
|
# A completed abort must not run cleanup again.
|
|
if self._aborted:
|
|
return
|
|
|
|
# Mark the session as aborted before releasing resources.
|
|
self._aborted = True
|
|
self._stopping = True
|
|
self._started = False
|
|
|
|
try:
|
|
await self._cleanup()
|
|
|
|
finally:
|
|
# Remove transcripts and termination markers that may
|
|
# have been queued before the abort completed.
|
|
while True:
|
|
try:
|
|
self._events.get_nowait()
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
|
|
# Publish one termination marker for the event consumer.
|
|
self._events.put_nowait(None)
|
|
|
|
|
|
def events(self) -> AsyncIterator[TranscriptEvent]:
|
|
return self._event_stream()
|
|
|
|
async def _event_stream(
|
|
self,
|
|
) -> AsyncIterator[TranscriptEvent]:
|
|
while True:
|
|
event = await self._events.get()
|
|
|
|
if event is None:
|
|
break
|
|
|
|
yield event
|
|
|
|
async def _create_session(self) -> dict[str, Any]:
|
|
if self._http is None:
|
|
raise RuntimeError("HTTP client is not initialized.")
|
|
|
|
response = await self._http.post(
|
|
f"http://{self.server}/v1/realtime/transcription_sessions"
|
|
)
|
|
|
|
response.raise_for_status()
|
|
|
|
session = response.json()
|
|
|
|
if not isinstance(session, dict):
|
|
raise RuntimeError(
|
|
"NVIDIA NIM returned an invalid session response."
|
|
)
|
|
|
|
return session
|
|
|
|
async def _connect(self) -> ClientConnection:
|
|
return await websockets.connect(
|
|
f"ws://{self.server}"
|
|
"/v1/realtime"
|
|
"?intent=transcription"
|
|
)
|
|
|
|
async def _initialize_session(
|
|
self,
|
|
session: dict[str, Any],
|
|
) -> None:
|
|
if self._websocket is None:
|
|
raise RuntimeError("NVIDIA ASR WebSocket is not connected.")
|
|
|
|
response = await self._websocket.recv()
|
|
event = self._decode_event(response)
|
|
|
|
if event.get("type") != "conversation.created":
|
|
raise RuntimeError(
|
|
"Unexpected first NIM event: "
|
|
f"{event.get('type')!r}"
|
|
)
|
|
|
|
session["input_audio_format"] = "pcm16"
|
|
|
|
input_audio_transcription = dict(
|
|
session.get("input_audio_transcription", {})
|
|
)
|
|
|
|
input_audio_transcription["language"] = self.language_code
|
|
session["input_audio_transcription"] = (
|
|
input_audio_transcription
|
|
)
|
|
|
|
session["input_audio_params"] = {
|
|
"sample_rate_hz": CANONICAL_AUDIO_FORMAT.sample_rate,
|
|
"num_channels": CANONICAL_AUDIO_FORMAT.channels,
|
|
}
|
|
|
|
recognition_config = dict(
|
|
session.get("recognition_config", {})
|
|
)
|
|
|
|
recognition_config[
|
|
"enable_automatic_punctuation"
|
|
] = self.automatic_punctuation
|
|
|
|
session["recognition_config"] = recognition_config
|
|
|
|
await self._send_json(
|
|
{
|
|
"type": "transcription_session.update",
|
|
"session": session,
|
|
}
|
|
)
|
|
|
|
response = await self._websocket.recv()
|
|
event = self._decode_event(response)
|
|
|
|
if event.get("type") != "transcription_session.updated":
|
|
raise RuntimeError(
|
|
"NVIDIA NIM did not confirm session update: "
|
|
f"{event.get('type')!r}"
|
|
)
|
|
|
|
async def _receive_events(self) -> None:
|
|
websocket = self._websocket
|
|
|
|
if websocket is None:
|
|
return
|
|
|
|
try:
|
|
async for raw_message in websocket:
|
|
# Stop processing provider events after an explicit abort.
|
|
if self._aborted:
|
|
break
|
|
event = self._decode_event(raw_message)
|
|
event_type = event.get("type", "")
|
|
|
|
if (
|
|
event_type
|
|
== "conversation.item.input_audio_transcription.delta"
|
|
):
|
|
text = event.get("delta", "")
|
|
|
|
if text:
|
|
await self._events.put(
|
|
TranscriptEvent(
|
|
type="partial",
|
|
text=text,
|
|
)
|
|
)
|
|
|
|
elif (
|
|
event_type
|
|
== "conversation.item.input_audio_transcription.completed"
|
|
):
|
|
text = event.get("transcript", "")
|
|
|
|
if not text:
|
|
continue
|
|
|
|
is_last_result = event.get(
|
|
"is_last_result",
|
|
False,
|
|
)
|
|
|
|
await self._events.put(
|
|
TranscriptEvent(
|
|
type=(
|
|
"final"
|
|
if is_last_result
|
|
else "partial"
|
|
),
|
|
text=text,
|
|
)
|
|
)
|
|
|
|
if is_last_result:
|
|
break
|
|
|
|
elif event_type.lower().startswith("error"):
|
|
error = event.get("error", {})
|
|
|
|
raise RuntimeError(
|
|
error.get(
|
|
"message",
|
|
"Unknown NVIDIA ASR error.",
|
|
)
|
|
)
|
|
|
|
except websockets.ConnectionClosed as exc:
|
|
if not self._stopping:
|
|
raise RuntimeError(
|
|
"NVIDIA ASR WebSocket closed unexpectedly."
|
|
) from exc
|
|
|
|
finally:
|
|
if not self._aborted:
|
|
await self._events.put(None)
|
|
|
|
async def _send_json(self, message: dict[str, Any]) -> None:
|
|
if self._websocket is None:
|
|
raise RuntimeError(
|
|
"NVIDIA ASR WebSocket is not connected."
|
|
)
|
|
|
|
await self._websocket.send(
|
|
json.dumps(
|
|
message,
|
|
separators=(",", ":"),
|
|
)
|
|
)
|
|
|
|
@staticmethod
|
|
def _decode_event(
|
|
message: str | bytes,
|
|
) -> dict[str, Any]:
|
|
if isinstance(message, bytes):
|
|
message = message.decode("utf-8")
|
|
|
|
try:
|
|
event = json.loads(message)
|
|
except json.JSONDecodeError as exc:
|
|
raise RuntimeError(
|
|
"Received invalid JSON from NVIDIA ASR."
|
|
) from exc
|
|
|
|
if not isinstance(event, dict):
|
|
raise RuntimeError(
|
|
"Received invalid event from NVIDIA ASR."
|
|
)
|
|
|
|
return event
|
|
|
|
async def _cleanup(self) -> None:
|
|
"""
|
|
Release NVIDIA transport resources.
|
|
|
|
Lifecycle flags are controlled by start_session(),
|
|
stop_session(), and abort_session().
|
|
"""
|
|
self._started = False
|
|
|
|
current_task = asyncio.current_task()
|
|
|
|
if self._receiver_task is not None:
|
|
if self._receiver_task is not current_task:
|
|
if not self._receiver_task.done():
|
|
self._receiver_task.cancel()
|
|
|
|
try:
|
|
await self._receiver_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
self._receiver_task = None
|
|
|
|
if self._websocket is not None:
|
|
try:
|
|
await self._websocket.close()
|
|
finally:
|
|
self._websocket = None
|
|
|
|
if self._http is not None:
|
|
await self._http.aclose()
|
|
self._http = None |