tei_annotator/providers/base.py

"""Shared utilities and abstract base class for LLM connectors."""

from __future__ import annotations

import json
import os
import socket
import sys
import threading
import time
import urllib.error
import urllib.request
from abc import ABC, abstractmethod
from typing import Callable

_DEBUG_HTTP = os.environ.get("TEI_DEBUG_HTTP", "") == "1"


def _dbg(msg: str) -> None:
    if _DEBUG_HTTP:
        print(f"[DBG {time.strftime('%H:%M:%S')}] {msg}", file=sys.stderr, flush=True)


class ProviderTimeoutError(RuntimeError):
    """Raised when a provider's HTTP call exceeds its timeout."""

    def __init__(self, url: str, timeout: float, elapsed: float) -> None:
        super().__init__(
            f"Provider did not respond within {timeout:.0f}s "
            f"(waited {elapsed:.0f}s) — {url}"
        )
        self.url = url
        self.timeout = timeout
        self.elapsed = elapsed


class _RateLimiter:
    """Enforce a minimum interval between calls (thread-safe)."""

    def __init__(self, rate_per_minute: int) -> None:
        self._interval = 60.0 / rate_per_minute
        self._lock = threading.Lock()
        self._last: float = 0.0

    def acquire(self) -> None:
        with self._lock:
            now = time.monotonic()
            wait = self._interval - (now - self._last)
            if wait > 0:
                time.sleep(wait)
            self._last = time.monotonic()


def _post_json(url: str, payload: dict, headers: dict, timeout: int = 120) -> dict:
    """POST JSON with a hard total-elapsed timeout.

    urllib's ``timeout=`` is a per-read inactivity timeout — a server that
    trickles back bytes never trips it.  We enforce a total wall-clock cap
    via a daemon-thread watchdog and raise :class:`ProviderTimeoutError`
    when it fires.
    """
    body = json.dumps(payload).encode()
    req = urllib.request.Request(url, data=body, headers=headers, method="POST")
    _dbg(f"_post_json POST {url} (timeout={timeout}s, body={len(body)} bytes)")
    t0 = time.monotonic()

    result: dict = {}

    def _worker() -> None:
        try:
            # urllib's inactivity timeout slightly above the watchdog so
            # the watchdog wins the race and produces a clean message.
            with urllib.request.urlopen(req, timeout=timeout + 30) as resp:
                result["status"] = resp.status
                result["raw"] = resp.read()
        except BaseException as exc:  # noqa: BLE001 — pass any failure up
            result["exc"] = exc

    worker = threading.Thread(target=_worker, daemon=True)
    worker.start()
    worker.join(timeout=timeout)

    if worker.is_alive():
        elapsed = time.monotonic() - t0
        _dbg(f"_post_json WATCHDOG TIMEOUT after {elapsed:.1f}s (limit={timeout}s)")
        # The daemon thread keeps running; it will either complete on its
        # own or die with the process. urllib will time out at timeout+30s.
        raise ProviderTimeoutError(url, timeout, elapsed)

    exc = result.get("exc")
    if exc is not None:
        elapsed = time.monotonic() - t0
        if isinstance(exc, urllib.error.HTTPError):
            detail = exc.read().decode(errors="replace")
            _dbg(f"_post_json HTTPError {exc.code} after {elapsed:.1f}s: {detail[:200]}")
            raise RuntimeError(f"HTTP {exc.code}: {detail}") from exc
        if isinstance(exc, (socket.timeout, TimeoutError)):
            _dbg(f"_post_json TIMEOUT after {elapsed:.1f}s (limit={timeout}s)")
            raise ProviderTimeoutError(url, timeout, elapsed) from exc
        if isinstance(exc, urllib.error.URLError) and isinstance(
            exc.reason, (socket.timeout, TimeoutError)
        ):
            _dbg(f"_post_json TIMEOUT (URLError) after {elapsed:.1f}s")
            raise ProviderTimeoutError(url, timeout, elapsed) from exc
        _dbg(f"_post_json EXC after {elapsed:.1f}s: {type(exc).__name__}: {exc}")
        raise exc

    status, raw = result["status"], result["raw"]
    _dbg(f"_post_json POST {url} -> {status} in {time.monotonic()-t0:.1f}s "
         f"({len(raw)} bytes)")
    return json.loads(raw)


class Connector(ABC):
    """Base class for all LLM provider connectors."""

    @property
    @abstractmethod
    def id(self) -> str:
        """Short machine-readable identifier, e.g. 'hf', 'gemini', 'kisski'."""

    @property
    @abstractmethod
    def name(self) -> str:
        """Human-readable provider name shown in the UI dropdown group label."""

    @property
    @abstractmethod
    def description(self) -> str:
        """One-sentence description of the provider."""

    # Subclasses declare their model list here.
    # Prefix a model name with '*' to mark it as premium-only.
    # The '*' is stripped before the ID is passed to any API call.
    _MODELS: list[str] = []

    @abstractmethod
    def is_available(self) -> bool:
        """Return True iff the required credentials are present in the environment."""

    def models(self) -> list[str]:
        """Return all model IDs (premium and standard), with '*' stripped."""
        return [m.lstrip("*") for m in self._MODELS]

    def standard_models(self) -> list[str]:
        """Return only non-premium model IDs (those not prefixed with '*')."""
        return [m for m in self._MODELS if not m.startswith("*")]

    @property
    def default_model(self) -> str:
        """The model pre-selected in the UI (override to customise)."""
        return self.standard_models()[0] if self.standard_models() else self.models()[0]

    def get_model_statuses(self) -> list[dict]:
        """
        Return [{"model": str, "demand": int, "status": str}, ...].
        demand: 0 = free, 1-2 = moderate, 3+ = busy.
        Returns [] if this provider does not expose status information.
        """
        return []

    @abstractmethod
    def make_call_fn(self, model_id: str, timeout: int = 300) -> Callable[[str], str]:
        """Return a call_fn(prompt: str) -> str for the given model."""