94 lines
3.0 KiB
Python
94 lines
3.0 KiB
Python
#!/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/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())
|