Add a provider

"""Acme real-time speech-to-text.

The first message is a JSON config; audio follows as binary PCM frames. The
server sends {"type": "partial" | "final", "text", "start", "end"}.
`finalize()` sends {"type": "flush"}; the server answers with a final that
has "flushed": true, which marks the end of the turn.

`language` becomes `lang`; `vocab` becomes `boost`.

Docs: https://docs.acme.example/stt/streaming
"""

import json
from collections.abc import AsyncIterator
from typing import Any

from websockets.asyncio.client import ClientConnection, connect

from cuelo.provider import ProviderError, primary_language, ws_url
from cuelo.runtime.audio import SAMPLE_WIDTH, Audio, PCMBuffer
from cuelo.stt.interface import STT, STTSession, Transcript

CHUNK_SECONDS = 0.1


class AcmeSTT(STT):
    default_base_url = "https://stream.acme.example"
    default_api_key_env = "ACME_API_KEY"
    supports = frozenset({"language", "vocab"})

    def __init__(self, *, model: str = "acme-rt-2", **kwargs: Any) -> None:
        super().__init__(**kwargs)   # language, vocab, sample_rate, base_url, api_key, api_key_env
        self.model = model

    def config(self) -> dict[str, Any]:
        config: dict[str, Any] = {
            "model": self.model,
            "encoding": "pcm_s16le",
            "sample_rate": self.sample_rate,
        }
        if self.language:
            config["lang"] = primary_language(self.language)   # "en-US" -> "en"
        if self.vocab:
            config["boost"] = self.vocab
        return config

    async def connect(self) -> STTSession:
        ws = await connect(
            f"{ws_url(self.base_url or '')}/v2/listen",
            additional_headers={"Authorization": f"Bearer {self.api_key}"},
        )
        await ws.send(json.dumps(self.config()))
        return _Session(ws, self.sample_rate)


class _Session(STTSession):
    def __init__(self, ws: ClientConnection, sample_rate: int) -> None:
        self.ws = ws
        self.sample_rate = sample_rate
        # Batch the harness's small frames into ~100 ms chunks.
        self.pcm = PCMBuffer(min_bytes=int(sample_rate * CHUNK_SECONDS) * SAMPLE_WIDTH)

    async def send(self, audio: Audio) -> None:
        if audio.sample_rate != self.sample_rate:
            raise ValueError(f"expected {self.sample_rate} Hz audio, got {audio.sample_rate}")
        if data := self.pcm.push(audio.pcm):
            await self.ws.send(data)

    async def finalize(self) -> None:
        if data := self.pcm.flush():
            await self.ws.send(data)
        await self.ws.send(json.dumps({"type": "flush"}))

    async def transcripts(self) -> AsyncIterator[Transcript]:
        async for raw in self.ws:
            msg = json.loads(raw)
            if msg.get("type") == "error":
                raise ProviderError(f"Acme STT: {msg.get('message')}")
            if msg.get("type") in ("partial", "final") and msg.get("text"):
                yield Transcript(
                    msg["text"],
                    is_final=msg["type"] == "final",
                    finalized=bool(msg.get("flushed")),
                    start=msg.get("start"),
                    end=msg.get("end"),
                )

    async def close(self) -> None:
        await self.ws.close()

[project.optional-dependencies]
acme = [
    "websockets>=17.4",
]

"stt": {
  "properties": {
    "provider": {
      "enum": ["smallest", "elevenlabs", "deepgram", "cartesia", "inworld", "assemblyai", "soniox", "acme"]
    }
  }
}

"x-cuelo-providers": {
  "stt": {
    "acme": {
      "class": "cuelo.stt.acme.AcmeSTT",
      "extra": "acme",
      "env": "ACME_API_KEY"
    }
  }
}

make schema

func orderedNames(stage string) []string {
	switch stage {
	case "stt":
		return []string{"smallest", "elevenlabs", "deepgram", "cartesia", "inworld", "assemblyai", "soniox", "acme"}

import json
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager

import pytest
from websockets.asyncio.server import ServerConnection, serve

from cuelo.provider import ProviderError
from cuelo.runtime.audio import Audio
from cuelo.stt.acme import AcmeSTT

PCM = b"\x01\x00" * 1600  # 100 ms at 16 kHz


@asynccontextmanager
async def ws_server(handler) -> AsyncIterator[str]:
    async with serve(handler, "127.0.0.1", 0) as server:
        port = next(iter(server.sockets)).getsockname()[1]
        yield f"http://127.0.0.1:{port}"


async def test_acme_stt() -> None:
    seen: dict = {}

    async def handler(ws: ServerConnection) -> None:
        seen["auth"] = ws.request.headers["Authorization"]
        seen["config"] = json.loads(await ws.recv())
        audio = await ws.recv()
        seen["audio"] = len(audio)
        await ws.send(json.dumps({"type": "partial", "text": "hel"}))
        assert json.loads(await ws.recv()) == {"type": "flush"}
        await ws.send(json.dumps({"type": "final", "text": "hello", "flushed": True}))
        await ws.close()

    async with ws_server(handler) as url:
        stt = AcmeSTT(base_url=url, api_key="k", language="en-US", vocab=["Cuelo"])
        session = await stt.connect()
        await session.send(Audio(PCM, 16_000))
        await session.finalize()
        out = [t async for t in session.transcripts()]

    assert seen["auth"] == "Bearer k"
    assert seen["config"]["lang"] == "en" and seen["config"]["boost"] == ["Cuelo"]
    assert seen["audio"] == len(PCM)
    assert [(t.text, t.is_final, t.finalized) for t in out] == [
        ("hel", False, False),
        ("hello", True, True),
    ]


async def test_acme_stt_error() -> None:
    async def handler(ws: ServerConnection) -> None:
        await ws.recv()
        await ws.send(json.dumps({"type": "error", "message": "invalid model"}))

    async with ws_server(handler) as url:
        session = await AcmeSTT(base_url=url, api_key="k").connect()
        with pytest.raises(ProviderError, match="invalid model"):
            async for _ in session.transcripts():
                pass

CUELO_STT=acme cuelo start