binary PCM frames (16 kHz)          JSON: Begin, SpeechStarted, Turn, Termination
 microphone ───────────────────────────▶ wss://streaming.assemblyai.com/v3/ws ───────────┐
                                                                                          │ final Turn
                                                                                          ▼
                                     POST https://llm-gateway.assemblyai.com/v1/chat/completions
                                          (messages + tools ─▶ reply, or tool_calls to run)
                                                                                          │ reply text
                                                                                          ▼
 speaker ◀─────────────────────────────── wss://tts.us.assemblyai.com/v1/ws ◀── Generate, Flush, Cancel
          JSON Audio frames (base64 PCM, 24 kHz)   plus FlushDone, Cancelled, Termination, Error

brew install portaudio

pip install -U "websockets>=14" httpx sounddevice
export ASSEMBLYAI_API_KEY=your_key_here
python agent.py

[tts] session 285ab483-997d-439d-a207-411b373a4ae1: 24000 Hz pcm_s16le
Listening. Use headphones, then press Ctrl+C to stop.
[stt] session 6b1e0f62-0c5a-4d8e-9f0e-2c4a8d1e7b13
[tts] reply 0: 3120 ms of audio
[stt] speech started
[user] What's the weather like in Lisbon today?
[tool] get_weather({'city': 'Lisbon'})
[agent] It's sunny in Lisbon today, with a high around 22 degrees.
[tts] reply 1: 3440 ms of audio
[stt] speech started
[user] And what about Madrid?
[tool] get_weather({'city': 'Madrid'})
[agent] Madrid is sunny too, with a high around 22 degrees, and it should stay dry all afternoon.

[agent] barge-in: stopping speech
[user] Actually, never mind, thanks.
[tts] cancelled after 1.84 s
[agent] No problem. Have a great day!
[tts] reply 3: 1680 ms of audio
^C[stt] terminated cleanly
[tts] terminated cleanly

# agent.py: a voice agent on raw WebSockets. No AssemblyAI SDK.
import asyncio
import base64
import json
import os
import string
import threading
from urllib.parse import urlencode

import httpx
import sounddevice as sd
import websockets

API_KEY = os.environ["ASSEMBLYAI_API_KEY"]
HEADERS = {"Authorization": API_KEY}  # the raw key, with no "Bearer" prefix

MIC_RATE = 16000
MIC_CHUNK = 800  # 50 ms of 16 kHz audio per binary frame

STT_URL = "wss://streaming.assemblyai.com/v3/ws?" + urlencode({
    "speech_model": "universal-3-6-pro",
    "sample_rate": MIC_RATE,
})
TTS_URL = "wss://tts.us.assemblyai.com/v1/ws?" + urlencode({
    "voice_id": "49c443a4-8a1d-5f7e-b70b-98bdbd9823c3",  # Audrey
    "language": "english",
    "sample_rate": 24000,
    "encoding": "pcm_s16le",
})
LLM_URL = "https://llm-gateway.assemblyai.com/v1/chat/completions"
LLM_MODEL = "claude-haiku-4-5-20251001"  # any model from the LLM Gateway model list

SYSTEM_PROMPT = (
    "You are a friendly voice assistant. Answer in one or two short spoken "
    "sentences, with no markdown or lists. Use get_weather for weather questions."
)
GREETING = "Hi there! Ask me about the weather anywhere in the world."

async def send_microphone(stt):
    """Stream the microphone to the STT socket as binary frames."""
    loop = asyncio.get_running_loop()
    frames = asyncio.Queue()

    def on_audio(indata, frame_count, time_info, status):
        loop.call_soon_threadsafe(frames.put_nowait, bytes(indata))

    with sd.RawInputStream(samplerate=MIC_RATE, channels=1, dtype="int16",
                           blocksize=MIC_CHUNK, callback=on_audio):
        while True:
            await stt.send(await frames.get())  # bytes -> binary frame

# "yes" and "no" are left out on purpose: they are real answers. Edit for your domain.
BACKCHANNELS = frozenset({
    "mhm", "mm", "mmhm", "uh", "huh", "um", "hmm", "ah", "oh",
    "yeah", "yep", "okay", "ok", "right", "gotcha",
})


def is_backchannel(text):
    words = text.lower().translate(str.maketrans("", "", string.punctuation)).split()
    return len(words) < 2 or all(w in BACKCHANNELS for w in words)


async def read_transcripts(stt, agent):
    """Handle every STT frame: Begin, SpeechStarted, Turn and Termination."""
    async for raw in stt:
        msg = json.loads(raw)
        kind = msg.get("type")

        if kind == "Begin":
            print(f"[stt] session {msg['id']}")
        elif kind == "SpeechStarted":
            print("[stt] speech started")  # barge-in waits for words, below
        elif kind == "Turn":
            text = msg.get("transcript", "").strip()
            if not text:
                continue
            if agent.speaking:
                if is_backchannel(text):
                    continue  # "mhm" while the agent talks is not an interruption
                await agent.interrupt()
            if msg.get("end_of_turn"):
                print(f"\n[user] {text}")
                agent.turns.put_nowait(text)
            else:
                print(f"[user ...] {text}", end="\r", flush=True)
        elif kind == "Termination":
            print(f"[stt] closed after {msg.get('audio_duration_seconds')} s of audio")
            return
        # Ignore frame types you do not handle; the protocol can add more.

class Speaker:
    """Plays raw 16-bit PCM as it arrives. clear() drops whatever hasn't played."""

    def __init__(self, sample_rate):
        self._buf = bytearray()
        self._lock = threading.Lock()
        self._stream = sd.RawOutputStream(samplerate=sample_rate, channels=1,
                                          dtype="int16", callback=self._fill)
        self._stream.start()

    def _fill(self, outdata, frame_count, time_info, status):
        n = len(outdata)
        with self._lock:
            chunk = bytes(self._buf[:n])
            del self._buf[:n]
        outdata[:len(chunk)] = chunk
        outdata[len(chunk):] = b"\x00" * (n - len(chunk))  # silence when idle

    def play(self, pcm):
        with self._lock:
            self._buf += pcm

    def clear(self):
        with self._lock:
            self._buf.clear()

    @property
    def busy(self):
        return len(self._buf) > 0

    def close(self):
        self._stream.stop()
        self._stream.close()

class Agent:
    """Owns the TTS socket: sends Generate/Flush/Cancel and plays the audio."""

    def __init__(self, tts, http, sample_rate):
        self.tts = tts
        self.http = http
        self.speaker = Speaker(sample_rate)
        self.turns = asyncio.Queue()
        self.messages = [{"role": "system", "content": SYSTEM_PROMPT}]
        self.in_flight = False  # a request was flushed and its FlushDone hasn't arrived
        self.dropping = 0       # Cancels sent whose Cancelled hasn't arrived yet

    @property
    def speaking(self):
        return self.in_flight or self.speaker.busy

    async def say(self, text):
        self.in_flight = True
        await self.tts.send(json.dumps({"type": "Generate", "text": text}))
        await self.tts.send(json.dumps({"type": "Flush"}))

    async def interrupt(self):
        print("\n[agent] barge-in: stopping speech")
        self.dropping += 1
        self.in_flight = False
        self.speaker.clear()
        await self.tts.send(json.dumps({"type": "Cancel"}))


async def read_speech(tts, agent):
    """Handle every TTS frame after Begin."""
    async for raw in tts:
        msg = json.loads(raw)
        kind = msg.get("type")

        if kind == "Audio":
            if agent.dropping:
                continue  # the tail of a cancelled reply: drop it
            agent.speaker.play(base64.b64decode(msg["audio"]))
        elif kind == "FlushDone":
            agent.in_flight = False
            print(f"[tts] reply {msg['flush_id']}: {msg['audio_duration_ms']} ms of audio")
        elif kind == "Cancelled":
            agent.dropping -= 1
            print(f"[tts] cancelled after {msg['audio_duration_seconds']} s")
        elif kind == "Termination":
            print(f"[tts] closed: {msg['total_input_char_length']} characters, "
                  f"{msg['audio_duration_seconds']} s of audio")
            return
        elif kind == "Error":
            raise RuntimeError(f"TTS {msg['error_code']} {msg['error_type']}: {msg['error']}")
        # WordBoundaries is on by default on the raw wire; this agent ignores it.


async def keep_tts_alive(agent):
    """The TTS session closes after two idle minutes; KeepAlive resets the timer."""
    while True:
        await asyncio.sleep(60)
        if not agent.speaking:
            await agent.tts.send(json.dumps({"type": "KeepAlive"}))

TOOLS = [{
    "type": "function",
    "function": {
        "name": "get_weather",
        "description": "Get the current weather for a city.",
        "parameters": {
            "type": "object",
            "properties": {"city": {"type": "string"}},
            "required": ["city"],
        },
    },
}]


def get_weather(city):
    return {"city": city, "conditions": "sunny", "high_c": 22}  # swap in a real API


async def think(http, messages):
    """One LLM turn over plain HTTP, running tool calls until the model answers."""
    while True:
        response = await http.post(LLM_URL, headers=HEADERS, json={
            "model": LLM_MODEL,
            "messages": messages,
            "tools": TOOLS,
            "max_tokens": 300,
        })
        response.raise_for_status()
        message = response.json()["choices"][0]["message"]
        calls = message.get("tool_calls") or []
        if not calls:
            answer = message.get("content") or ""
            messages.append({"role": "assistant", "content": answer})
            return answer

        messages.append({"role": "assistant", "content": message.get("content"),
                         "tool_calls": calls})
        for call in calls:
            args = call["function"]["arguments"]
            args = json.loads(args or "{}") if isinstance(args, str) else args
            print(f"[tool] {call['function']['name']}({args})")
            if call["function"]["name"] == "get_weather":
                result = get_weather(**args)
            else:
                result = {"error": "unknown tool"}
            messages.append({"role": "tool", "tool_call_id": call["id"],
                             "content": json.dumps(result)})


async def reply_loop(agent):
    """Each final user turn becomes one LLM call and one TTS request."""
    while True:
        text = await agent.turns.get()
        agent.messages.append({"role": "user", "content": text})
        answer = await think(agent.http, agent.messages)
        print(f"[agent] {answer}")
        await agent.say(answer)

async def wait_for_termination(ws):
    async for raw in ws:
        if json.loads(raw).get("type") == "Termination":
            return


async def main():
    async with websockets.connect(STT_URL, additional_headers=HEADERS) as stt, \
               websockets.connect(TTS_URL, additional_headers=HEADERS) as tts, \
               httpx.AsyncClient(timeout=30) as http:

        # TTS speaks first: Begin, or an Error in its place.
        begin = json.loads(await tts.recv())
        if begin["type"] == "Error":
            raise SystemExit(f"TTS {begin['error_code']}: {begin['error']}")
        config = begin["configuration"]
        print(f"[tts] session {begin['id']}: {config['sample_rate']} Hz {config['encoding']}")

        agent = Agent(tts, http, config["sample_rate"])
        await agent.say(GREETING)

        tasks = [asyncio.create_task(job) for job in (
            send_microphone(stt),
            read_transcripts(stt, agent),
            read_speech(tts, agent),
            reply_loop(agent),
            keep_tts_alive(agent),
        )]
        print("Listening. Use headphones, then press Ctrl+C to stop.")
        try:
            await asyncio.gather(*tasks)
        except asyncio.CancelledError:
            pass
        except websockets.ConnectionClosedError as e:
            print(f"socket closed: {e}")  # close code and reason
        finally:
            for task in tasks:
                task.cancel()
            agent.speaker.close()
            # Terminate both sessions and wait for their summaries.
            for name, ws in (("stt", stt), ("tts", tts)):
                try:
                    await ws.send(json.dumps({"type": "Terminate"}))
                    await asyncio.wait_for(wait_for_termination(ws), timeout=5)
                    print(f"[{name}] terminated cleanly")
                except (websockets.ConnectionClosed, asyncio.TimeoutError):
                    pass


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass