Speak in realtime

import asyncio
import base64
import json
import os
import time
from pathlib import Path

import websockets

MODEL = "voxtral-mini-tts-3-0"
URL = f"wss://api.mistral.ai/v1/audio/speech/realtime?model={MODEL}"
HEADERS = {"Authorization": f"Bearer {os.environ['MISTRAL_API_KEY']}"}


async def receive_audio(ws, label, start):
    """Receive audio for one generation. Returns the raw bytes."""
    chunks = []
    first_audio = None
    while True:
        msg = json.loads(await ws.recv())
        if msg["type"] == "speech.audio.delta":
            if first_audio is None:
                first_audio = time.perf_counter() - start
                print(f"  [{label}] First audio after {first_audio:.2f}s")
            chunks.append(base64.b64decode(msg["audio_data"]))
        elif msg["type"] == "speech.audio.done":
            print(f"  [{label}] Done ({sum(len(c) for c in chunks)} bytes)")
            return b"".join(chunks)
        elif msg["type"] == "error":
            raise RuntimeError(msg["error"]["message"])


async def main():
    start = time.perf_counter()
    async with websockets.connect(URL, additional_headers=HEADERS) as ws:
        # Wait for session.created
        msg = json.loads(await ws.recv())
        assert msg["type"] == "session.created"
        print(f"Connected ({time.perf_counter() - start:.2f}s)")

        # Configure the session
        await ws.send(json.dumps({
            "type": "session.update",
            "session": {"voice": "voice_447", "response_format": "wav"},
        }))
        msg = json.loads(await ws.recv())
        assert msg["type"] == "session.updated"

        # Generation 1: send text in pieces, then flush
        await ws.send(json.dumps({"type": "input_text.append", "text": "Hello "}))
        await ws.send(json.dumps({"type": "input_text.append", "text": "world. "}))
        await ws.send(json.dumps({
            "type": "input_text.append",
            "text": "This is a test of the realtime TTS API.",
        }))
        await ws.send(json.dumps({"type": "input_text.flush"}))
        audio1 = await receive_audio(ws, "gen1", start)
        Path("realtime_gen1.wav").write_bytes(audio1)

        # Generation 2: send text, then end (closes the connection)
        await ws.send(json.dumps({"type": "input_text.append", "text": "Goodbye."}))
        await ws.send(json.dumps({"type": "input_text.end"}))
        audio2 = await receive_audio(ws, "gen2", start)
        Path("realtime_gen2.wav").write_bytes(audio2)

    print(f"Total: {time.perf_counter() - start:.2f}s")
    print("Saved realtime_gen1.wav and realtime_gen2.wav")


asyncio.run(main())