feat: establish voice assistant V1 baseline
This commit is contained in:
327
tools/voice_pipeline_smoke_test.py
Normal file
327
tools/voice_pipeline_smoke_test.py
Normal 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())
|
||||
Reference in New Issue
Block a user