Multi-context WebSocket

import asyncio
import base64
import json

import websockets

URL = (
    "wss://api.chariot.in/v1/tts/ws-multi"
    "?voice_id=VOICE_ID&model_type=chariot-multilingual-v1"
)
HEADERS = {"chariotai-api-key": "YOUR_API_KEY"}


async def main():
    async with websockets.connect(URL, additional_headers=HEADERS) as ws:
        cancelled = set()

        async def send(msg):
            await ws.send(json.dumps(msg))

        async def receive():
            async for raw in ws:
                event = json.loads(raw)
                kind, ctx = event["type"], event.get("context_id")
                if kind == "audio.chunk":
                    if ctx not in cancelled:          # drop audio still in flight after a cancel
                        play_pcm(ctx, base64.b64decode(event["audio"]))  # 44.1 kHz, mono, s16le
                elif kind == "context.done":
                    print(ctx, "ended:", event["reason"], "credits:", event["credits_utilized"])
                elif kind == "error":
                    print("error:", event["code"], event["message"])
                elif kind == "session.done":
                    return

        receiver = asyncio.create_task(receive())

        # Turn 1: stream the reply, flushing at the end of each sentence.
        await send({"type": "input.text", "context_id": "turn-1", "text": "Your order shipped on Monday. "})
        await send({"type": "input.flush", "context_id": "turn-1"})
        await send({"type": "input.text", "context_id": "turn-1", "text": "It should arrive by Thursday evening. "})
        await send({"type": "input.flush", "context_id": "turn-1"})

        await asyncio.sleep(1.0)  # the user starts speaking

        # Barge-in: stop turn 1 (and its playback), answer in turn 2 on the same socket.
        cancelled.add("turn-1")
        stop_playback("turn-1")
        await send({"type": "context.cancel", "context_id": "turn-1"})
        await send({"type": "input.text", "context_id": "turn-2", "text": "Sure, which account? "})
        await send({"type": "context.close", "context_id": "turn-2"})

        # Finish once every context is done.
        await send({"type": "session.close"})
        await receiver


asyncio.run(main())