327 lines
8.3 KiB
Python
327 lines
8.3 KiB
Python
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()) |