197 lines
6.7 KiB
Python
197 lines
6.7 KiB
Python
"""Core DB primitives for the ``proficiency`` table.
|
|
|
|
Kept separate from the public writer surface so the empirical-Bayes
|
|
recomputation module can share the same single write path without creating
|
|
a circular import.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from collections.abc import Generator
|
|
from contextlib import contextmanager
|
|
from datetime import datetime, timezone
|
|
|
|
from config import RouterConfig
|
|
from proficiency import blend
|
|
|
|
|
|
@contextmanager
|
|
def _proficiency_tx(
|
|
conn: sqlite3.Connection,
|
|
) -> Generator[None, None, None]:
|
|
"""Run a block inside a single SQLite transaction.
|
|
|
|
If the caller has already started a transaction, this manager is a no-op
|
|
except for rolling back on exception. Otherwise it ensures columns are
|
|
present, issues one ``conn.commit()`` at successful exit, and rolls back on
|
|
error. Python's sqlite3 module implicitly begins a transaction before the
|
|
first data-modifying statement, so we do not issue an explicit ``BEGIN``.
|
|
"""
|
|
own_transaction = not conn.in_transaction
|
|
if own_transaction:
|
|
ensure_columns(conn)
|
|
try:
|
|
yield
|
|
if own_transaction:
|
|
conn.commit()
|
|
except Exception:
|
|
if own_transaction:
|
|
conn.rollback()
|
|
raise
|
|
|
|
|
|
def _now() -> str:
|
|
return datetime.now(timezone.utc).isoformat()
|
|
|
|
|
|
def ensure_columns(conn: sqlite3.Connection) -> None:
|
|
"""Add columns 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".
|
|
|
|
This function may run DDL; callers that need atomic behaviour should call
|
|
it *before* beginning an explicit transaction.
|
|
"""
|
|
columns = {row[1] for row in conn.execute("PRAGMA table_info(proficiency)")}
|
|
migrated = False
|
|
|
|
if "inherited_from" not in columns:
|
|
conn.execute("ALTER TABLE proficiency ADD COLUMN inherited_from TEXT")
|
|
conn.commit()
|
|
_backfill_inherited(conn)
|
|
if "outcome_score" not in columns:
|
|
conn.execute("ALTER TABLE proficiency ADD COLUMN outcome_score REAL")
|
|
migrated = True
|
|
if "outcome_samples" not in columns:
|
|
conn.execute(
|
|
"ALTER TABLE proficiency ADD COLUMN outcome_samples INTEGER DEFAULT 0"
|
|
)
|
|
migrated = True
|
|
|
|
if migrated and not conn.in_transaction:
|
|
conn.commit()
|
|
|
|
|
|
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
|
|
) -> sqlite3.Row | None:
|
|
conn.row_factory = sqlite3.Row
|
|
return conn.execute(
|
|
"""
|
|
SELECT leaderboard_score, self_eval_score, self_eval_samples,
|
|
outcome_score, outcome_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: float | None,
|
|
self_eval_score: float | None,
|
|
self_eval_samples: int,
|
|
outcome_score: float | None = None,
|
|
outcome_samples: int = 0,
|
|
inherited_from: str | None = None,
|
|
*,
|
|
_blended_override: bool = False,
|
|
_blended_score: float | None = None,
|
|
_blended_source: str | None = None,
|
|
) -> None:
|
|
ensure_columns(conn)
|
|
|
|
if _blended_override:
|
|
blended, source = _blended_score, _blended_source
|
|
else:
|
|
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, outcome_score, outcome_samples, 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,
|
|
outcome_score = excluded.outcome_score,
|
|
outcome_samples = excluded.outcome_samples,
|
|
last_updated = excluded.last_updated
|
|
""",
|
|
(
|
|
model_id,
|
|
provider,
|
|
category,
|
|
leaderboard_score,
|
|
self_eval_score,
|
|
self_eval_samples,
|
|
blended,
|
|
source,
|
|
inherited_from,
|
|
outcome_score,
|
|
outcome_samples,
|
|
_now(),
|
|
),
|
|
)
|