Python example

"""A minimal Talking Avatar Realistic Avatars client: one LiveAvatar per conversation.

pip install websockets
"""
import asyncio
import itertools
import json
import struct

import websockets

API_URL = "wss://api.talking-avatar.dev"
HEADER = struct.Struct("<dI")  # play time (ms since the epoch; 0 = queue it), sample rate
PIECE = 256 * 1024  # messages can be up to 1 MB


class AvatarError(Exception):
    def __init__(self, code, message):
        super().__init__(f"{message} ({code})")
        self.code = code


class LiveAvatar:
    @classmethod
    async def start(cls, api_key, avatar, sample_rate=24000, livekit=None, url=API_URL):
        """Opens a session; returns once the avatar is live. livekit: {"url", "token"} to join your room instead."""
        ws = await websockets.connect(url, max_size=2**20)
        start = {"type": "start", "apiKey": api_key, "avatar": avatar, "sampleRate": sample_rate}
        if livekit:
            start.update(livekitUrl=livekit["url"], token=livekit["token"], audio={"publish": True})
        else:
            start["output"] = "direct"
        await ws.send(json.dumps(start))
        error = None
        try:
            async for message in ws:
                msg = json.loads(message)
                if msg["type"] == "ready":
                    return cls(ws, sample_rate, msg)
                if msg["type"] == "error":
                    error = msg["message"]
        except websockets.ConnectionClosed:
            pass
        raise AvatarError(ws.close_code, error or ws.close_reason or "connection closed")

    def __init__(self, ws, sample_rate, ready):
        self.ws = ws
        self.sample_rate = sample_rate
        self.width, self.height = ready["width"], ready["height"]
        self.on_event = lambda msg: None  # speech_started, speech_ended, republished, error
        self.closed = asyncio.get_running_loop().create_future()  # resolves with (code, reason)
        self._pending = {}  # request id -> future
        self._clearing = []  # futures waiting for "cleared"
        self._ids = itertools.count(1)
        self._carry = b""  # a byte split from its sample at the end of the last chunk
        self._reader = asyncio.create_task(self._read())

    async def push_audio(self, pcm: bytes):
        """Your agent's voice as it streams: PCM16 mono, little-endian, at sample_rate. Chunks of any size."""
        pcm = self._carry + pcm
        whole = len(pcm) - len(pcm) % 2
        pcm, self._carry = pcm[:whole], pcm[whole:]
        for at in range(0, len(pcm), PIECE):
            await self.ws.send(HEADER.pack(0, self.sample_rate) + pcm[at:at + PIECE])

    async def clear(self):
        """Barge-in: drops the audio not yet played. Returns how many ms of the interrupted reply were heard."""
        self._carry = b""
        cleared = asyncio.get_running_loop().create_future()
        self._clearing.append(cleared)
        await self.ws.send(json.dumps({"type": "clear"}))
        return await cleared

    async def handle_view(self, body):
        """Direct output: your page's route. Pass the request body, return the result as JSON."""
        if body.get("viewId") and body.get("sdp"):
            await self._request({"type": "view_answer", "viewId": body["viewId"], "sdp": body["sdp"]})
            return {"viewId": body["viewId"]}
        offer = await self._request({"type": "view"})
        return {"viewId": offer["viewId"], "sdp": offer["sdp"]}

    async def close(self):
        """Ends the session: billing stops."""
        await self.ws.close(1000)
        await self.closed

    async def _request(self, msg):
        msg["id"] = f"r{next(self._ids)}"
        reply = self._pending[msg["id"]] = asyncio.get_running_loop().create_future()
        await self.ws.send(json.dumps(msg))
        return await asyncio.wait_for(reply, 15)

    async def _read(self):
        try:
            async for message in self.ws:
                msg = json.loads(message)
                reply = self._pending.pop(msg.get("id"), None)
                if reply:
                    if msg["type"] == "error":
                        reply.set_exception(AvatarError(0, msg["message"]))
                    else:
                        reply.set_result(msg)
                elif msg["type"] == "cleared" and self._clearing:
                    self._clearing.pop(0).set_result(msg["playedMs"])
                else:
                    self.on_event(msg)
        except websockets.ConnectionClosed:
            pass
        finally:
            ended = AvatarError(self.ws.close_code, self.ws.close_reason or "session ended")
            for waiting in [*self._pending.values(), *self._clearing]:
                if not waiting.done():
                    waiting.set_exception(ended)
            self.closed.set_result((self.ws.close_code, self.ws.close_reason))

import os
from contextlib import asynccontextmanager

from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, JSONResponse

from talking_avatar import LiveAvatar


@asynccontextmanager
async def lifespan(app: FastAPI):
    # One avatar per conversation: here, one for the whole demo
    app.state.avatar = await LiveAvatar.start(os.environ["TALKING_AVATAR_API_KEY"], "ava4", sample_rate=24000)
    yield
    await app.state.avatar.close()  # billing stops


app = FastAPI(lifespan=lifespan)


@app.post("/avatar")
async def avatar_view(request: Request):
    try:
        return await request.app.state.avatar.handle_view(await request.json())
    except Exception as error:
        return JSONResponse({"error": str(error)}, status_code=503)


@app.get("/")
async def page():
    return HTMLResponse("""
<div id="avatar" style="width: 640px; aspect-ratio: 16 / 9"></div>
<button id="start">Start</button>
<script type="module">
  import { LiveAvatarView } from "https://cdn.jsdelivr.net/npm/@interviewflowai/talking-avatar@0.2.0/dist/live.js";
  const view = new LiveAvatarView({ container: document.getElementById("avatar"), endpoint: "/avatar" });
  document.getElementById("start").onclick = () => view.connect();
</script>
""")