6krrt/proficiency_store.py

"""DB access for the ``proficiency`` table.

Pairs with the pure ``proficiency.py`` the way ``tier.py`` pairs with
``tiering.py``: the blending arithmetic stays testable without a database,
and every write goes through here so ``blended_score`` and ``source`` can
never drift out of step with the two inputs that produced them.

Both writers use this — ``leaderboard.py`` sets priors, ``eval_proficiency.py``
folds in measured runs — and each write re-blends the row it touched.
"""

from __future__ import annotations

import sqlite3
from datetime import datetime, timezone
from typing import Optional

from config import RouterConfig
from proficiency import accumulate, blend


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def ensure_columns(conn: sqlite3.Connection) -> None:
    """Add a column an older router.db predates. Idempotent, and cheap.

    schema.sql is CREATE TABLE IF NOT EXISTS, so it defines a NEW database and
    silently does nothing to an existing one. Anything added later therefore
    needs this, or the first write against a database created last week fails
    with "no such column".
    """
    columns = {row[1] for row in conn.execute("PRAGMA table_info(proficiency)")}
    if "inherited_from" not in columns:
        conn.execute("ALTER TABLE proficiency ADD COLUMN inherited_from TEXT")
        conn.commit()
        _backfill_inherited(conn)


def _backfill_inherited(conn: sqlite3.Connection) -> None:
    """One-time: mark rows this harness could not have measured directly.

    ADD COLUMN gives every existing row NULL, which reads as "measured here"
    and would leave exactly the rows this provenance was added for frozen
    forever -- the migration would ship the fix and none of the repair.

    Provenance that was never recorded cannot be recovered in general. It can
    be for the rows that matter, and not by guessing: ``eval_identities``
    selects ``latency_class='standard'`` rows, plus flex rows that have NO
    standard equivalent. So a flex row WITH a standard equivalent was never a
    candidate for direct evaluation, whatever its sample count says. That is
    the harness's own selection rule read backwards.

    Anything else keeps NULL, which is the safe direction: it means "do not
    overwrite", so a real measurement is never lost to this.
    """
    pairs = conn.execute(
        """
        SELECT v.model_id, v.provider, s.model_id
          FROM models v
          JOIN models s
            ON s.base_model_id = v.base_model_id
           AND s.provider = v.provider
           AND s.reasoning_mode = v.reasoning_mode
           AND s.context_variant = v.context_variant
         WHERE v.latency_class = 'flex' AND s.latency_class = 'standard'
        """
    ).fetchall()
    for variant_id, provider, source_id in pairs:
        conn.execute(
            """
            UPDATE proficiency SET inherited_from = ?
             WHERE model_id = ? AND provider = ? AND inherited_from IS NULL
            """,
            (source_id, variant_id, provider),
        )
    conn.commit()


def _read_row(
    conn: sqlite3.Connection, model_id: str, provider: str, category: str
) -> Optional[sqlite3.Row]:
    conn.row_factory = sqlite3.Row
    return conn.execute(
        """
        SELECT leaderboard_score, self_eval_score, self_eval_samples, inherited_from
          FROM proficiency
         WHERE model_id = ? AND provider = ? AND category = ?
        """,
        (model_id, provider, category),
    ).fetchone()


def _write(
    conn: sqlite3.Connection,
    cfg: RouterConfig,
    model_id: str,
    provider: str,
    category: str,
    leaderboard_score: Optional[float],
    self_eval_score: Optional[float],
    self_eval_samples: int,
    inherited_from: Optional[str] = None,
) -> None:
    ensure_columns(conn)
    blended, source = blend(
        leaderboard_score,
        self_eval_score,
        self_eval_samples,
        leaderboard_weight=cfg.proficiency.leaderboard_weight,
        self_eval_weight=cfg.proficiency.self_eval_weight,
        min_samples=cfg.proficiency.self_eval_min_samples,
    )
    conn.execute(
        """
        INSERT INTO proficiency (
            model_id, provider, category, leaderboard_score,
            self_eval_score, self_eval_samples, blended_score, source,
            inherited_from, last_updated
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ON CONFLICT(model_id, provider, category) DO UPDATE SET
            leaderboard_score = excluded.leaderboard_score,
            self_eval_score = excluded.self_eval_score,
            self_eval_samples = excluded.self_eval_samples,
            blended_score = excluded.blended_score,
            source = excluded.source,
            inherited_from = excluded.inherited_from,
            last_updated = excluded.last_updated
        """,
        (
            model_id,
            provider,
            category,
            leaderboard_score,
            self_eval_score,
            self_eval_samples,
            blended,
            source,
            inherited_from,
            _now(),
        ),
    )


def set_leaderboard(
    conn: sqlite3.Connection,
    cfg: RouterConfig,
    model_id: str,
    provider: str,
    category: str,
    score: Optional[float],
) -> None:
    """Set a category's leaderboard prior, preserving any self-eval history."""
    existing = _read_row(conn, model_id, provider, category)
    _write(
        conn,
        cfg,
        model_id,
        provider,
        category,
        score,
        existing["self_eval_score"] if existing else None,
        existing["self_eval_samples"] if existing else 0,
        # A prior says nothing about where the self-eval half came from.
        existing["inherited_from"] if existing else None,
    )


def add_self_eval(
    conn: sqlite3.Connection,
    cfg: RouterConfig,
    model_id: str,
    provider: str,
    category: str,
    scores: list[float],
) -> None:
    """Fold an eval run's per-task scores into the running self-eval mean.

    Accumulates rather than replaces, so re-running the harness tightens the
    estimate instead of discarding what came before — and so
    ``self_eval_samples`` keeps meaning "how much evidence stands behind
    this", which is what the blending threshold is gating on.
    """
    # Clamped at the one place every write passes through, which is what this
    # module exists to be. score_judge already bounded its output; score_code
    # did not, and a harness bug let a model's own demo output push a two-check
    # task to 1.50. A blended_score above 1.0 raises `best` in
    # rank_candidates and shifts every other candidate's quality band, so the
    # damage is not confined to the row that carries the bad number.
    scores = [min(1.0, max(0.0, s)) for s in scores]

    existing = _read_row(conn, model_id, provider, category)
    prev_score = existing["self_eval_score"] if existing else None
    prev_samples = existing["self_eval_samples"] if existing else 0

    new_score, new_samples = accumulate(prev_score, prev_samples, scores)
    _write(
        conn,
        cfg,
        model_id,
        provider,
        category,
        existing["leaderboard_score"] if existing else None,
        new_score,
        new_samples,
    )


def propagate_to_variants(
    conn: sqlite3.Connection, cfg: RouterConfig, source_model_id: str, provider: str
) -> int:
    """Copy an evaluated row's scores onto its equivalent serving variants.

    A ``-flex`` row is the same weights, same reasoning setting, same context
    pool, on a different queue — so its answer quality IS the standard row's
    and inheriting is correct. A ``-fast`` row is NOT equivalent (reasoning
    off or capped), and neither is a ``-short`` row, so matching is on
    (base_model_id, reasoning_mode, context_variant) rather than on family
    alone. Inheriting across those would attribute reasoning-on quality to a
    reasoning-off row.

    Anything measured DIRECTLY keeps its own scores; inheritance is the
    fallback, never an overwrite. A row that was inherited before is refreshed,
    which needs ``inherited_from`` to tell the two apart -- both carry
    self_eval_samples > 0, so the old "samples > 0 means leave it alone" test
    could not distinguish them and froze every variant at its first
    inheritance. Observed: kimi-k3 moved from 0.85 (n=2) to 0.957 (n=7) while
    kimi-k3-flex sat at 0.85 with a week-old timestamp, and the run reported
    "propagated 0 inherited rows". Flex rows serve `auto:batch`, so those
    requests were ranking on scores their family had left behind.

    Returns the number of rows written.
    """
    conn.row_factory = sqlite3.Row
    source_rows = conn.execute(
        """
        SELECT category, leaderboard_score, self_eval_score, self_eval_samples
          FROM proficiency
         WHERE model_id = ? AND provider = ?
        """,
        (source_model_id, provider),
    ).fetchall()
    if not source_rows:
        return 0

    variants = conn.execute(
        """
        SELECT v.model_id FROM models v
          JOIN models src
            ON src.base_model_id = v.base_model_id
           AND src.provider = v.provider
           AND src.reasoning_mode = v.reasoning_mode
           AND src.context_variant = v.context_variant
         WHERE src.model_id = ? AND v.provider = ? AND v.model_id != ?
        """,
        (source_model_id, provider, source_model_id),
    ).fetchall()

    written = 0
    for variant in variants:
        for row in source_rows:
            existing = _read_row(conn, variant["model_id"], provider, row["category"])
            measured_here = (
                existing is not None
                and (existing["self_eval_samples"] or 0) > 0
                and not existing["inherited_from"]
            )
            if measured_here:
                continue  # its own measurement outranks the family's
            _write(
                conn,
                cfg,
                variant["model_id"],
                provider,
                row["category"],
                row["leaderboard_score"],
                row["self_eval_score"],
                row["self_eval_samples"],
                source_model_id,
            )
            written += 1
    return written