6krrt/tier.py

#!/usr/bin/env python3
"""DB tiering pass for the local LLM router.

Reads every row from the ``models`` table, resolves each row's tier via
``tiering.resolve_tier`` (heuristic + config override map), and writes the
resolved tier back to ``models.tier``. This pass is the source of truth for
``tier`` and OVERWRITES existing/non-NULL tiers — overrides flow through
config, not DB edits. Idempotent: running twice yields identical tiers.

Run manually:
    python tier.py

The pure resolver lives in ``tiering.py``; this module owns the DB I/O and
the thin CLI so ``tiering.py`` stays free of I/O (see its docstring).
"""

from __future__ import annotations

import sqlite3
import sys

from config import RouterConfig, load_config
from tiering import resolve_tier

WARNING_UNIFORMLY_FALSE = (
    "WARNING: reasoning_default_enabled is uniformly False — verify poller field mapping"
)


def apply_tiering(conn: sqlite3.Connection, config: RouterConfig) -> None:
    """Resolve and write ``models.tier`` for every row in the table.

    Reads all rows, resolves each via ``resolve_tier``, and issues one
    ``UPDATE models SET tier = ? WHERE model_id = ? AND provider = ?`` per
    row, then commits. Only the ``tier`` column is touched.

    Sanity guard: if ``reasoning_default_enabled`` is True for ZERO rows
    across the whole table, emit a warning to stderr — a uniformly False
    field suggests a poller field-mapping bug that would silently zero
    tier-3. This is a warning, not a hard failure.
    """
    rows = conn.execute(
        "SELECT model_id, provider, supports_reasoning, reasoning_default_enabled, "
        "reasoning_mode, cost_per_1m_completion, pricing_tbd, context_window "
        "FROM models"
    ).fetchall()

    if not any(row[3] for row in rows):
        print(WARNING_UNIFORMLY_FALSE, file=sys.stderr)

    for (
        model_id,
        provider,
        supports_reasoning,
        reasoning_default_enabled,
        reasoning_mode,
        cost,
        pricing_tbd,
        context_window,
    ) in rows:
        tier = resolve_tier(
            {
                "model_id": model_id,
                "supports_reasoning": bool(supports_reasoning),
                "reasoning_default_enabled": bool(reasoning_default_enabled),
                "reasoning_mode": reasoning_mode,
                "cost_per_1m_completion": cost,
                "pricing_tbd": bool(pricing_tbd),
                "context_window": context_window,
            },
            config.tiering.cheap_completion_max,
            config.tiering.model_tiers,
            config.tiering.tier1_context_max,
        )
        conn.execute(
            "UPDATE models SET tier = ? WHERE model_id = ? AND provider = ?",
            (tier, model_id, provider),
        )
    conn.commit()


def main() -> int:
    cfg = load_config("config.yaml")
    conn = sqlite3.connect(cfg.database.path)
    conn.execute("PRAGMA foreign_keys = ON")
    apply_tiering(conn, cfg)
    conn.close()
    print("tiering applied")
    return 0


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