350 lines
12 KiB
Python
350 lines
12 KiB
Python
"""
|
||
Arena Dashboard — FastAPI backend
|
||
Таблица моделей OmniRoute с Arena Score (text + agent) и liveness-статусом.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import sqlite3
|
||
import time
|
||
import urllib.request
|
||
from contextlib import asynccontextmanager
|
||
from datetime import datetime, timezone
|
||
from typing import Optional
|
||
|
||
from fastapi import FastAPI, Query
|
||
from fastapi.responses import HTMLResponse, JSONResponse
|
||
from fastapi.staticfiles import StaticFiles
|
||
|
||
DB_PATH = os.environ.get("DB_PATH", "/data/arena.db")
|
||
OMNIROUTE_URL = os.environ.get("OMNIROUTE_URL", "http://192.168.1.101:20128")
|
||
OMNIROUTE_KEY = os.environ.get("OMNIROUTE_API_KEY", "")
|
||
ARENA_CACHE_TEXT = os.environ.get("ARENA_CACHE_TEXT", "/data/arena_text.json")
|
||
ARENA_CACHE_AGENT = os.environ.get("ARENA_CACHE_AGENT", "/data/arena_agent.json")
|
||
|
||
|
||
# ─── DB setup ───────────────────────────────────────────────────────────────
|
||
|
||
def get_db() -> sqlite3.Connection:
|
||
db = sqlite3.connect(DB_PATH, check_same_thread=False)
|
||
db.row_factory = sqlite3.Row
|
||
return db
|
||
|
||
|
||
def init_db():
|
||
db = get_db()
|
||
db.executescript("""
|
||
CREATE TABLE IF NOT EXISTS models (
|
||
id TEXT PRIMARY KEY,
|
||
provider TEXT NOT NULL,
|
||
name TEXT NOT NULL,
|
||
arena_text_score TEXT,
|
||
arena_text_rank INTEGER,
|
||
arena_agent_score TEXT,
|
||
arena_agent_rank INTEGER,
|
||
updated_at TEXT
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS checks (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
model_id TEXT NOT NULL,
|
||
checked_at TEXT NOT NULL,
|
||
ok INTEGER NOT NULL,
|
||
latency_ms INTEGER
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_checks_model ON checks(model_id);
|
||
CREATE INDEX IF NOT EXISTS idx_checks_time ON checks(checked_at);
|
||
""")
|
||
db.commit()
|
||
db.close()
|
||
|
||
|
||
# ─── Arena data loading ──────────────────────────────────────────────────────
|
||
|
||
def load_arena_cache(path: str) -> dict[str, dict]:
|
||
"""Возвращает dict {normalized_model_name: {rank, score}} из кеш-файла арены."""
|
||
if not os.path.exists(path):
|
||
return {}
|
||
try:
|
||
with open(path) as f:
|
||
data = json.load(f)
|
||
result = {}
|
||
for m in data.get("models", []):
|
||
result[m["model"].lower()] = m
|
||
return result
|
||
except Exception:
|
||
return {}
|
||
|
||
|
||
# ─── OmniRoute helpers ────────────────────────────────────────────────────────
|
||
|
||
def fetch_omni_models() -> list[dict]:
|
||
req = urllib.request.Request(
|
||
f"{OMNIROUTE_URL}/v1/models",
|
||
headers={"Authorization": f"Bearer {OMNIROUTE_KEY}"},
|
||
)
|
||
with urllib.request.urlopen(req, timeout=15) as r:
|
||
data = json.loads(r.read())
|
||
return data.get("data", [])
|
||
|
||
|
||
def is_skip(model_id: str) -> bool:
|
||
return any(x in model_id.lower() for x in [
|
||
"image", "embed", "flux", "tts", "stt", "audio", "whisper",
|
||
"vision", "guard", "safety", "moderat",
|
||
])
|
||
|
||
|
||
# ─── Arena matching (simple) ─────────────────────────────────────────────────
|
||
|
||
import re
|
||
|
||
|
||
def _norm(s: str) -> str:
|
||
return re.sub(r"[^a-z0-9]", "", s.lower())
|
||
|
||
|
||
def _norm_keep(s: str) -> str:
|
||
return re.sub(r"[^a-z0-9\-]", "", s.lower())
|
||
|
||
|
||
def arena_match(model_id: str, arena: dict[str, dict]) -> Optional[dict]:
|
||
"""Сопоставить модель OmniRoute с записью арены."""
|
||
raw = model_id.split("/")[-1].lower()
|
||
base = re.sub(r"-(it|preview)(:[^:]*)?$", "", raw)
|
||
base = re.sub(r"(:[^:]*)$", "", base)
|
||
|
||
# 1) точное совпадение нормализованного имени
|
||
bn = _norm(base)
|
||
for k, v in arena.items():
|
||
if _norm(k) == bn:
|
||
return v
|
||
|
||
# 2) без суффиксов тиров
|
||
cleaned = re.sub(r"-(thinking|tiered|disabled|low|medium|high|xhigh|max|ultra|mini)", "", base)
|
||
cleaned = re.sub(r"-(thinking|tiered|disabled|low|medium|high|xhigh|max|ultra|mini)", "", cleaned)
|
||
cn = _norm(cleaned)
|
||
for k, v in arena.items():
|
||
if _norm(k) == cn:
|
||
return v
|
||
|
||
# 3) containment
|
||
for candidate in (base, cleaned):
|
||
gk = _norm_keep(candidate)
|
||
if len(gk) >= 12 and (re.search(r"\d", gk) or "-" in gk):
|
||
for k, v in arena.items():
|
||
am = _norm_keep(k)
|
||
if gk == am or am.endswith("-" + gk) or ("-" + gk) in am or (gk in am and len(gk) >= 12):
|
||
return v
|
||
return None
|
||
|
||
|
||
# ─── model sync ──────────────────────────────────────────────────────────────
|
||
|
||
def sync_models():
|
||
"""Обновить список моделей в DB из OmniRoute + арены."""
|
||
try:
|
||
omni = fetch_omni_models()
|
||
except Exception as e:
|
||
print(f"[sync] OmniRoute error: {e}")
|
||
return
|
||
|
||
arena_text = load_arena_cache(ARENA_CACHE_TEXT)
|
||
arena_agent = load_arena_cache(ARENA_CACHE_AGENT)
|
||
|
||
db = get_db()
|
||
now = datetime.now(timezone.utc).isoformat()
|
||
|
||
for m in omni:
|
||
mid = m["id"]
|
||
if is_skip(mid):
|
||
continue
|
||
# openrouter — только :free
|
||
if mid.startswith("openrouter/") and not mid.endswith(":free"):
|
||
continue
|
||
|
||
provider = mid.split("/")[0]
|
||
name = mid
|
||
|
||
ta = arena_match(mid, arena_text)
|
||
ga = arena_match(mid, arena_agent)
|
||
|
||
db.execute("""
|
||
INSERT INTO models (id, provider, name, arena_text_score, arena_text_rank,
|
||
arena_agent_score, arena_agent_rank, updated_at)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||
ON CONFLICT(id) DO UPDATE SET
|
||
arena_text_score = excluded.arena_text_score,
|
||
arena_text_rank = excluded.arena_text_rank,
|
||
arena_agent_score = excluded.arena_agent_score,
|
||
arena_agent_rank = excluded.arena_agent_rank,
|
||
updated_at = excluded.updated_at
|
||
""", (
|
||
mid, provider, name,
|
||
ta["score"] if ta else None,
|
||
ta["rank"] if ta else None,
|
||
ga["score"] if ga else None,
|
||
ga["rank"] if ga else None,
|
||
now,
|
||
))
|
||
|
||
db.commit()
|
||
db.close()
|
||
print(f"[sync] synced {len(omni)} models at {now}")
|
||
|
||
|
||
# ─── Stats helpers ────────────────────────────────────────────────────────────
|
||
|
||
def model_stats(db: sqlite3.Connection, model_id: str, hours: int = 24) -> dict:
|
||
"""Последний статус и % успехов за N часов."""
|
||
cutoff = datetime.fromtimestamp(time.time() - hours * 3600, timezone.utc).isoformat()
|
||
|
||
last = db.execute(
|
||
"SELECT ok, latency_ms, checked_at FROM checks WHERE model_id=? ORDER BY checked_at DESC LIMIT 1",
|
||
(model_id,)
|
||
).fetchone()
|
||
|
||
rows = db.execute(
|
||
"SELECT ok FROM checks WHERE model_id=? AND checked_at > ?",
|
||
(model_id, cutoff)
|
||
).fetchall()
|
||
|
||
total = len(rows)
|
||
success = sum(1 for r in rows if r["ok"])
|
||
pct = round(100 * success / total) if total else None
|
||
|
||
return {
|
||
"last_ok": bool(last["ok"]) if last else None,
|
||
"last_latency_ms": last["latency_ms"] if last else None,
|
||
"last_checked": last["checked_at"] if last else None,
|
||
"uptime_pct": pct,
|
||
"checks_count": total,
|
||
}
|
||
|
||
|
||
# ─── Lifespan ────────────────────────────────────────────────────────────────
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI):
|
||
init_db()
|
||
sync_models()
|
||
yield
|
||
|
||
|
||
app = FastAPI(title="Arena Dashboard", lifespan=lifespan)
|
||
|
||
|
||
# ─── API routes ───────────────────────────────────────────────────────────────
|
||
|
||
@app.get("/api/models")
|
||
def api_models(
|
||
provider: Optional[str] = Query(None),
|
||
arena_only: bool = Query(False),
|
||
sort: str = Query("arena_text_rank"),
|
||
):
|
||
db = get_db()
|
||
q = "SELECT * FROM models WHERE 1=1"
|
||
params: list = []
|
||
if provider:
|
||
q += " AND provider = ?"
|
||
params.append(provider)
|
||
if arena_only:
|
||
q += " AND arena_text_score IS NOT NULL"
|
||
|
||
rows = db.execute(q, params).fetchall()
|
||
|
||
result = []
|
||
for r in rows:
|
||
mid = r["id"]
|
||
stats = model_stats(db, mid)
|
||
result.append({
|
||
"id": mid,
|
||
"provider": r["provider"],
|
||
"name": r["name"],
|
||
"arena_text_score": r["arena_text_score"],
|
||
"arena_text_rank": r["arena_text_rank"],
|
||
"arena_agent_score": r["arena_agent_score"],
|
||
"arena_agent_rank": r["arena_agent_rank"],
|
||
**stats,
|
||
})
|
||
|
||
db.close()
|
||
|
||
# сортировка
|
||
def sort_key(x):
|
||
if sort == "arena_text_rank":
|
||
return (x["arena_text_rank"] is None, x["arena_text_rank"] or 9999)
|
||
if sort == "uptime_pct":
|
||
return (x["uptime_pct"] is None, -(x["uptime_pct"] or 0))
|
||
if sort == "provider":
|
||
return x["provider"]
|
||
return x.get(sort) or ""
|
||
|
||
result.sort(key=sort_key)
|
||
return JSONResponse(result)
|
||
|
||
|
||
@app.get("/api/providers")
|
||
def api_providers():
|
||
db = get_db()
|
||
rows = db.execute("SELECT DISTINCT provider FROM models ORDER BY provider").fetchall()
|
||
db.close()
|
||
return [r["provider"] for r in rows]
|
||
|
||
|
||
@app.post("/api/sync")
|
||
def api_sync():
|
||
sync_models()
|
||
return {"ok": True}
|
||
|
||
|
||
@app.get("/api/check/{model_id:path}")
|
||
def api_check(model_id: str):
|
||
"""Мгновенная проверка одной модели."""
|
||
import time as _t
|
||
t0 = _t.time()
|
||
ok = False
|
||
latency = None
|
||
try:
|
||
body = json.dumps({
|
||
"model": model_id,
|
||
"messages": [{"role": "user", "content": "What is 2+2? Reply with just the number."}],
|
||
"max_tokens": 5,
|
||
}).encode()
|
||
req = urllib.request.Request(
|
||
f"{OMNIROUTE_URL}/v1/chat/completions",
|
||
data=body,
|
||
headers={
|
||
"Authorization": f"Bearer {OMNIROUTE_KEY}",
|
||
"Content-Type": "application/json",
|
||
},
|
||
)
|
||
with urllib.request.urlopen(req, timeout=30) as r:
|
||
resp = json.loads(r.read())
|
||
content = resp["choices"][0]["message"]["content"]
|
||
ok = bool(content.strip())
|
||
latency = int((_t.time() - t0) * 1000)
|
||
except Exception as e:
|
||
print(f"[check] {model_id}: {e}")
|
||
|
||
now = datetime.now(timezone.utc).isoformat()
|
||
db = get_db()
|
||
db.execute(
|
||
"INSERT INTO checks (model_id, checked_at, ok, latency_ms) VALUES (?,?,?,?)",
|
||
(model_id, now, int(ok), latency),
|
||
)
|
||
db.commit()
|
||
db.close()
|
||
return {"model_id": model_id, "ok": ok, "latency_ms": latency, "checked_at": now}
|
||
|
||
|
||
@app.get("/", response_class=HTMLResponse)
|
||
def index():
|
||
with open("/app/static/index.html") as f:
|
||
return f.read()
|
||
|
||
|
||
app.mount("/static", StaticFiles(directory="/app/static"), name="static")
|