feat: establish voice assistant V1 baseline

This commit is contained in:
2026-10-09 18:31:21 +00:00
commit cebcb3fd61
57 changed files with 9501 additions and 0 deletions

View File

@@ -0,0 +1,246 @@
import asyncio
import json
import wave
import websockets
GATEWAY_URL = "ws://127.0.0.1:8000/ws/audio"
AUDIO_FILE = (
"/home/ubuntu/voice-assistant/audios/chafik_normalize.wav"
)
CHUNK_FRAMES = 1600
async def main() -> None:
print("========================================")
print("STT GATEWAY SMOKE TEST")
print("========================================")
print(f"Gateway : {GATEWAY_URL}")
print(f"Audio : {AUDIO_FILE}")
print(f"Chunk : {CHUNK_FRAMES} frames")
print()
async with websockets.connect(GATEWAY_URL) as websocket:
# -------------------------------------------------------------
# 1. session.start
# -------------------------------------------------------------
await websocket.send(
json.dumps(
{
"type": "session.start",
"audio": {
"encoding": "pcm_s16le",
"sample_rate": 16000,
"channels": 1,
},
}
)
)
response = json.loads(
await websocket.recv()
)
print(
f"[RECV] {response['type']}"
)
if response["type"] != "session.ready":
raise RuntimeError(
f"Expected session.ready, got {response}"
)
session_id = response["session_id"]
# -------------------------------------------------------------
# 2. audio.start
# -------------------------------------------------------------
await websocket.send(
json.dumps(
{
"type": "audio.start",
}
)
)
response = json.loads(
await websocket.recv()
)
print(
f"[RECV] {response['type']}"
)
if response["type"] != "turn.started":
raise RuntimeError(
f"Expected turn.started, got {response}"
)
turn_id = response["turn_id"]
print(
f"Session : {session_id}"
)
print(
f"Turn : {turn_id}"
)
print()
# -------------------------------------------------------------
# 3. Start a receiver task.
#
# This is necessary because partial/final events can arrive
# while we are still sending audio.
# -------------------------------------------------------------
received_events: list[dict] = []
async def receive_events() -> None:
while True:
message = await websocket.recv()
event = json.loads(message)
received_events.append(event)
event_type = event.get("type")
if event_type == "stt.partial":
print(
f"[PARTIAL] {event.get('text', '')}"
)
elif event_type == "stt.final":
print(
f"[FINAL ] {event.get('text', '')}"
)
elif event_type == "turn.ended":
print(
"[RECV] turn.ended"
)
break
elif event_type == "error":
print(
"[ERROR]",
event,
)
break
else:
print(
f"[RECV] {event_type}"
)
receiver_task = asyncio.create_task(
receive_events()
)
# -------------------------------------------------------------
# 4. Stream WAV as:
#
# audio.chunk JSON
# binary PCM
#
# The sequence number belongs to the logical chunk.
# -------------------------------------------------------------
with wave.open(AUDIO_FILE, "rb") as wav:
sequence_number = 0
while True:
audio = wav.readframes(CHUNK_FRAMES)
if not audio:
break
await websocket.send(
json.dumps(
{
"type": "audio.chunk",
"turn_id": turn_id,
"sequence_number": sequence_number,
}
)
)
await websocket.send(audio)
sequence_number += 1
print()
print(
f"Audio sent: {sequence_number} chunks"
)
# -------------------------------------------------------------
# 5. audio.stop
# -------------------------------------------------------------
await websocket.send(
json.dumps(
{
"type": "audio.stop",
"turn_id": turn_id,
}
)
)
# -------------------------------------------------------------
# 6. Wait for:
#
# stt.final
# turn.ended
#
# The receiver task handles both.
# -------------------------------------------------------------
await receiver_task
print()
print("========================================")
print("RESULT")
print("========================================")
partials = [
event
for event in received_events
if event.get("type") == "stt.partial"
]
finals = [
event
for event in received_events
if event.get("type") == "stt.final"
]
turn_ended = [
event
for event in received_events
if event.get("type") == "turn.ended"
]
print(
f"Partial events : {len(partials)}"
)
print(
f"Final events : {len(finals)}"
)
print(
f"Turn ended : {len(turn_ended)}"
)
if finals:
print()
print("Final transcript:")
print(
finals[-1].get("text", "")
)
print()
print("STT Gateway smoke test finished.")
if __name__ == "__main__":
asyncio.run(main())

114
tools/stt_smoke_test.py Normal file
View File

@@ -0,0 +1,114 @@
import argparse
import asyncio
from pathlib import Path
from riva.client.asr import AudioChunkFileIterator
from riva.client.realtime import RealtimeClientASR
NIM_SERVER = "127.0.0.1:9000"
INPUT_FILE = "/home/ubuntu/voice-assistant/audios/chafik_normalize.wav"
CHUNK_FRAMES = 1600
def build_args() -> argparse.Namespace:
return argparse.Namespace(
input_file=INPUT_FILE,
mic=False,
list_devices=False,
duration=None,
input_device=None,
sample_rate_hz=16000,
num_channels=1,
file_streaming_chunk=CHUNK_FRAMES,
output_text=None,
prompt="",
server=NIM_SERVER,
ssl_root_cert=None,
ssl_client_cert=None,
ssl_client_key=None,
use_ssl=False,
metadata=None,
options=None,
max_message_length=67108864,
word_time_offsets=False,
max_alternatives=1,
profanity_filter=False,
automatic_punctuation=True,
no_verbatim_transcripts=False,
language_code="pt-BR",
model_name="",
boosted_lm_words=None,
boosted_lm_score=4.0,
speaker_diarization=False,
diarization_max_speakers=3,
start_history=-1,
start_threshold=-1.0,
stop_history=-1,
stop_threshold=-1.0,
stop_history_eou=-1,
stop_threshold_eou=-1.0,
custom_configuration="",
endpoint="/v1/realtime",
query_params="intent=transcription",
mint_api_key=None,
)
async def main() -> None:
input_path = Path(INPUT_FILE)
if not input_path.exists():
raise FileNotFoundError(
f"Audio file not found: {input_path}"
)
args = build_args()
print(f"NIM server : {args.server}")
print(f"Audio : {input_path}")
print(f"Language : {args.language_code}")
print(f"Chunk : {CHUNK_FRAMES} frames")
client = RealtimeClientASR(args=args)
audio_iterator = AudioChunkFileIterator(
input_file=str(input_path),
chunk_n_frames=CHUNK_FRAMES,
delay_callback=None,
)
try:
print("Connecting to NIM...")
await client.connect()
print("Connected.")
send_task = asyncio.create_task(
client.send_audio_chunks(audio_iterator)
)
receive_task = asyncio.create_task(
client.receive_responses()
)
await asyncio.gather(
send_task,
receive_task,
)
finally:
await client.disconnect()
print("Disconnected.")
if __name__ == "__main__":
asyncio.run(main())

View File

View File

@@ -0,0 +1,327 @@
import asyncio
import json
import wave
from pathlib import Path
import websockets
GATEWAY_URL = "ws://127.0.0.1:8001/ws/audio"
AUDIO_FILE = Path(
"/home/ubuntu/voice-assistant/audios/chafik_normalize.wav"
)
OUTPUT_FILE = Path(
"/tmp/voice_pipeline_response.pcm"
)
CHUNK_FRAMES = 1600
async def receive_messages(
websocket,
state: dict,
) -> None:
while True:
message = await websocket.recv()
if isinstance(message, bytes):
state["tts_chunks"] += 1
state["tts_bytes"] += len(message)
state["tts_audio"].extend(message)
print(
f"[TTS AUDIO] "
f"chunk={state['tts_chunks']} "
f"bytes={len(message)}"
)
continue
event = json.loads(message)
event_type = event.get("type")
if event_type == "stt.partial":
state["partial_count"] += 1
print(f"[PARTIAL] {event.get('text', '')}")
elif event_type == "stt.final":
state["final_text"] = event.get("text")
print(f"[FINAL ] {state['final_text']}")
elif event_type == "agent.response":
state["agent_response"] = event.get("text")
print(
f"[AGENT ] "
f"{state['agent_response']}"
)
elif event_type == "turn.ended":
state["turn_ended"] = True
print("[RECV] turn.ended")
return
elif event_type == "error":
raise RuntimeError(
f"Gateway error: {event}"
)
else:
print(f"[RECV] {event_type}")
async def send_audio(
websocket,
turn_id: int,
) -> int:
sent_chunks = 0
with wave.open(str(AUDIO_FILE), "rb") as audio:
sample_rate = audio.getframerate()
channels = audio.getnchannels()
sample_width = audio.getsampwidth()
if sample_rate != 16000:
raise RuntimeError(
f"Expected 16 kHz audio, got {sample_rate} Hz."
)
if channels != 1:
raise RuntimeError(
f"Expected mono audio, got {channels} channels."
)
if sample_width != 2:
raise RuntimeError(
f"Expected 16-bit PCM audio, got "
f"{sample_width * 8}-bit."
)
while True:
pcm = audio.readframes(CHUNK_FRAMES)
if not pcm:
break
await websocket.send(
json.dumps(
{
"type": "audio.chunk",
"turn_id": turn_id,
"sequence_number": sent_chunks,
}
)
)
await websocket.send(pcm)
sent_chunks += 1
return sent_chunks
async def main() -> None:
print("=" * 40)
print("VOICE PIPELINE SMOKE TEST")
print("=" * 40)
print(f"Gateway : {GATEWAY_URL}")
print(f"Audio : {AUDIO_FILE}")
print(f"Chunk : {CHUNK_FRAMES} frames")
print()
state = {
"partial_count": 0,
"final_text": None,
"agent_response": None,
"tts_chunks": 0,
"tts_bytes": 0,
"tts_audio": bytearray(),
"turn_ended": False,
}
async with websockets.connect(GATEWAY_URL) as websocket:
# ---------------------------------------------------------
# session.start
# ---------------------------------------------------------
await websocket.send(
json.dumps(
{
"type": "session.start",
"audio": {
"encoding": "pcm_s16le",
"sample_rate": 16000,
"channels": 1,
},
}
)
)
message = await websocket.recv()
if isinstance(message, bytes):
raise RuntimeError(
"Expected JSON session.ready."
)
session_ready = json.loads(message)
print(
f"[RECV] "
f"{session_ready.get('type')}"
)
if session_ready.get("type") != "session.ready":
raise RuntimeError(
f"Expected session.ready, got "
f"{session_ready}"
)
session_id = session_ready["session_id"]
# ---------------------------------------------------------
# audio.start
# ---------------------------------------------------------
await websocket.send(
json.dumps(
{
"type": "audio.start",
}
)
)
message = await websocket.recv()
if isinstance(message, bytes):
raise RuntimeError(
"Expected JSON turn.started."
)
turn_started = json.loads(message)
print(
f"[RECV] "
f"{turn_started.get('type')}"
)
if turn_started.get("type") != "turn.started":
raise RuntimeError(
f"Expected turn.started, got "
f"{turn_started}"
)
turn_id = turn_started["turn_id"]
print()
print(f"Session : {session_id}")
print(f"Turn : {turn_id}")
print()
# ---------------------------------------------------------
# Run sender and receiver concurrently
# ---------------------------------------------------------
receiver_task = asyncio.create_task(
receive_messages(
websocket,
state,
)
)
try:
sent_chunks = await send_audio(
websocket,
turn_id,
)
except Exception:
receiver_task.cancel()
raise
print()
print(
f"Audio sent: {sent_chunks} chunks"
)
# ---------------------------------------------------------
# Tell Gateway that user audio is complete.
# ---------------------------------------------------------
await websocket.send(
json.dumps(
{
"type": "audio.stop",
"turn_id": turn_id,
}
)
)
# ---------------------------------------------------------
# Wait for the entire pipeline to finish.
# ---------------------------------------------------------
await asyncio.wait_for(
receiver_task,
timeout=120.0,
)
# ---------------------------------------------------------
# Validate the complete STT -> Agent -> TTS pipeline.
# ---------------------------------------------------------
if not state["turn_ended"]:
raise RuntimeError("Turn did not complete.")
if not state["final_text"]:
raise RuntimeError("STT final transcript was not received.")
if not state["agent_response"]:
raise RuntimeError("Agent response was not received.")
if state["tts_chunks"] == 0:
raise RuntimeError("No TTS audio was received.")
if state["tts_bytes"] == 0:
raise RuntimeError("TTS output is empty.")
# -------------------------------------------------------------
# Save TTS audio
# -------------------------------------------------------------
OUTPUT_FILE.write_bytes(
bytes(state["tts_audio"])
)
print()
print("=" * 40)
print("RESULT")
print("=" * 40)
print(f"Session : {session_id}")
print(f"Turn : {turn_id}")
print(
f"Audio sent : {sent_chunks}"
)
print(
f"Partial events : {state['partial_count']}"
)
print(
f"Final transcript: {state['final_text']}"
)
print(
f"Agent response : {state['agent_response']}"
)
print(
f"TTS chunks : {state['tts_chunks']}"
)
print(
f"TTS bytes : {state['tts_bytes']}"
)
print(
f"TTS output : {OUTPUT_FILE}"
)
print(
f"Turn ended : {state['turn_ended']}"
)
print()
print(
"Voice pipeline smoke test finished."
)
if __name__ == "__main__":
asyncio.run(main())