feat: establish voice assistant V1 baseline
This commit is contained in:
430
src/voice_assistant/stt/nvidia.py
Normal file
430
src/voice_assistant/stt/nvidia.py
Normal file
@@ -0,0 +1,430 @@
|
||||
"""
|
||||
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
|
||||
Reference in New Issue
Block a user