JSON quickstart

uv init --python 3.11 coda-json
cd coda-json
uv add 'websockets>=15,<16'

import asyncio
import base64
import json
import os

from websockets.asyncio.client import connect

URL = "wss://coda.api.rime.ai/ws"
CONTEXT_ID = "turn-1"
# Replace this list with your model's tokens, unchanged.
TOKENS = ["Your ", "order ", "shipped ", "on ", "Tuesday", ". ",
          "It ", "should ", "arrive ", "by ", "Friday", "."]


async def write_audio(output, chunk):
    task = asyncio.create_task(asyncio.to_thread(output.write, chunk))
    try:
        await asyncio.shield(task)
    except asyncio.CancelledError:
        # A worker thread cannot be cancelled. Finish the write before closing.
        while not task.done():
            try:
                await asyncio.shield(task)
            except asyncio.CancelledError:
                pass
        task.result()
        raise


async def receive_event(connection):
    raw = await connection.recv()
    if not isinstance(raw, str):
        raise RuntimeError("Expected a JSON text message")
    event = json.loads(raw)
    if "error" in event:
        raise RuntimeError(f"Rime error: {event['error']}")
    return event


async def send_tokens(connection):
    for token in TOKENS:
        await connection.send(json.dumps({"contextId": CONTEXT_ID, "text": token}))
    await connection.send(json.dumps({"contextId": CONTEXT_ID, "end": {}}))


async def receive_audio(connection, output):
    while True:
        event = await receive_event(connection)
        if event.get("contextId") != CONTEXT_ID:
            raise RuntimeError("Unexpected context ID")
        if "started" in event:
            print("Request ID:", event["started"].get("requestId"))
        elif "audio" in event:
            await write_audio(output, base64.b64decode(event["audio"], validate=True))
        elif "done" in event:
            return
        elif "cancelled" in event:
            raise RuntimeError("Synthesis was cancelled")


async def main():
    headers = {"Authorization": f"Bearer {os.environ['RIME_API_KEY']}"}
    # This example allows two minutes for the whole operation.
    async with asyncio.timeout(120):
        async with connect(
            URL, additional_headers=headers, subprotocols=["rime.v1.json"]
        ) as connection:
            if connection.subprotocol != "rime.v1.json":
                raise RuntimeError("JSON subprotocol was not selected")
            event = await receive_event(connection)
            if event.get("ready", {}).get("protocol") != 1:
                raise RuntimeError("Expected WebSocket protocol version 1")
            await connection.send(json.dumps({
                "contextId": CONTEXT_ID,
                "start": {
                    "speaker": "lyra",
                    "language": "en",
                    "text": "",
                    "splitStrategy": "SPLIT_STRATEGY_ACCUMULATE_SENTENCE",
                    "audioParameters": {
                        "audioFormat": "audio/pcm",
                        "samplingRate": 24000,
                    },
                },
            }))
            with open("output.pcm", "wb", buffering=0) as output:
                async with asyncio.TaskGroup() as tasks:
                    tasks.create_task(send_tokens(connection))
                    tasks.create_task(receive_audio(connection, output))
    print("Saved output.pcm")


if __name__ == "__main__":
    asyncio.run(main())

uv run stream_json.py

uv run python - <<'PY'
import wave
from pathlib import Path

pcm = Path("output.pcm").read_bytes()
with wave.open("output.wav", "wb") as out:
    out.setnchannels(1)
    out.setsampwidth(2)
    out.setframerate(24000)
    out.writeframes(pcm)
PY