diff --git a/README.md b/README.md index 2bb5e43..5fcb281 100644 --- a/README.md +++ b/README.md @@ -76,6 +76,7 @@ latency budget (best first). Lower is better on every metric: | Model | False cutoffs @ 300 ms | False cutoffs @ 600 ms | Latency @ 5% cutoff | Latency @ 10% cutoff | | --- | ---: | ---: | ---: | ---: | | **LiveKit Turn Detector v1** | **9.9%** | **4.5%** | **543 ms** | **295 ms** | +| Baton | 12.3% | 4.8% | 577 ms | 350 ms | | Deepgram Flux | 12.9% | 9.9% | 1151 ms | 548 ms | | ultraVAD | 27.7% | 11.9% | 899 ms | 663 ms | | Gradium | 55.6% | 12.6% | 913 ms | 656 ms | diff --git a/eot_harness/.env.example b/eot_harness/.env.example index 21fb866..94e3126 100644 --- a/eot_harness/.env.example +++ b/eot_harness/.env.example @@ -10,3 +10,6 @@ XAI_API_KEY= SPEECHMATICS_API_KEY= OPENAI_API_KEY= GRADIUM_API_KEY= +BATON_API_KEY= +# Optional: override the Baton endpoint (e.g. a local server). +# BATON_BASE_URL=https://baton.joinin.ai diff --git a/eot_harness/baton_adapter.py b/eot_harness/baton_adapter.py new file mode 100644 index 0000000..7b80a14 --- /dev/null +++ b/eot_harness/baton_adapter.py @@ -0,0 +1,224 @@ +"""Streaming adapter for Baton, JoinIn AI's end-of-turn model. + +Baton is a hosted service; this adapter is a thin client for it. + +It scores each turn with a single stateless ``POST /v1/turn`` -- the whole turn in, +the full p_eot grid out. That is one request per turn rather than ~11k (one per grid +point), and because nothing is held open between calls the harness's ``--concurrency`` +maps straight onto horizontally scaled instances. + +The adapter passes ``row["messages"]`` as prior conversational context and sends only +the audio. It never reads ``row["words"]``, so the model is given nothing about the +turn it is being asked to judge beyond the audio up to that moment -- the same +constraint it operates under in production. There is a unit test asserting this. + +Requires BATON_API_KEY. Request a beta key at hello@joinin.ai. +""" + +from __future__ import annotations + +import asyncio +import base64 +import os +from typing import Any + +import numpy as np + +from .languages import supports_any_benchmark_language +from .io import decode_audio +from .streaming_stt import ( + build_event_prediction_rows, + resample_audio, + resolve_api_key, +) + +DEFAULT_BASE_URL = "https://baton.joinin.ai" +TURN_PATH = "/v1/turn" +DEFAULT_INFERENCE_INTERVAL = 0.1 +SAMPLE_RATE = 16000 + + +class BatonAdapter: + """Scores a turn with one stateless POST to the Baton hosted API.""" + + display_name = "Baton" + score_point = 0.2 + + def __init__( + self, + *, + api_key: str | None = None, + base_url: str | None = None, + model: str = "baton-v1", + # A cold instance can take minutes to become ready, so a request arriving + # during a scale-up may wait. Retrying past that is the difference between a + # warm-up blip and a silently dropped turn -- the harness is often invoked + # with --skip-errors, which would quietly shrink the eval set rather than + # fail loudly. + max_retries: int = 4, + retry_backoff: float = 5.0, + # A 429 is not a failure, it is "come back later" -- capacity, not breakage. + # It WILL succeed given time, so it gets its own far more patient budget: + # a dropped turn silently shrinks the eval set, which is worse than a slow run. + capacity_retries: int = 20, + capacity_backoff_cap: float = 30.0, + timeout: float = 300.0, + ) -> None: + self._api_key = api_key + self._base_url = (base_url or os.environ.get("BATON_BASE_URL") or DEFAULT_BASE_URL).rstrip("/") + self.model = model + self.max_retries = int(max_retries) + self.retry_backoff = float(retry_backoff) + self.capacity_retries = int(capacity_retries) + self.capacity_backoff_cap = float(capacity_backoff_cap) + self.timeout = float(timeout) + + @property + def adapter_id(self) -> str: + return f"joinin/{self.model}" + + def supports_language(self, lang_code: str) -> bool: + # English is materially stronger than the rest; it still scores every + # benchmark language. + return supports_any_benchmark_language(lang_code) + + async def predict_turn( + self, + row: dict[str, Any], + *, + inference_interval: float = DEFAULT_INFERENCE_INTERVAL, + ) -> dict[str, Any]: + api_key = self._api_key or resolve_api_key("BATON_API_KEY") + pcm_bytes, audio_sec = _prepare_pcm16_audio(row, sample_rate=SAMPLE_RATE) + + events: list[dict[str, Any]] = [] + attempt = capacity_attempt = 0 + while True: + try: + events = await self._score_events( + api_key=api_key, + pcm_bytes=pcm_bytes, + # Prior turns only. Sending anything about the turn under + # judgement would leak the future and flatter every number. + # Normalised here as well as server-side: the benchmark's + # non-English splits carry content: None, which a strict schema + # rejects with 422. Doing it client-side means the run works + # against a server that has not been updated yet. + messages=_clean_messages(row.get("messages")), + inference_interval=inference_interval, + ) + break + except Exception as exc: + # Capacity gets its own budget. Every turn must be scored: a turn + # dropped here vanishes from the metrics with nothing in the output + # saying so, and would read as the model failing rather than the + # server being busy. + if _is_capacity_error(exc) and capacity_attempt < self.capacity_retries: + delay = min(self.capacity_backoff_cap, + self.retry_backoff * (2 ** capacity_attempt)) + capacity_attempt += 1 + await asyncio.sleep(delay) + continue + if _is_transient_error(exc) and attempt < self.max_retries: + attempt += 1 + await asyncio.sleep(self.retry_backoff * attempt) + continue + raise + + return { + "id": row["id"], + "audio_sec": audio_sec, + "events": events, + "prediction_rows": build_event_prediction_rows( + row, + events, + inference_interval=inference_interval, + ), + } + + async def _score_events( + self, + *, + api_key: str, + pcm_bytes: bytes, + messages: list[dict[str, str]], + inference_interval: float, + ) -> list[dict[str, Any]]: + aiohttp = _import_aiohttp() + url = f"{self._base_url}{TURN_PATH}" + payload = { + "audio_pcm16_b64": base64.b64encode(pcm_bytes).decode("ascii"), + "messages": messages, + "inference_interval": float(inference_interval), + "sample_rate": SAMPLE_RATE, + } + timeout = aiohttp.ClientTimeout(total=self.timeout) + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.post(url, json=payload, headers={"X-API-Key": api_key}) as resp: + if resp.status != 200: + detail = await resp.text() + raise RuntimeError(f"Baton returned HTTP {resp.status}: {detail[:400]}") + body = await resp.json() + return [ + {"timestamp": float(e["timestamp"]), "p_eot": float(e["p_eot"])} + for e in body.get("events", []) + if e.get("timestamp") is not None and e.get("p_eot") is not None + ] + + +def _import_aiohttp(): + try: + import aiohttp + except ImportError as exc: # pragma: no cover + raise RuntimeError("BatonAdapter requires aiohttp (uv run --with aiohttp ...)") from exc + return aiohttp + + +def _prepare_pcm16_audio(row: dict[str, Any], *, sample_rate: int = SAMPLE_RATE) -> tuple[bytes, float]: + """PCM16-encode a row's audio using Baton's scale convention. + + Deliberately not ``streaming_stt.prepare_pcm16_audio``: that helper scales + by 32767, which shifts every sample by one LSB. That is tolerable for some + consumers, but it is a measurable quality cost for anything sensitive to exact + sample values. Encoding as round(x * 32768) clipped into int16 range is what + Baton's API contract specifies. + """ + audio = row.get("audio") or {} + if "array" in audio: + # already-decoded shape (the batch path, and hand-built test rows) + array = np.asarray(audio["array"], dtype=np.float32) + orig_sr = int(audio["sampling_rate"]) + else: + # A STREAMING adapter is handed the raw row, and the harness loads the + # dataset with Audio(decode=False) -- so `audio` is {bytes, path} and must + # be decoded here. Assuming the decoded shape passes every unit test built + # on synthetic rows and then KeyErrors on the first real turn. + array, orig_sr = decode_audio(audio) + array = np.asarray(array, dtype=np.float32) + orig_sr = int(orig_sr) + if orig_sr != sample_rate: + array = resample_audio(array, orig_sr, sample_rate) + scaled = np.rint(np.clip(array, -1.0, 1.0) * 32768.0) + pcm = np.clip(scaled, -32768.0, 32767.0).astype(np.int16) + return pcm.tobytes(), float(len(array) / sample_rate) + + +def _clean_messages(messages) -> list[dict[str, str]]: + """Coerce the dataset's message list into {role, content} strings.""" + out: list[dict[str, str]] = [] + for m in messages or []: + out.append({"role": str(m.get("role") or "user"), "content": str(m.get("content") or "")}) + return out + + +def _is_capacity_error(exc: BaseException) -> bool: + """429 / 503: the server is busy or scaling, not broken. Always worth waiting for.""" + m = str(exc) + return "429" in m or "503" in m + + +def _is_transient_error(exc: BaseException) -> bool: + if isinstance(exc, (asyncio.TimeoutError, ConnectionError, OSError)): + return True + message = str(exc).lower() + return any(token in message for token in ("timeout", "temporarily", "502", "503", "504", "429", "reset")) diff --git a/tests/test_baton_adapter.py b/tests/test_baton_adapter.py new file mode 100644 index 0000000..5ee860c --- /dev/null +++ b/tests/test_baton_adapter.py @@ -0,0 +1,246 @@ +from __future__ import annotations + +import asyncio +import base64 +import json + +import numpy as np +import pytest + +from eot_harness import baton_adapter as mod +from eot_harness.baton_adapter import BatonAdapter, _is_transient_error, _prepare_pcm16_audio + + +class FakeHTTP: + """Minimal aiohttp double: records the request, replays a scripted response.""" + + def __init__(self, body=None, *, status=200, text=""): + self._body = body if body is not None else {"events": []} + self._status = status + self._text = text + self.url = None + self.payload = None + self.headers = None + + # -- aiohttp module surface -- + def ClientTimeout(self, **kw): + return kw + + def ClientSession(self, **kw): + return self + + # -- session -- + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + def post(self, url, *, json=None, headers=None): + self.url = url + self.payload = json + self.headers = headers + return self + + # -- response -- + @property + def status(self): + return self._status + + async def json(self): + return self._body + + async def text(self): + return self._text + + +def _install(monkeypatch, http): + monkeypatch.setattr(mod, "_import_aiohttp", lambda: http) + return http + + +def _row(**over): + row = { + "id": "turn-1", + "audio": {"array": np.zeros(16000, dtype=np.float32), "sampling_rate": 16000}, + "silence_spans": [{"start": 0.5, "end": 0.9}], + "messages": [{"role": "assistant", "content": "how can I help?"}], + "words": [{"word": "leaked", "end": 0.4}], + } + row.update(over) + return row + + +def _run(adapter, row, **kw): + return asyncio.run(adapter.predict_turn(row, **kw)) + + +# ---- identity ------------------------------------------------------------- + +def test_adapter_id_is_namespaced(): + assert BatonAdapter().adapter_id == "joinin/baton-v1" + + +def test_adapter_id_tracks_model_override(): + assert BatonAdapter(model="baton-v2").adapter_id == "joinin/baton-v2" + + +def test_score_point_matches_published_basis(): + assert BatonAdapter.score_point == 0.2 + assert BatonAdapter.display_name == "Baton" + + +def test_supports_every_benchmark_language(): + adapter = BatonAdapter() + assert adapter.supports_language("en") + assert adapter.supports_language("de") + + +# ---- PCM16 encoding ------------------------------------------------------- + +def test_pcm16_uses_32768_scale_not_32767(): + row = _row(audio={"array": np.array([0.5], dtype=np.float32), "sampling_rate": 16000}) + pcm, _ = _prepare_pcm16_audio(row) + assert np.frombuffer(pcm, dtype=np.int16)[0] == 16384 # 0.5 * 32768 + + +def test_pcm16_clips_to_int16_range(): + row = _row(audio={"array": np.array([1.0, -1.0, 2.0], dtype=np.float32), "sampling_rate": 16000}) + pcm, _ = _prepare_pcm16_audio(row) + assert np.frombuffer(pcm, dtype=np.int16).tolist() == [32767, -32768, 32767] + + +def test_pcm16_reports_duration_in_seconds(): + row = _row(audio={"array": np.zeros(8000, dtype=np.float32), "sampling_rate": 16000}) + _, audio_sec = _prepare_pcm16_audio(row) + assert audio_sec == pytest.approx(0.5) + + +# ---- event mapping -------------------------------------------------------- + +def test_events_are_mapped_from_response(monkeypatch): + http = _install(monkeypatch, FakeHTTP({"events": [ + {"timestamp": 0.1, "p_eot": 0.01}, + {"timestamp": 0.2, "p_eot": 0.87}, + ]})) + out = _run(BatonAdapter(api_key="k"), _row()) + assert out["events"] == [ + {"timestamp": 0.1, "p_eot": 0.01}, + {"timestamp": 0.2, "p_eot": 0.87}, + ] + assert out["id"] == "turn-1" + assert http.url.endswith("/v1/turn") + assert http.headers == {"X-API-Key": "k"} + + +def test_incomplete_events_are_skipped(monkeypatch): + _install(monkeypatch, FakeHTTP({"events": [ + {"words": ["partial"]}, + {"timestamp": 0.3}, + {"p_eot": 0.5}, + {"timestamp": 0.4, "p_eot": 0.9}, + ]})) + out = _run(BatonAdapter(api_key="k"), _row()) + assert out["events"] == [{"timestamp": 0.4, "p_eot": 0.9}] + + +def test_non_200_raises_with_detail(monkeypatch): + _install(monkeypatch, FakeHTTP(status=401, text="invalid or missing X-API-Key")) + with pytest.raises(RuntimeError, match="401"): + _run(BatonAdapter(api_key="k", max_retries=0), _row()) + + +def test_prediction_rows_are_built(monkeypatch): + _install(monkeypatch, FakeHTTP({"events": [{"timestamp": 0.6, "p_eot": 0.4}]})) + out = _run(BatonAdapter(api_key="k"), _row()) + assert isinstance(out["prediction_rows"], list) + + +# ---- the contract that matters ------------------------------------------- + +def test_only_prior_messages_are_sent_never_the_in_progress_words(monkeypatch): + """Baton must not see the turn it is being asked to judge.""" + http = _install(monkeypatch, FakeHTTP({"events": [{"timestamp": 0.1, "p_eot": 0.1}]})) + _run(BatonAdapter(api_key="k"), _row()) + assert http.payload["messages"] == [{"role": "assistant", "content": "how can I help?"}] + assert "leaked" not in json.dumps(http.payload["messages"]) + + +def test_audio_is_sent_base64_in_one_request(monkeypatch): + http = _install(monkeypatch, FakeHTTP({"events": [{"timestamp": 0.1, "p_eot": 0.1}]})) + _run(BatonAdapter(api_key="k"), _row()) + # 1 s of 16 kHz mono PCM16 = 32000 bytes, in a single call. + assert len(base64.b64decode(http.payload["audio_pcm16_b64"])) == 32000 + assert http.payload["sample_rate"] == 16000 + + +def test_inference_interval_is_forwarded(monkeypatch): + http = _install(monkeypatch, FakeHTTP({"events": [{"timestamp": 0.1, "p_eot": 0.1}]})) + _run(BatonAdapter(api_key="k"), _row(), inference_interval=0.25) + assert http.payload["inference_interval"] == 0.25 + + +# ---- auth and retries ----------------------------------------------------- + +def test_missing_api_key_raises(monkeypatch): + monkeypatch.delenv("BATON_API_KEY", raising=False) + _install(monkeypatch, FakeHTTP()) + with pytest.raises(Exception): + _run(BatonAdapter(), _row()) + + +def test_transient_errors_are_classified(): + assert _is_transient_error(asyncio.TimeoutError()) + assert _is_transient_error(ConnectionError("reset by peer")) + assert _is_transient_error(RuntimeError("503 unavailable")) + assert not _is_transient_error(ValueError("bad audio")) + + +def test_capacity_errors_are_retried_until_they_succeed(monkeypatch): + """A 429 must never cost a turn: it is the server being busy, not broken. + + A dropped turn vanishes from the metrics with nothing in the output saying so, + and would read as the model failing rather than the service being loaded. + """ + calls = {"n": 0} + + class Flaky(FakeHTTP): + def post(self, url, *, json=None, headers=None): + # count per REQUEST, not per status read -- the error message reads + # .status too, which would double-count + calls["n"] += 1 + self._status = 429 if calls["n"] <= 3 else 200 + return super().post(url, json=json, headers=headers) + + http = _install(monkeypatch, Flaky({"events": [{"timestamp": 0.1, "p_eot": 0.7}]})) + real_sleep = asyncio.sleep # capture before patching, or it recurses + monkeypatch.setattr(mod.asyncio, "sleep", lambda *_a, **_k: real_sleep(0)) + out = _run(BatonAdapter(api_key="k", retry_backoff=0.0), _row()) + assert out["events"] == [{"timestamp": 0.1, "p_eot": 0.7}] + assert calls["n"] == 4, f"expected 3 rejections then success, got {calls['n']} requests" + + +def test_capacity_retries_are_separate_from_transient_retries(): + a = BatonAdapter() + assert a.capacity_retries > a.max_retries, "429 needs a more patient budget" + assert mod._is_capacity_error(RuntimeError("Baton returned HTTP 429: busy")) + assert mod._is_capacity_error(RuntimeError("HTTP 503")) + assert not mod._is_capacity_error(RuntimeError("HTTP 401 invalid key")) + + +@pytest.mark.skipif(not hasattr(mod, "decode_audio"), reason="harness io not importable") +def test_accepts_the_raw_undecoded_audio_shape(monkeypatch): + """Streaming adapters get Audio(decode=False): {bytes, path}, not {array, ...}. + + Assuming the decoded shape passes every synthetic-row test and then KeyErrors on + the first real dataset turn. + """ + import io as _io, wave as _wave + buf = _io.BytesIO() + with _wave.open(buf, "wb") as w: + w.setnchannels(1); w.setsampwidth(2); w.setframerate(16000) + w.writeframes(np.zeros(1600, dtype="