engine.py

"""Every model the Den uses, called through Hugging Face Inference Providers.

    review(pitch, panel)   the panel's written review         GLM-5.3-Flash, with DeepSeek-V4.1-Flash as backup
    transcribe(samples)    16 kHz mono float32 in, text out   Whisper large-v3-turbo, with Qwen3-ASR-1.7B as backup
    speak(turns)           (judge, line) pairs in, audio out  Kokoro-82M, with when each line starts

All of it is billed to the Space owner's credits through the HF_TOKEN secret, so visitors need no
account, and a spending guard stops the Den before those credits run out.
"""
from __future__ import annotations

import io
import json
import threading
import time
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable, Sequence, TypeVar

import numpy as np
import soundfile as sf
from huggingface_hub import InferenceClient

import artifacts
import config
import contract
from contract import Review
from judges import JUDGES, system_prompt

T = TypeVar("T")


def log(stage: str, started: float, **extra) -> None:
    detail = " ".join(f"{k}={v}" for k, v in extra.items())
    print(f"[den] {stage} {time.perf_counter() - started:.2f}s {detail}".rstrip(), flush=True)


def _client(provider: str, timeout: float = config.TIMEOUT) -> InferenceClient:
    return InferenceClient(provider=provider, api_key=config.HF_TOKEN, timeout=timeout)


def _first(attempts: Sequence[Callable[[], T]], head_start: float) -> T:
    """Run attempts in order of preference and return the first to succeed.

    Each attempt gets `head_start` seconds to itself before the next one starts alongside it, and a
    failure starts the next one at once. Providers stall now and then; this way one slow provider costs a
    visitor a few seconds rather than a timeout, and the preferred one still wins whenever it is healthy.
    """
    pool = ThreadPoolExecutor(max_workers=len(attempts))
    waiting, running, error = list(attempts), set(), None
    try:
        while waiting or running:
            while waiting and (not running or head_start <= 0):   # No head start: all start together.
                running.add(pool.submit(waiting.pop(0)))
            done, running = wait(running, timeout=head_start if waiting else None, return_when=FIRST_COMPLETED)
            for future in done:
                if future.exception() is None:
                    return future.result()
                error = future.exception()
            if waiting and (not done or error):
                running.add(pool.submit(waiting.pop(0)))
        raise error
    finally:
        pool.shutdown(wait=False, cancel_futures=True)   # A slower attempt finishes in the background.


# ---------------------------------------------------------------- spending


OUT_OF_BUDGET = ValueError("The panel has used up its budget, so it's taking a break. Sorry!")


class Budget:
    """Dollars spent this calendar month, kept on disk so a restart doesn't forget them."""

    def __init__(self, path: Path, limit: float):
        self._path, self._limit, self._lock = path, limit, threading.Lock()

    def _read(self) -> float:
        try:
            data = json.loads(self._path.read_text())
            return float(data["usd"]) if data.get("month") == _month() else 0.0
        except (OSError, ValueError, KeyError, TypeError):
            return 0.0

    def check(self) -> None:
        """Before any call: refuse if the Space has no token, or this month's budget is spent."""
        if not config.HF_TOKEN:
            raise ValueError("The panel isn't connected yet. Your pitch is safe - please try again shortly.")
        if self._read() >= self._limit:
            raise OUT_OF_BUDGET

    def add(self, usd: float) -> float:
        with self._lock:
            total = self._read() + usd
            self._path.parent.mkdir(parents=True, exist_ok=True)
            self._path.write_text(json.dumps({"month": _month(), "usd": round(total, 6)}))
            return total


def _month() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m")


BUDGET = Budget(Path(config.SPEND_FILE), config.BUDGET_USD)


def _status(error: Exception) -> int | None:
    """The HTTP status, if any. Hugging Face refuses every provider with 402 once the account's credits are spent."""
    return getattr(getattr(error, "response", None), "status_code", None)


def _give_up(stage: str, error: Exception, message: str) -> ValueError:
    """Every provider failed. Logs only the error's type and status, never its text, which can echo a pitch."""
    print(f"[den] {stage} failed everywhere: {type(error).__name__} {_status(error)}", flush=True)
    return OUT_OF_BUDGET if _status(error) == 402 else ValueError(message)


# ---------------------------------------------------------------- the review


# What Hugging Face actually bills, read from this account's usage: audio models by the second of request
# time (a stalled request costs as much as it waits, so timeouts are short), Together a flat cent per request,
# the other chat providers a fraction of that. Dollars.
PER_SECOND = {"deepinfra": 0.0013, "fal-ai": 0.0013, "hf-inference": 0.00025}


def _billed(provider: str, call: Callable[[], T]) -> T:
    """Run one audio request and count what it costs, whether or not it succeeds."""
    started = time.perf_counter()
    try:
        return call()
    finally:
        BUDGET.add((time.perf_counter() - started) * PER_SECOND[provider])


@dataclass(frozen=True)
class Route:
    """One chat model at one provider, and roughly what one review costs there."""
    provider: str
    model: str
    cost: float
    options: dict = field(default_factory=dict)
    schema_in_prompt: bool = False   # For a provider with JSON mode but no JSON-schema mode.


# GLM-5.3-Flash writes the liveliest judges. It always thinks unless told otherwise, and each provider has its
# own switch: Together and Z.ai honour a low reasoning effort (Together answered 20 of 20 reviews in 3-5s with
# ten visitors at once; Z.ai 18 of 20 in 6-9s). DeepSeek on Baseten is the last resort, in case GLM is down.
GLM_QUICK = {"reasoning_effort": "low", "response_format": {"type": "json_object"}}
ROUTES = (
    Route("together", "zai-org/GLM-5.3-Flash", 0.01, GLM_QUICK, schema_in_prompt=True),
    Route("zai-org", "zai-org/GLM-5.3-Flash", 0.001, GLM_QUICK, schema_in_prompt=True),
    Route("baseten", "deepseek-ai/DeepSeek-V4.1-Flash", 0.002,
          {"chat_template_kwargs": {"enable_thinking": False},
           "response_format": {"type": "json_schema",
                               "json_schema": {"name": "review", "strict": True, "schema": contract.review_schema()}}}),
)


def _ask(route: Route, pitch: str, panel: Sequence[str]) -> Review:
    started = time.perf_counter()
    system = system_prompt(panel)
    if route.schema_in_prompt:
        system += "\n\nReply with only a JSON object matching this schema:\n" + json.dumps(contract.review_schema())
    user = json.dumps({"pitch": pitch, "excerpts": dict(enumerate(contract.excerpts(pitch)))}, ensure_ascii=False)
    reply = _client(route.provider).chat_completion(
        [{"role": "system", "content": system}, {"role": "user", "content": user}], model=route.model,
        max_tokens=config.MAX_TOKENS, temperature=config.TEMPERATURE, top_p=config.TOP_P,
        presence_penalty=config.PRESENCE_PENALTY, extra_body=route.options)
    BUDGET.add(route.cost)   # A reply arrived, so it was billed, even if it turns out to be unreadable.
    review = contract.read(reply.choices[0].message.content or "", panel, pitch)
    log("review", started, route=f"{route.provider}/{route.model}")
    return review


def review(pitch: str, panel: Sequence[str]) -> Review:
    BUDGET.check()
    try:
        return _first([lambda route=route: _ask(route, pitch, panel) for route in ROUTES], config.REVIEW_HEAD_START)
    except Exception as error:   # Busy, slow, refused or unreadable, everywhere.
        raise _give_up("review", error, "The panel is swamped right now. Your pitch is safe - give it a moment and "
                                        "try again.") from None


# ---------------------------------------------------------------- speech


# Hugging Face's own Whisper server is steady at about 3s per minute of audio and cheap. DeepInfra's Qwen3-ASR
# is usually quicker but sometimes stalls, so it only joins in when Whisper is slow or fails.
ASR = (("hf-inference", "openai/whisper-large-v3-turbo"), ("deepinfra", "Qwen/Qwen3-ASR-1.7B"))
# DeepInfra has Kokoro's British voices. fal, with American voices only, is used only if DeepInfra fails: never
# alongside it, because running both bills both.
TTS = (("deepinfra", "voice"), ("fal-ai", "backup_voice"))
TTS_MODEL = "hexgrad/Kokoro-82M"


def transcribe(samples: np.ndarray) -> str:
    BUDGET.check()
    started = time.perf_counter()
    path = artifacts.save_flac(samples, config.ASR_SAMPLE_RATE)   # hf-inference needs a file to know the format

    def attempt(provider: str, model: str) -> str:
        return _billed(provider, lambda: _client(provider, config.ASR_TIMEOUT).automatic_speech_recognition(
            path, model=model).text)

    minutes = len(samples) / config.ASR_SAMPLE_RATE / 60
    try:   # Whisper's head start grows with the recording, so a long pitch isn't paid for twice.
        text = _first([lambda p=p, m=m: attempt(p, m) for p, m in ASR],
                      config.ASR_HEAD_START + config.ASR_SECONDS_PER_MINUTE * minutes)
    except Exception as error:
        raise _give_up("transcribe", error, "We couldn't transcribe that just now. Try again, or type your pitch.") from None
    log("transcribe", started, audio_seconds=round(minutes * 60, 1))
    return text.strip()


def _voice_with(provider: str, voice_field: str, turns: Sequence[tuple[str, str]]) -> tuple[int, np.ndarray, list[float]]:
    """One whole clip from one provider, every line at once, so accents never switch mid-argument."""
    started = time.perf_counter()

    def line(turn: tuple[str, str]) -> tuple[np.ndarray, int]:
        judge, text = turn
        audio = _billed(provider, lambda: _client(provider, config.LINE_TIMEOUT).text_to_speech(
            text, model=TTS_MODEL, extra_body={"voice": getattr(JUDGES[judge], voice_field)}))
        return sf.read(io.BytesIO(audio), dtype="float32")

    with ThreadPoolExecutor(max_workers=len(turns)) as pool:
        lines = list(pool.map(line, turns))
    rate = lines[0][1]
    pause = np.zeros(int(config.GAP_SECONDS * rate), dtype="float32")
    audio = np.concatenate([part for samples, _ in lines for part in (samples, pause)])
    starts = list(np.cumsum([0] + [len(samples) + len(pause) for samples, _ in lines[:-1]]) / rate)
    log("speak", started, provider=provider, speech_seconds=round(len(audio) / rate, 1))
    return rate, audio, [round(float(s), 2) for s in starts]


def speak(turns: Sequence[tuple[str, str]]) -> tuple[int, np.ndarray, list[float]]:
    """The panel's lines as one clip. Returns (sample rate, samples, the second each line starts)."""
    BUDGET.check()
    for provider, voice_field in TTS:
        started = time.perf_counter()
        try:
            return _voice_with(provider, voice_field, turns)
        except Exception as error:
            log("speak", started, provider=provider, failed=type(error).__name__)
            if _status(error) == 402:
                raise OUT_OF_BUDGET from None
    raise ValueError("The panel lost their voices just now. Your review is below.")