6krrt/poller.py

#!/usr/bin/env python3
"""
Pricing/catalog poller for the local LLM router.

Fetches the NeuralWatt model catalog (unauthenticated public endpoint),
normalizes it into the `models` table in router.db, and flags rows that have
gone stale.

Run manually:
    python poller.py

Run on a schedule (cron example, every 2 hours):
    0 */2 * * * /usr/bin/python3 /path/to/poller.py >> /var/log/router-poller.log 2>&1

Does NOT touch energy data — energy is only available per-completion, not from
a models list, so the dispatcher writes `energy_observations` instead.

NeuralWatt is the only provider. The `provider` column and the (model_id,
provider) primary key are kept so a second provider can be added without a
migration.
"""

from __future__ import annotations

import sqlite3
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Optional

import requests

from config import RouterConfig, load_config

NEURALWATT_MODELS_URL = "https://api.neuralwatt.com/v1/models"

REQUEST_TIMEOUT = 20  # seconds

# Serving-class suffixes. NeuralWatt ships one base model as several catalog
# rows that differ only by these tokens, and they combine freely — hence ids
# like 'glm-5.2-short-fast-flex'. They are stripped from the end of the id one
# segment at a time so a base name that merely *contains* a lookalike token is
# never misread (e.g. 'deepseek-v4-flash' is not a '-fast' row).
SUFFIX_FLEX = "flex"
SUFFIX_FAST = "fast"
SUFFIX_SHORT = "short"
SERVING_SUFFIXES = frozenset({SUFFIX_FLEX, SUFFIX_FAST, SUFFIX_SHORT})


def parse_serving_class(model_id: str) -> tuple[str, str, str]:
    """Derive (latency_class, reasoning_mode, context_variant) from a model id.

    Returns the schema defaults ('standard', 'default', 'full') for a base
    model. Suffixes are matched as whole '-'-delimited segments only.
    """
    segments = model_id.lower().split("-")
    found = set()
    while len(segments) > 1 and segments[-1] in SERVING_SUFFIXES:
        found.add(segments.pop())

    return (
        "flex" if SUFFIX_FLEX in found else "standard",
        "reduced" if SUFFIX_FAST in found else "default",
        "short" if SUFFIX_SHORT in found else "full",
    )


def parse_base_model_id(model_id: str) -> str:
    """Reduce a catalog id to the model family underneath it.

    ``glm-5.2-short-fast-flex`` and ``glm-5.2`` are the same weights served
    differently, and ``deepseek-ai/DeepSeek-V4-Flash`` is the HF-style
    duplicate of ``deepseek-v4-flash``. Proficiency is a property of the
    weights, not of the queue they sit in, so scores are keyed on this and
    every serving variant inherits from its family. Leaderboard priors work
    the same way — no benchmark rates a ``-flex`` row separately.

    Note this deliberately collapses ``-fast`` too, even though reasoning
    being off does change answer quality. The eval runner scores ``-fast``
    variants separately and overrides the inherited value; the family is the
    fallback, not the final word.
    """
    namespace_stripped = model_id.rsplit("/", 1)[-1]
    segments = namespace_stripped.lower().split("-")
    while len(segments) > 1 and segments[-1] in SERVING_SUFFIXES:
        segments.pop()
    return "-".join(segments)


def parse_access_level(display_name: Optional[str], description: Optional[str]) -> str:
    """Derive an access level from the catalog's prose.

    NeuralWatt exposes no structured gating field — restricted models are only
    marked in free text ("Private preview (grant-gated)", "(Canary)"). Routing
    to one earns a 403 at dispatch, so this is parsed defensively: anything
    that looks gated is treated as gated.
    """
    blob = f"{display_name or ''} {description or ''}".lower()
    if "grant-gated" in blob or "private preview" in blob:
        return "preview"
    if "canary" in blob:
        return "canary"
    return "public"


@dataclass
class ModelRow:
    model_id: str
    provider: str
    base_model_id: str
    display_name: Optional[str]
    cost_per_1m_prompt: Optional[float]
    cost_per_1m_completion: Optional[float]
    cost_per_1m_prompt_cached: Optional[float]
    context_window: Optional[int]
    max_output_tokens: Optional[int]
    supports_tools: bool
    supports_json_mode: bool
    supports_vision: bool
    supports_reasoning: bool
    reasoning_default_enabled: bool
    latency_class: str
    reasoning_mode: str
    context_variant: str
    access_level: str
    pricing_tbd: bool
    deprecated: bool

    def effective_context_window(self, cfg: RouterConfig) -> Optional[int]:
        """Usable context, after the safety factor and an output reserve.

        ``context.per_model_overrides`` wins where it is set, which is what it
        is for: the global factor is a guess that has to hold for the whole
        catalog, while a row someone has actually measured deserves its own
        number.

        Compared with ``is not None`` rather than ``or``, so an override of 0
        reserve tokens means zero rather than silently falling through to the
        default. Present-but-falsy is a trap this project's own eval set tests
        models on; the router should not walk into it.
        """
        if not self.context_window:
            return None
        override = cfg.context.per_model_overrides.get(self.model_id)

        factor = cfg.context.safety_factor
        if override is not None and override.safety_factor is not None:
            factor = override.safety_factor

        # 11 of 19 catalog rows report no max_output_tokens, so the configured
        # reserve carries most of the catalog.
        if override is not None and override.output_reserve_tokens is not None:
            reserve = override.output_reserve_tokens
        else:
            reserve = self.max_output_tokens or cfg.context.default_output_reserve_tokens

        usable = int(self.context_window * factor) - reserve
        return max(usable, 0)


def fetch_neuralwatt() -> list[ModelRow]:
    resp = requests.get(NEURALWATT_MODELS_URL, timeout=REQUEST_TIMEOUT)
    resp.raise_for_status()
    payload = resp.json()

    rows = []
    for m in payload.get("data", []):
        model_id = m.get("id")
        meta = m.get("metadata", {}) or {}
        pricing = meta.get("pricing", {}) or {}
        caps = meta.get("capabilities", {}) or {}
        limits = meta.get("limits", {}) or {}
        reasoning = meta.get("reasoning") or {}

        supports_reasoning = bool(caps.get("reasoning"))
        # capabilities.reasoning only means "the API accepts a reasoning
        # param" and is true for nearly the whole catalog. default_enabled is
        # the discriminating signal; a few models (the kimi-k2.7-code family)
        # expose no reasoning block at all, so fall back to the capability.
        default_enabled = reasoning.get("default_enabled")
        if default_enabled is None:
            default_enabled = supports_reasoning

        latency_class, reasoning_mode, context_variant = parse_serving_class(model_id)

        rows.append(
            ModelRow(
                model_id=model_id,
                provider="neuralwatt",
                base_model_id=parse_base_model_id(model_id),
                display_name=meta.get("display_name"),
                cost_per_1m_prompt=pricing.get("input_per_million"),
                cost_per_1m_completion=pricing.get("output_per_million"),
                cost_per_1m_prompt_cached=pricing.get("cached_input_per_million"),
                context_window=limits.get("max_context_length")
                or m.get("max_model_len"),
                max_output_tokens=limits.get("max_output_tokens"),
                supports_tools=bool(caps.get("tools")),
                supports_json_mode=bool(caps.get("json_mode")),
                supports_vision=bool(caps.get("vision")),
                supports_reasoning=supports_reasoning,
                reasoning_default_enabled=bool(default_enabled),
                latency_class=latency_class,
                reasoning_mode=reasoning_mode,
                context_variant=context_variant,
                access_level=parse_access_level(
                    meta.get("display_name"), meta.get("description")
                ),
                pricing_tbd=bool(pricing.get("pricing_tbd")),
                deprecated=bool(meta.get("deprecated")),
            )
        )
    return rows


def upsert(conn: sqlite3.Connection, rows: list[ModelRow], cfg: RouterConfig) -> None:
    now = datetime.now(timezone.utc).isoformat()
    for r in rows:
        availability = "deprecated" if r.deprecated else "active"
        conn.execute(
            """
            INSERT INTO models (
                model_id, provider, base_model_id, display_name,
                cost_per_1m_prompt, cost_per_1m_completion, cost_per_1m_prompt_cached,
                context_window, effective_context_window, max_output_tokens,
                supports_tools, supports_json_mode, supports_vision, supports_reasoning,
                reasoning_default_enabled, latency_class, reasoning_mode,
                context_variant, access_level,
                pricing_tbd, deprecated, availability, last_updated
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ON CONFLICT(model_id, provider) DO UPDATE SET
                base_model_id = excluded.base_model_id,
                display_name = excluded.display_name,
                cost_per_1m_prompt = excluded.cost_per_1m_prompt,
                cost_per_1m_completion = excluded.cost_per_1m_completion,
                cost_per_1m_prompt_cached = excluded.cost_per_1m_prompt_cached,
                context_window = excluded.context_window,
                effective_context_window = excluded.effective_context_window,
                max_output_tokens = excluded.max_output_tokens,
                supports_tools = excluded.supports_tools,
                supports_json_mode = excluded.supports_json_mode,
                supports_vision = excluded.supports_vision,
                supports_reasoning = excluded.supports_reasoning,
                reasoning_default_enabled = excluded.reasoning_default_enabled,
                latency_class = excluded.latency_class,
                reasoning_mode = excluded.reasoning_mode,
                context_variant = excluded.context_variant,
                access_level = excluded.access_level,
                pricing_tbd = excluded.pricing_tbd,
                deprecated = excluded.deprecated,
                availability = excluded.availability,
                last_updated = excluded.last_updated
            """,
            (
                r.model_id,
                r.provider,
                r.base_model_id,
                r.display_name,
                r.cost_per_1m_prompt,
                r.cost_per_1m_completion,
                r.cost_per_1m_prompt_cached,
                r.context_window,
                r.effective_context_window(cfg),
                r.max_output_tokens,
                int(r.supports_tools),
                int(r.supports_json_mode),
                int(r.supports_vision),
                int(r.supports_reasoning),
                int(r.reasoning_default_enabled),
                r.latency_class,
                r.reasoning_mode,
                r.context_variant,
                r.access_level,
                int(r.pricing_tbd),
                int(r.deprecated),
                availability,
                now,
            ),
        )
    conn.commit()


def mark_stale(conn: sqlite3.Connection, cfg: RouterConfig) -> None:
    """Flag rows that weren't touched by this poll run as stale, rather than
    silently leaving old data looking current."""
    conn.execute(
        """
        UPDATE models
        SET availability = 'stale'
        WHERE availability = 'active'
          AND julianday('now') - julianday(last_updated) > ?
        """,
        (cfg.freshness.stale_after_days,),
    )
    conn.commit()


def main() -> int:
    cfg = load_config("config.yaml")
    conn = sqlite3.connect(cfg.database.path)
    conn.execute("PRAGMA foreign_keys = ON")

    try:
        rows = fetch_neuralwatt()
    except requests.RequestException as e:
        print(f"[neuralwatt] FAILED: {e}", file=sys.stderr)
        conn.close()
        return 1

    upsert(conn, rows, cfg)
    print(f"[neuralwatt] upserted {len(rows)} models")

    mark_stale(conn, cfg)
    conn.close()
    print(f"done, {len(rows)} rows upserted total")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())