feat: учёт token usage в checks с агрегатами и UI
Build and deploy / build-deploy (push) Has been cancelled

- миграция checks: nullable input_tokens/output_tokens/total_tokens/usage_json
  (идемпотентно, старые строки остаются NULL; probe мигрирует сам — cron стартует раньше uvicorn)
- probe.py и /api/check фиксируют usage ДО валидации контента: токены
  сохраняются при пустом/сломанном ответе; таймауты и ответы без usage = NULL, не нули
- usage_json хранит сырой usage провайдера (OpenAI/Anthropic/camelCase форматы)
- /api/checks и сегментная история несут токены; новый /api/token_stats:
  rolling 24h + all time, покрытие known/unknown, разбивка по провайдерам
  (провайдер из models с fallback на префикс model_id)
- UI: статы «Токены 24h» и «Покрытие usage» в шапке, токены в тултипах
  сегментов и логе проверок
- статика через __file__-относительный путь (работает и в /app/app, и локально)
- тесты (30, реальные HTTP против мока OmniRoute): миграция/идемпотентность,
  NULL при missing/zero usage, usage при упавшем контенте, таймаут, агрегаты,
  оба пути записи (probe.run + /api/check)
This commit is contained in:
ai_coder
2026-10-25 17:27:38 +00:00
parent 0a0455af73
commit 5c8c13891f
11 changed files with 1007 additions and 21 deletions
+4
View File
@@ -0,0 +1,4 @@
.venv/
.pytest_cache/
__pycache__/
.smoke/
+146 -9
View File
@@ -17,6 +17,8 @@ from fastapi import FastAPI, Query
from fastapi.responses import HTMLResponse, JSONResponse from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles from fastapi.staticfiles import StaticFiles
from usage_utils import extract_usage
DB_PATH = os.environ.get("DB_PATH", "/data/arena.db") DB_PATH = os.environ.get("DB_PATH", "/data/arena.db")
OMNIROUTE_URL = os.environ.get("OMNIROUTE_URL", "http://192.168.1.101:20128") OMNIROUTE_URL = os.environ.get("OMNIROUTE_URL", "http://192.168.1.101:20128")
OMNIROUTE_KEY = os.environ.get("OMNIROUTE_API_KEY", "") OMNIROUTE_KEY = os.environ.get("OMNIROUTE_API_KEY", "")
@@ -25,6 +27,10 @@ ARENA_CACHE_AGENT = os.environ.get("ARENA_CACHE_AGENT", "/data/arena_agent.json"
# Сколько последних проверок показывать сегментами в колонке Uptime. # Сколько последних проверок показывать сегментами в колонке Uptime.
HISTORY_LIMIT = int(os.environ.get("HISTORY_LIMIT", "40")) HISTORY_LIMIT = int(os.environ.get("HISTORY_LIMIT", "40"))
# Статика — относительно этого файла: работает и в контейнере (/app/app),
# и при локальном запуске/тестах из чекаута репозитория.
STATIC_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static")
# ─── DB setup ─────────────────────────────────────────────────────────────── # ─── DB setup ───────────────────────────────────────────────────────────────
@@ -98,7 +104,11 @@ def init_db():
model_id TEXT NOT NULL, model_id TEXT NOT NULL,
checked_at TEXT NOT NULL, checked_at TEXT NOT NULL,
ok INTEGER NOT NULL, ok INTEGER NOT NULL,
latency_ms INTEGER latency_ms INTEGER,
input_tokens INTEGER,
output_tokens INTEGER,
total_tokens INTEGER,
usage_json TEXT
); );
CREATE INDEX IF NOT EXISTS idx_checks_model ON checks(model_id); CREATE INDEX IF NOT EXISTS idx_checks_model ON checks(model_id);
@@ -112,6 +122,17 @@ def init_db():
if "error" not in cols: if "error" not in cols:
db.execute("ALTER TABLE checks ADD COLUMN error TEXT") db.execute("ALTER TABLE checks ADD COLUMN error TEXT")
# Идемпотентная миграция: учёт token usage. Старые строки остаются NULL
# (неизвестное потребление не подменяется нулями), повторный запуск — no-op.
for col, decl in (
("input_tokens", "INTEGER"),
("output_tokens", "INTEGER"),
("total_tokens", "INTEGER"),
("usage_json", "TEXT"),
):
if col not in cols:
db.execute(f"ALTER TABLE checks ADD COLUMN {col} {decl}")
db.commit() db.commit()
db.close() db.close()
@@ -360,7 +381,7 @@ def model_stats(db: sqlite3.Connection, model_id: str, hours: int = 24) -> dict:
).fetchone() ).fetchone()
rows = db.execute( rows = db.execute(
"SELECT ok, latency_ms, checked_at, error FROM checks " "SELECT ok, latency_ms, checked_at, error, input_tokens, output_tokens, total_tokens FROM checks "
"WHERE model_id=? AND checked_at > ? ORDER BY checked_at ASC", "WHERE model_id=? AND checked_at > ? ORDER BY checked_at ASC",
(model_id, cutoff) (model_id, cutoff)
).fetchall() ).fetchall()
@@ -370,13 +391,16 @@ def model_stats(db: sqlite3.Connection, model_id: str, hours: int = 24) -> dict:
pct = round(100 * success / total) if total else None pct = round(100 * success / total) if total else None
# История последних проверок (старые → новые) для сегментной витрины. # История последних проверок (старые → новые) для сегментной витрины.
# Каждый элемент — детали проверки: ok, latency, время, текст ошибки. # Каждый элемент — детали проверки: ok, latency, время, текст ошибки, токены.
history = [ history = [
{ {
"ok": bool(r["ok"]), "ok": bool(r["ok"]),
"latency_ms": r["latency_ms"], "latency_ms": r["latency_ms"],
"at": r["checked_at"], "at": r["checked_at"],
"error": r["error"], "error": r["error"],
"input_tokens": r["input_tokens"],
"output_tokens": r["output_tokens"],
"total_tokens": r["total_tokens"],
} }
for r in rows[-HISTORY_LIMIT:] for r in rows[-HISTORY_LIMIT:]
] ]
@@ -391,6 +415,93 @@ def model_stats(db: sqlite3.Connection, model_id: str, hours: int = 24) -> dict:
} }
# ─── Token usage aggregates ───────────────────────────────────────────────────
def _token_agg_rows(db: sqlite3.Connection, where: str, params: list) -> list[dict]:
"""GROUP BY провайдер по checks с суммами токенов и покрытием known/unknown."""
# Провайдер — из таблицы models; если модели уже нет, выводим из model_id.
prov_expr = (
"COALESCE(m.provider, "
"CASE WHEN instr(c.model_id, '/') > 0 "
"THEN substr(c.model_id, 1, instr(c.model_id, '/') - 1) "
"ELSE 'unknown' END)"
)
return [
{
"provider": r["provider"],
"checks": r["checks"],
"known": r["known"],
"unknown": r["unknown"],
"coverage_pct": round(100 * r["known"] / r["checks"]) if r["checks"] else None,
"input_tokens": r["input_tokens"],
"output_tokens": r["output_tokens"],
"total_tokens": r["total_tokens"],
}
for r in db.execute(
f"""
SELECT {prov_expr} AS provider,
COUNT(*) AS checks,
SUM(CASE WHEN c.total_tokens IS NOT NULL THEN 1 ELSE 0 END) AS known,
SUM(CASE WHEN c.total_tokens IS NULL THEN 1 ELSE 0 END) AS unknown,
SUM(c.input_tokens) AS input_tokens,
SUM(c.output_tokens) AS output_tokens,
SUM(c.total_tokens) AS total_tokens
FROM checks c
LEFT JOIN models m ON c.model_id = m.id
WHERE {where}
GROUP BY provider
ORDER BY total_tokens DESC
""",
params,
).fetchall()
]
def _merge_aggs(rows: list[dict]) -> dict:
"""Свернуть группы провайдеров в одну сводку."""
checks = sum(r["checks"] for r in rows)
known = sum(r["known"] for r in rows)
return {
"checks": checks,
"known": known,
"unknown": checks - known,
"coverage_pct": round(100 * known / checks) if checks else None,
"input_tokens": sum(r["input_tokens"] or 0 for r in rows),
"output_tokens": sum(r["output_tokens"] or 0 for r in rows),
"total_tokens": sum(r["total_tokens"] or 0 for r in rows),
}
def token_stats(db: sqlite3.Connection) -> dict:
"""Сводка token usage: скользящие 24h и всё время + разбивка по провайдерам."""
cutoff24 = datetime.fromtimestamp(time.time() - 24 * 3600, timezone.utc).isoformat()
now = datetime.now(timezone.utc).isoformat()
prov24 = _token_agg_rows(db, "c.checked_at > ?", [cutoff24])
prov_all = _token_agg_rows(db, "1=1", [])
all_map = {r["provider"]: r for r in prov_all}
return {
"generated_at": now,
"window_24h": _merge_aggs(prov24),
"all_time": _merge_aggs(prov_all),
"providers": [
{
"provider": r["provider"],
"window_24h": r,
"all_time": all_map.get(r["provider"], _merge_aggs([])),
}
for r in prov24
] + [
{
"provider": r["provider"],
"window_24h": _merge_aggs([]),
"all_time": r,
}
for r in prov_all if r["provider"] not in {p["provider"] for p in prov24}
],
}
# ─── Lifespan ──────────────────────────────────────────────────────────────── # ─── Lifespan ────────────────────────────────────────────────────────────────
@asynccontextmanager @asynccontextmanager
@@ -493,6 +604,8 @@ def api_check(model_id: str):
ok = False ok = False
latency = None latency = None
error = None error = None
usage = {"input_tokens": None, "output_tokens": None,
"total_tokens": None, "usage_json": None}
try: try:
body = json.dumps({ body = json.dumps({
"model": model_id, "model": model_id,
@@ -509,6 +622,9 @@ def api_check(model_id: str):
) )
with urllib.request.urlopen(req, timeout=30) as r: with urllib.request.urlopen(req, timeout=30) as r:
resp = json.loads(r.read()) resp = json.loads(r.read())
# Usage фиксируем ДО валидации контента: даже если парсинг choices
# упадёт, реально потраченные токены сохраняются в checks.
usage = extract_usage(resp)
content = resp["choices"][0]["message"]["content"] content = resp["choices"][0]["message"]["content"]
ok = bool(content.strip()) ok = bool(content.strip())
except Exception as e: except Exception as e:
@@ -520,8 +636,12 @@ def api_check(model_id: str):
now = datetime.now(timezone.utc).isoformat() now = datetime.now(timezone.utc).isoformat()
db = get_db() db = get_db()
db.execute( db.execute(
"INSERT INTO checks (model_id, checked_at, ok, latency_ms, error) VALUES (?,?,?,?,?)", "INSERT INTO checks (model_id, checked_at, ok, latency_ms, error, "
(model_id, now, int(ok), latency, error), "input_tokens, output_tokens, total_tokens, usage_json) "
"VALUES (?,?,?,?,?,?,?,?,?)",
(model_id, now, int(ok), latency, error,
usage["input_tokens"], usage["output_tokens"],
usage["total_tokens"], usage["usage_json"]),
) )
db.commit() db.commit()
stats = model_stats(db, model_id) stats = model_stats(db, model_id)
@@ -531,6 +651,9 @@ def api_check(model_id: str):
"ok": ok, "ok": ok,
"latency_ms": latency, "latency_ms": latency,
"checked_at": now, "checked_at": now,
"input_tokens": usage["input_tokens"],
"output_tokens": usage["output_tokens"],
"total_tokens": usage["total_tokens"],
"history": stats["history"], "history": stats["history"],
"uptime_pct": stats["uptime_pct"], "uptime_pct": stats["uptime_pct"],
} }
@@ -542,8 +665,8 @@ def api_checks(model_id: str):
limit = int(os.environ.get("CHECKS_LIMIT", "100")) limit = int(os.environ.get("CHECKS_LIMIT", "100"))
db = get_db() db = get_db()
rows = db.execute( rows = db.execute(
"SELECT ok, latency_ms, checked_at, error FROM checks " "SELECT ok, latency_ms, checked_at, error, input_tokens, output_tokens, total_tokens, usage_json "
"WHERE model_id=? ORDER BY checked_at DESC LIMIT ?", "FROM checks WHERE model_id=? ORDER BY checked_at DESC LIMIT ?",
(model_id, limit), (model_id, limit),
).fetchall() ).fetchall()
db.close() db.close()
@@ -555,17 +678,31 @@ def api_checks(model_id: str):
"latency_ms": r["latency_ms"], "latency_ms": r["latency_ms"],
"checked_at": r["checked_at"], "checked_at": r["checked_at"],
"error": r["error"], "error": r["error"],
"input_tokens": r["input_tokens"],
"output_tokens": r["output_tokens"],
"total_tokens": r["total_tokens"],
"usage_json": r["usage_json"],
} }
for r in rows for r in rows
], ],
} }
@app.get("/api/token_stats")
def api_token_stats():
"""Агрегат token usage: скользящие 24h и всё время, покрытие known/unknown,
разбивка по провайдерам."""
db = get_db()
result = token_stats(db)
db.close()
return JSONResponse(result)
@app.get("/", response_class=HTMLResponse) @app.get("/", response_class=HTMLResponse)
def index(): def index():
with open("/app/app/static/index.html") as f: with open(os.path.join(STATIC_DIR, "index.html")) as f:
return f.read() return f.read()
app.mount("/static", StaticFiles(directory="/app/app/static"), name="static") app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
+42 -7
View File
@@ -12,6 +12,8 @@ import time
import urllib.request import urllib.request
from datetime import datetime, timezone from datetime import datetime, timezone
from usage_utils import extract_usage
DB_PATH = os.environ.get("DB_PATH", "/data/arena.db") DB_PATH = os.environ.get("DB_PATH", "/data/arena.db")
OMNIROUTE_URL = os.environ.get("OMNIROUTE_URL", "http://192.168.1.101:20128") OMNIROUTE_URL = os.environ.get("OMNIROUTE_URL", "http://192.168.1.101:20128")
OMNIROUTE_KEY = os.environ.get("OMNIROUTE_API_KEY", "") OMNIROUTE_KEY = os.environ.get("OMNIROUTE_API_KEY", "")
@@ -19,9 +21,17 @@ TIMEOUT = int(os.environ.get("CHECK_TIMEOUT", "30"))
CONCURRENCY = int(os.environ.get("CHECK_CONCURRENCY", "5")) CONCURRENCY = int(os.environ.get("CHECK_CONCURRENCY", "5"))
def check_model(model_id: str) -> tuple[bool, int | None, str | None]: def check_model(model_id: str) -> tuple[bool, int | None, str | None, dict]:
"""Возвращает (ok, latency_ms, error, usage).
usage извлекается ДО валидации контента: если парсинг choices падает,
потраченные токены всё равно сохраняются. Таймаут/ошибка сети без тела
ответа дают usage=None по всем полям (NULL, а не нули).
"""
t0 = time.time() t0 = time.time()
error = None error = None
usage = {"input_tokens": None, "output_tokens": None,
"total_tokens": None, "usage_json": None}
try: try:
body = json.dumps({ body = json.dumps({
"model": model_id, "model": model_id,
@@ -38,6 +48,8 @@ def check_model(model_id: str) -> tuple[bool, int | None, str | None]:
) )
with urllib.request.urlopen(req, timeout=TIMEOUT) as r: with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
resp = json.loads(r.read()) resp = json.loads(r.read())
# Usage — до любой валидации контента.
usage = extract_usage(resp)
content = resp["choices"][0]["message"]["content"] content = resp["choices"][0]["message"]["content"]
ok = bool(content.strip()) ok = bool(content.strip())
except Exception as e: except Exception as e:
@@ -45,11 +57,30 @@ def check_model(model_id: str) -> tuple[bool, int | None, str | None]:
print(f" ✗ {model_id}: {error}", file=sys.stderr) print(f" ✗ {model_id}: {error}", file=sys.stderr)
ok = False ok = False
latency = int((time.time() - t0) * 1000) latency = int((time.time() - t0) * 1000)
return ok, latency, error return ok, latency, error, usage
def _ensure_usage_columns(db: sqlite3.Connection):
"""Идемпотентно добавить колонки token usage, если БД ещё со старой схемой.
Cron стартует раньше uvicorn, поэтому probe не может полагаться на то,
что init_db из main.py уже прогнал миграцию.
"""
cols = [r[1] for r in db.execute("PRAGMA table_info(checks)").fetchall()]
for col, decl in (
("input_tokens", "INTEGER"),
("output_tokens", "INTEGER"),
("total_tokens", "INTEGER"),
("usage_json", "TEXT"),
):
if col not in cols:
db.execute(f"ALTER TABLE checks ADD COLUMN {col} {decl}")
db.commit()
def run(): def run():
db = sqlite3.connect(DB_PATH, check_same_thread=False) db = sqlite3.connect(DB_PATH, check_same_thread=False)
_ensure_usage_columns(db)
models = [r[0] for r in db.execute("SELECT id FROM models").fetchall()] models = [r[0] for r in db.execute("SELECT id FROM models").fetchall()]
db.close() db.close()
@@ -58,8 +89,8 @@ def run():
from concurrent.futures import ThreadPoolExecutor, as_completed from concurrent.futures import ThreadPoolExecutor, as_completed
def do(mid): def do(mid):
ok, lat, err = check_model(mid) ok, lat, err, usage = check_model(mid)
return mid, ok, lat, err return mid, ok, lat, err, usage
db = sqlite3.connect(DB_PATH, check_same_thread=False) db = sqlite3.connect(DB_PATH, check_same_thread=False)
now = datetime.now(timezone.utc).isoformat() now = datetime.now(timezone.utc).isoformat()
@@ -68,10 +99,14 @@ def run():
futs = {ex.submit(do, m): m for m in models} futs = {ex.submit(do, m): m for m in models}
done = ok_count = 0 done = ok_count = 0
for f in as_completed(futs): for f in as_completed(futs):
mid, ok, lat, err = f.result() mid, ok, lat, err, usage = f.result()
db.execute( db.execute(
"INSERT INTO checks (model_id, checked_at, ok, latency_ms, error) VALUES (?,?,?,?,?)", "INSERT INTO checks (model_id, checked_at, ok, latency_ms, error, "
(mid, now, int(ok), lat, err), "input_tokens, output_tokens, total_tokens, usage_json) "
"VALUES (?,?,?,?,?,?,?,?,?)",
(mid, now, int(ok), lat, err,
usage["input_tokens"], usage["output_tokens"],
usage["total_tokens"], usage["usage_json"]),
) )
done += 1 done += 1
if ok: if ok:
+43 -5
View File
@@ -245,6 +245,8 @@
<div class="stat"><div class="val" id="s-ranked">—</div><div class="lbl">В рейтинге</div></div> <div class="stat"><div class="val" id="s-ranked">—</div><div class="lbl">В рейтинге</div></div>
<div class="stat"><div class="val" id="s-alive">—</div><div class="lbl">Живых</div></div> <div class="stat"><div class="val" id="s-alive">—</div><div class="lbl">Живых</div></div>
<div class="stat"><div class="val" id="s-checked">—</div><div class="lbl">Проверено сегодня</div></div> <div class="stat"><div class="val" id="s-checked">—</div><div class="lbl">Проверено сегодня</div></div>
<div class="stat"><div class="val" id="s-tokens-24h">—</div><div class="lbl">Токены 24ч</div></div>
<div class="stat"><div class="val" id="s-tokens-cover">—</div><div class="lbl">Покрытие usage</div></div>
</div> </div>
<div class="search-wrap"> <div class="search-wrap">
@@ -351,9 +353,10 @@ function showSegTip(el) {
const tip = segTip(); const tip = segTip();
const ok = data.ok; const ok = data.ok;
const lat = data.latency_ms !== null && data.latency_ms !== undefined ? `${data.latency_ms} мс` : '—'; const lat = data.latency_ms !== null && data.latency_ms !== undefined ? `${data.latency_ms} мс` : '—';
const toks = usageStr(data);
tip.innerHTML = ` tip.innerHTML = `
<div class="tip-head ${ok ? 'tip-ok' : 'tip-fail'}">${ok ? '✓ Успех' : '✗ Провал'}</div> <div class="tip-head ${ok ? 'tip-ok' : 'tip-fail'}">${ok ? '✓ Успех' : '✗ Провал'}</div>
<div class="tip-meta">${fmtTime(data.at)} · ${lat}</div> <div class="tip-meta">${fmtTime(data.at)} · ${lat}${toks ? ` · ${escapeHtml(toks)}` : ''}</div>
${data.error ? `<div class="tip-err">${escapeHtml(data.error)}</div>` : ''} ${data.error ? `<div class="tip-err">${escapeHtml(data.error)}</div>` : ''}
`; `;
tip.style.display = 'block'; tip.style.display = 'block';
@@ -413,6 +416,37 @@ function renderRow(m) {
</tr>`; </tr>`;
} }
function fmtTokens(n) {
if (n === null || n === undefined) return null;
if (n >= 1e9) return (n / 1e9).toFixed(1) + 'B';
if (n >= 1e6) return (n / 1e6).toFixed(1) + 'M';
if (n >= 1e3) return (n / 1e3).toFixed(1) + 'k';
return String(n);
}
// Строка токенов для тултипа/лога: "in 12 → out 5" или null если usage неизвестен.
function usageStr(u) {
if (!u || u.total_tokens === null || u.total_tokens === undefined) return null;
const inp = u.input_tokens !== null && u.input_tokens !== undefined ? u.input_tokens : '?';
const out = u.output_tokens !== null && u.output_tokens !== undefined ? u.output_tokens : '?';
return `in ${inp} → out ${out} (${u.total_tokens})`;
}
async function loadTokenStats() {
try {
const res = await fetch('/api/token_stats');
const ts = await res.json();
const w = ts.window_24h || {};
const t24 = document.getElementById('s-tokens-24h');
const cov = document.getElementById('s-tokens-cover');
const tot = fmtTokens(w.total_tokens);
t24.textContent = tot === null ? '—' : tot;
t24.title = `in ${fmtTokens(w.input_tokens) ?? '—'} / out ${fmtTokens(w.output_tokens) ?? '—'} · всё время: ${fmtTokens((ts.all_time || {}).total_tokens) ?? '—'}`;
cov.textContent = w.coverage_pct === null || w.coverage_pct === undefined ? '—' : w.coverage_pct + '%';
cov.title = `24ч: известно ${w.known ?? 0} из ${w.checks ?? 0} проверок (неизвестно ${w.unknown ?? 0})`;
} catch (e) { /* статы токенов не критичны */ }
}
function updateStats(data) { function updateStats(data) {
document.getElementById('s-total').textContent = data.length; document.getElementById('s-total').textContent = data.length;
document.getElementById('s-ranked').textContent = data.filter(m => m.arena_text_score).length; document.getElementById('s-ranked').textContent = data.filter(m => m.arena_text_score).length;
@@ -498,6 +532,7 @@ async function loadData() {
document.getElementById('last-updated').textContent = `Обновлено: ${now.toLocaleTimeString('ru')}`; document.getElementById('last-updated').textContent = `Обновлено: ${now.toLocaleTimeString('ru')}`;
applyFilters(); applyFilters();
loadTokenStats();
} }
async function loadProviders() { async function loadProviders() {
@@ -527,6 +562,7 @@ async function checkOne(modelId, btn) {
if (data.uptime_pct !== undefined) allData[idx].uptime_pct = data.uptime_pct; if (data.uptime_pct !== undefined) allData[idx].uptime_pct = data.uptime_pct;
} }
applyFilters(); applyFilters();
loadTokenStats();
} finally { } finally {
btn.textContent = '▶ Тест'; btn.textContent = '▶ Тест';
btn.classList.remove('checking'); btn.classList.remove('checking');
@@ -562,15 +598,17 @@ async function openCheckHistory(modelId, btn) {
if (!checks.length) { if (!checks.length) {
body.innerHTML = '<div class="no-data" style="text-align:center;padding:20px;">Нет данных о проверках</div>'; body.innerHTML = '<div class="no-data" style="text-align:center;padding:20px;">Нет данных о проверках</div>';
} else { } else {
body.innerHTML = checks.map(c => ` body.innerHTML = checks.map(c => {
const toks = usageStr(c);
return `
<div class="check-log-item"> <div class="check-log-item">
<div class="log-line1"> <div class="log-line1">
<span class="log-status ${c.ok ? 'ok' : 'fail'}">${c.ok ? '✓ Успех' : '✗ Провал'}</span> <span class="log-status ${c.ok ? 'ok' : 'fail'}">${c.ok ? '✓ Успех' : '✗ Провал'}</span>
<span class="log-meta">${fmtTime(c.checked_at)} · ${c.latency_ms !== null && c.latency_ms !== undefined ? c.latency_ms + ' мс' : '—'}</span> <span class="log-meta">${fmtTime(c.checked_at)} · ${c.latency_ms !== null && c.latency_ms !== undefined ? c.latency_ms + ' мс' : '—'}${toks ? ` · ${escapeHtml(toks)} т.` : ''}</span>
</div> </div>
${c.error ? `<div class="log-err">${escapeHtml(c.error)}</div>` : ''} ${c.error ? `<div class="log-err">${escapeHtml(c.error)}</div>` : ''}
</div> </div>`;
`).join(''); }).join('');
} }
modal.classList.add('open'); modal.classList.add('open');
+83
View File
@@ -0,0 +1,83 @@
"""
Общий модуль извлечения token usage из ответа OmniRoute.
Используется и cron-пробой (probe.py), и ручной проверкой (main.py).
Правила:
* usage извлекается ДО валидации контента — чтобы провал content-парсинга
не терял реально потраченные токены;
* если провайдер не отдал usage (или отдал полностью нулевой) — все поля
остаются NULL: неизвестное потребление не подменяется нулями;
* usage_json хранит сырой словарь usage как пришел от провайдера.
"""
from __future__ import annotations
import json
from typing import Optional
# Возможные имена полей в разных диалектах (OpenAI / Anthropic / camelCase).
_INPUT_KEYS = ("prompt_tokens", "input_tokens", "promptTokens", "inputTokens")
_OUTPUT_KEYS = ("completion_tokens", "output_tokens", "completionTokens", "outputTokens")
_TOTAL_KEYS = ("total_tokens", "totalTokens")
def _first_int(usage: dict, keys: tuple[str, ...]) -> Optional[int]:
"""Первое числовое значение из ключей; не-числа и None игнорируются."""
for k in keys:
v = usage.get(k)
if isinstance(v, bool): # bool — подкласс int, отсекаем явно
continue
if isinstance(v, (int, float)) and v >= 0:
return int(v)
return None
def extract_usage(resp) -> dict:
"""Достать usage из тела ответа chat/completions.
Возвращает dict с ключами input_tokens / output_tokens / total_tokens /
usage_json. Отсутствующий или вырожденный (полностью нулевой) usage дает
None во всех полях — NULL, а не нулевая оценка.
"""
empty = {"input_tokens": None, "output_tokens": None,
"total_tokens": None, "usage_json": None}
usage = None
if isinstance(resp, dict):
u = resp.get("usage")
if isinstance(u, dict):
usage = u
elif isinstance(u, str):
# редко: usage пришел строкой-JSON
try:
parsed = json.loads(u)
if isinstance(parsed, dict):
usage = parsed
except (ValueError, TypeError):
usage = None
if not isinstance(usage, dict):
return empty
inp = _first_int(usage, _INPUT_KEYS)
out = _first_int(usage, _OUTPUT_KEYS)
tot = _first_int(usage, _TOTAL_KEYS)
# Полностью нулевой usage (0/0/0) трактуем как «не отдал» — провайдеры
# так отвечают в ошибочных путях; реальные токены всегда > 0.
if not any(v for v in (inp, out, tot) if v is not None):
return empty
# total — точная арифметика по отчетным значениям, не оценка.
if tot is None and inp is not None and out is not None:
tot = inp + out
try:
usage_json = json.dumps(usage, ensure_ascii=False, sort_keys=True)
except (TypeError, ValueError):
usage_json = None
return {
"input_tokens": inp,
"output_tokens": out,
"total_tokens": tot,
"usage_json": usage_json,
}
+180
View File
@@ -0,0 +1,180 @@
"""Общие фикстуры: реальный мок-сервер OmniRoute на localhost + чистые модули."""
from __future__ import annotations
import json
import os
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from types import SimpleNamespace
import pytest
from fastapi.testclient import TestClient
REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
APP_DIR = os.path.join(REPO, "app")
if APP_DIR not in sys.path:
sys.path.insert(0, APP_DIR)
MODULES = ("main", "probe", "usage_utils")
SAVED_ENV = {}
class OmniMock:
"""Управляемый мок OmniRoute: реальные HTTP-запросы на 127.0.0.1."""
def __init__(self):
self.state = {
# поведение /v1/chat/completions по умолчанию; переопределяется
# per-model: status, usage, no_usage, omit_choices, content, delay
"completions": {"status": 200, "usage": None, "content": "4"},
"models": {"data": [
{"id": "zai/test-model-a"},
{"id": "agy/test-model-b"},
]},
"global_delay": 0.0,
"per_model": {},
}
self.requests = []
self.lock = threading.Lock()
outer = self
class Handler(BaseHTTPRequestHandler):
def log_message(self, *a):
pass
def _send_json(self, status, payload):
data = json.dumps(payload).encode()
try:
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(data)))
self.end_headers()
self.wfile.write(data)
except (BrokenPipeError, ConnectionResetError):
pass # клиент отвалился по таймауту — это ожидаемо
def _read_body(self):
n = int(self.headers.get("Content-Length") or 0)
raw = self.rfile.read(n) if n else b""
try:
return json.loads(raw)
except ValueError:
return {}
def do_POST(self):
body = self._read_body()
with outer.lock:
outer.requests.append({"path": self.path, "body": body})
beh = dict(outer.state["completions"])
beh.update(outer.state["per_model"].get(body.get("model"), {}))
delay = beh.get("delay", 0.0) or outer.state["global_delay"]
if delay:
time.sleep(delay)
if beh.get("status", 200) != 200:
self._send_json(beh["status"], {"error": {"message": "mock error"}})
return
resp = {"id": "chatcmpl-mock", "object": "chat.completion",
"created": int(time.time()), "model": body.get("model", "?")}
if not beh.get("no_usage"):
resp["usage"] = beh.get("usage")
if not beh.get("omit_choices"):
resp["choices"] = [{"index": 0, "finish_reason": "stop",
"message": {"role": "assistant",
"content": beh.get("content", "4")}}]
self._send_json(200, resp)
def do_GET(self):
if self.path == "/v1/models":
with outer.lock:
payload = dict(outer.state["models"])
self._send_json(200, payload)
return
self._send_json(404, {"detail": "not found"})
self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
self.thread.start()
self.url = f"http://127.0.0.1:{self.server.server_address[1]}"
def set_completions(self, **kw):
with self.lock:
self.state["completions"].update(kw)
def set_per_model(self, model_id, beh):
with self.lock:
self.state["per_model"][model_id] = beh
def close(self):
self.server.shutdown()
self.server.server_close()
self.thread.join(timeout=5)
@pytest.fixture(scope="session")
def omni_server():
mock = OmniMock()
yield mock
mock.close()
@pytest.fixture(autouse=True)
def _reset_omni(omni_server):
"""Изолировать тесты: состояние мока не должно протекать между тестами."""
with omni_server.lock:
omni_server.state["per_model"] = {}
omni_server.state["global_delay"] = 0.0
omni_server.state["completions"] = {"status": 200, "usage": None, "content": "4"}
omni_server.requests.clear()
yield
def _set_env(db_path: str, omni_url: str):
env = {
"DB_PATH": db_path,
"OMNIROUTE_URL": omni_url,
"OMNIROUTE_API_KEY": "test-key",
"CHECK_TIMEOUT": "2",
"CHECK_CONCURRENCY": "4",
"ARENA_CACHE_TEXT": "/nonexistent/arena_text.json",
"ARENA_CACHE_AGENT": "/nonexistent/arena_agent.json",
}
for k, v in env.items():
SAVED_ENV[k] = os.environ.get(k)
os.environ[k] = v
def restore_env():
for k, v in SAVED_ENV.items():
if v is None:
os.environ.pop(k, None)
else:
os.environ[k] = v
@pytest.fixture
def fresh_main(tmp_path, omni_server):
"""Модуль main с env, привязанным к tmp-БД и мок-серверу (без запуска app).
init_db НЕ вызывается — тест миграции сам решает, что лежит в БД до него.
"""
db_path = str(tmp_path / "arena.db")
_set_env(db_path, omni_server.url)
for mod in MODULES:
sys.modules.pop(mod, None)
import main as m
try:
yield m
finally:
restore_env()
@pytest.fixture
def core(fresh_main):
"""Приложение с запущенным lifespan (init_db + sync_models) + модуль probe."""
import probe as probe_mod
with TestClient(fresh_main.app) as client:
yield SimpleNamespace(main=fresh_main, probe=probe_mod,
client=client, db_path=fresh_main.DB_PATH)
+69
View File
@@ -0,0 +1,69 @@
"""Реальный smoke: uvicorn + мок OmniRoute, проверка живых HTTP-запросов."""
import json
import os
import signal
import subprocess
import sys
import time
import urllib.request
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from conftest import OmniMock
REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
TMP = os.path.join(REPO, ".smoke")
os.makedirs(TMP, exist_ok=True)
mock = OmniMock()
mock.set_completions(usage={"prompt_tokens": 9, "completion_tokens": 4, "total_tokens": 13})
env = dict(os.environ)
env.update({
"DB_PATH": os.path.join(TMP, "arena.db"),
"OMNIROUTE_URL": mock.url,
"OMNIROUTE_API_KEY": "smoke",
"ARENA_CACHE_TEXT": "/nonexistent/a.json",
"ARENA_CACHE_AGENT": "/nonexistent/b.json",
})
proc = subprocess.Popen(
[os.path.join(REPO, ".venv/bin/uvicorn"), "main:app", "--host", "127.0.0.1", "--port", "18700"],
cwd=os.path.join(REPO, "app"), env=env,
stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
)
# pip install uvicorn into venv if missing handled outside
try:
base = "http://127.0.0.1:18700"
for _ in range(50):
try:
urllib.request.urlopen(f"{base}/api/providers", timeout=1)
break
except Exception:
time.sleep(0.2)
# 1. UI отдаётся (статика через __file__-путь)
ui = urllib.request.urlopen(f"{base}/", timeout=5).read().decode()
assert "Arena Dashboard" in ui and "s-tokens-24h" in ui and "loadTokenStats" in ui, "UI не содержит token-элементов"
# 2. Ручная проверка с usage
chk = json.loads(urllib.request.urlopen(f"{base}/api/check/zai/test-model-a", timeout=10).read())
assert chk["input_tokens"] == 9 and chk["total_tokens"] == 13, chk
# 3. Агрегат
ts = json.loads(urllib.request.urlopen(f"{base}/api/token_stats", timeout=5).read())
assert ts["window_24h"]["known"] >= 1 and ts["window_24h"]["total_tokens"] >= 13, ts
print("SMOKE_OK", json.dumps({
"check_usage": [chk["input_tokens"], chk["output_tokens"], chk["total_tokens"]],
"token_stats_24h": ts["window_24h"],
"providers": [(p["provider"], p["window_24h"]["total_tokens"]) for p in ts["providers"]],
}, ensure_ascii=False))
finally:
proc.send_signal(signal.SIGTERM)
try:
out, _ = proc.communicate(timeout=10)
print("--- uvicorn log tail ---")
print("\n".join(out.splitlines()[-8:]))
except subprocess.TimeoutExpired:
proc.kill()
mock.close()
+171
View File
@@ -0,0 +1,171 @@
"""Ручной путь (main.py /api/check) + API-агрегаты token usage."""
import json
import sqlite3
from datetime import datetime, timedelta, timezone
def _con(db_path):
con = sqlite3.connect(db_path)
con.row_factory = sqlite3.Row
return con
def _insert_check(db_path, model_id, at, ok=1, inp=None, out=None, tot=None, ujson=None):
con = _con(db_path)
con.execute(
"INSERT INTO checks (model_id, checked_at, ok, latency_ms, error, "
"input_tokens, output_tokens, total_tokens, usage_json) VALUES (?,?,?,?,?,?,?,?,?)",
(model_id, at, ok, 50, None, inp, out, tot, ujson),
)
con.commit()
con.close()
def _iso(**kw):
return (datetime.now(timezone.utc) + timedelta(**kw)).isoformat()
# ─── ручная проверка /api/check ───────────────────────────────────────────────
def test_api_check_captures_usage(core, omni_server):
omni_server.set_completions(usage={"prompt_tokens": 12, "completion_tokens": 3, "total_tokens": 15})
r = core.client.get("/api/check/zai/test-model-a")
assert r.status_code == 200
body = r.json()
assert body["ok"] is True
assert body["input_tokens"] == 12
assert body["output_tokens"] == 3
assert body["total_tokens"] == 15
row = _con(core.db_path).execute(
"SELECT * FROM checks WHERE model_id='zai/test-model-a'").fetchone()
assert row["input_tokens"] == 12 and row["total_tokens"] == 15
assert json.loads(row["usage_json"])["prompt_tokens"] == 12
def test_api_check_missing_usage_null(core, omni_server):
omni_server.set_completions(no_usage=True)
body = core.client.get("/api/check/zai/test-model-a").json()
assert body["ok"] is True
assert body["input_tokens"] is None and body["total_tokens"] is None
row = _con(core.db_path).execute(
"SELECT input_tokens, total_tokens, usage_json FROM checks WHERE model_id='zai/test-model-a'").fetchone()
assert row["input_tokens"] is None and row["total_tokens"] is None and row["usage_json"] is None
def test_api_check_failed_content_preserves_usage(core, omni_server):
omni_server.set_completions(usage={"prompt_tokens": 8, "completion_tokens": 1, "total_tokens": 9},
omit_choices=True)
body = core.client.get("/api/check/agy/test-model-b").json()
assert body["ok"] is False
assert body["input_tokens"] == 8 # usage пережил ошибку контента
row = _con(core.db_path).execute(
"SELECT ok, error, total_tokens FROM checks WHERE model_id='agy/test-model-b'").fetchone()
assert row["ok"] == 0 and row["error"]
assert row["total_tokens"] == 9
def test_api_check_no_response_stores_null(core, omni_server):
# недоступный OmniRoute: ответа нет — usage остаётся NULL (не нулями)
saved = core.main.OMNIROUTE_URL
core.main.OMNIROUTE_URL = "http://127.0.0.1:1" # connection refused, мгновенно
try:
body = core.client.get("/api/check/zai/test-model-a").json()
finally:
core.main.OMNIROUTE_URL = saved
assert body["ok"] is False
assert body["input_tokens"] is None and body["total_tokens"] is None
row = _con(core.db_path).execute(
"SELECT ok, error, total_tokens, usage_json FROM checks WHERE model_id='zai/test-model-a'").fetchone()
assert row["ok"] == 0 and row["error"]
assert row["total_tokens"] is None and row["usage_json"] is None
# ─── истории: usage в /api/checks и /api/models.history ──────────────────────
def test_api_checks_history_includes_usage(core, omni_server):
omni_server.set_completions(usage={"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7})
core.client.get("/api/check/zai/test-model-a")
omni_server.set_completions(no_usage=True)
core.client.get("/api/check/zai/test-model-a")
data = core.client.get("/api/checks/zai/test-model-a").json()
assert len(data["checks"]) == 2
with_usage = [c for c in data["checks"] if c["total_tokens"] == 7]
without = [c for c in data["checks"] if c["total_tokens"] is None]
assert len(with_usage) == 1 and with_usage[0]["input_tokens"] == 5
assert with_usage[0]["usage_json"]
assert len(without) == 1 and without[0]["input_tokens"] is None
# сегментная история тоже несёт токены (для тултипов UI)
models = core.client.get("/api/models").json()
model = next(m for m in models if m["id"] == "zai/test-model-a")
toks = [h["total_tokens"] for h in model["history"]]
assert set(toks) == {7, None}
# ─── /api/token_stats ─────────────────────────────────────────────────────────
def test_token_stats_empty_db(core):
ts = core.client.get("/api/token_stats").json()
assert ts["window_24h"]["checks"] == 0
assert ts["window_24h"]["coverage_pct"] is None
assert ts["all_time"]["checks"] == 0
assert ts["providers"] == []
def test_token_stats_windows_coverage_providers(core):
# zai: 2 свежих с usage + 1 старая (>24h) с usage
_insert_check(core.db_path, "zai/test-model-a", _iso(minutes=-10), inp=10, out=5, tot=15)
_insert_check(core.db_path, "zai/test-model-a", _iso(minutes=-20), inp=20, out=2, tot=22)
_insert_check(core.db_path, "zai/test-model-a", _iso(hours=-30), inp=100, out=50, tot=150)
# agy: 1 свежая без usage + 1 старая с usage
_insert_check(core.db_path, "agy/test-model-b", _iso(minutes=-5))
_insert_check(core.db_path, "agy/test-model-b", _iso(hours=-40), inp=7, out=3, tot=10)
# модель, которой нет в таблице models: провайдер выводится из model_id
_insert_check(core.db_path, "orphan/x", _iso(minutes=-1), inp=5, out=5, tot=10)
ts = core.client.get("/api/token_stats").json()
w = ts["window_24h"]
assert w["checks"] == 4 and w["known"] == 3 and w["unknown"] == 1
assert w["coverage_pct"] == 75
assert (w["input_tokens"], w["output_tokens"], w["total_tokens"]) == (35, 12, 47)
a = ts["all_time"]
assert a["checks"] == 6 and a["known"] == 5 and a["unknown"] == 1
assert (a["input_tokens"], a["output_tokens"], a["total_tokens"]) == (142, 65, 207)
provs = {p["provider"]: p for p in ts["providers"]}
assert set(provs) == {"zai", "agy", "orphan"}
assert provs["zai"]["window_24h"]["total_tokens"] == 37
assert provs["zai"]["all_time"]["total_tokens"] == 187
assert provs["agy"]["window_24h"]["known"] == 0
assert provs["agy"]["window_24h"]["unknown"] == 1
assert provs["agy"]["window_24h"]["total_tokens"] is None # сумма по NULL-группе
assert provs["agy"]["all_time"]["total_tokens"] == 10
assert provs["orphan"]["window_24h"]["total_tokens"] == 10 # fallback из model_id
def test_token_stats_counts_both_paths(core, omni_server):
"""probe.run() и ручной /api/check пишут usage, агрегат считает оба пути."""
omni_server.set_completions(usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15})
core.probe.run() # путь cron: zai 15 + agy 15
core.client.get("/api/check/zai/test-model-a") # путь ручной: +15
omni_server.set_completions(no_usage=True)
core.client.get("/api/check/agy/test-model-b") # путь ручной без usage: NULL
ts = core.client.get("/api/token_stats").json()
w = ts["window_24h"]
assert w["checks"] == 4 # 2 probe + 2 ручных
assert w["known"] == 3 # probe×2 + ручной zai
assert w["unknown"] == 1 # ручной agy без usage
assert w["total_tokens"] == 45 # 15 × 3 известных
assert w["input_tokens"] == 30 and w["output_tokens"] == 15
provs = {p["provider"]: p for p in ts["providers"]}
assert provs["zai"]["window_24h"]["total_tokens"] == 30 # probe + manual
assert provs["agy"]["window_24h"]["total_tokens"] == 15 # только probe
+77
View File
@@ -0,0 +1,77 @@
"""Миграция схемы checks: идемпотентность и сохранение старых NULL."""
import sqlite3
def _cols(db_path):
con = sqlite3.connect(db_path)
cols = [r[1] for r in con.execute("PRAGMA table_info(checks)").fetchall()]
con.close()
return cols
def _make_old_db(db_path, with_error_col=False):
"""БД со старой схемой checks (до учёта usage) и живыми данными."""
con = sqlite3.connect(db_path)
con.execute("""CREATE TABLE checks (
id INTEGER PRIMARY KEY AUTOINCREMENT,
model_id TEXT NOT NULL,
checked_at TEXT NOT NULL,
ok INTEGER NOT NULL,
latency_ms INTEGER)""")
if with_error_col:
con.execute("ALTER TABLE checks ADD COLUMN error TEXT")
con.execute(
"INSERT INTO checks (model_id, checked_at, ok, latency_ms, error) VALUES (?,?,?,?,?)"
if with_error_col else
"INSERT INTO checks (model_id, checked_at, ok, latency_ms) VALUES (?,?,?,?)",
("zai/old-model", "2026-01-01T00:00:00+00:00", 1, 123, None)
if with_error_col else ("zai/old-model", "2026-01-01T00:00:00+00:00", 1, 123),
)
con.commit()
con.close()
def test_migration_old_schema_preserves_null(fresh_main):
_make_old_db(fresh_main.DB_PATH, with_error_col=True)
fresh_main.init_db()
cols = _cols(fresh_main.DB_PATH)
for c in ("input_tokens", "output_tokens", "total_tokens", "usage_json"):
assert c in cols, f"колонка {c} не добавлена"
con = sqlite3.connect(fresh_main.DB_PATH)
con.row_factory = sqlite3.Row
row = con.execute("SELECT * FROM checks WHERE model_id='zai/old-model'").fetchone()
con.close()
assert row["ok"] == 1 and row["latency_ms"] == 123
# старые данные: неизвестное потребление = NULL, не нули
assert row["input_tokens"] is None
assert row["output_tokens"] is None
assert row["total_tokens"] is None
assert row["usage_json"] is None
def test_migration_from_ancient_schema_without_error(fresh_main):
_make_old_db(fresh_main.DB_PATH, with_error_col=False)
fresh_main.init_db() # не должен падать: добавит и error, и usage-колонки
cols = _cols(fresh_main.DB_PATH)
assert "error" in cols
assert "usage_json" in cols
def test_migration_idempotent(fresh_main):
_make_old_db(fresh_main.DB_PATH, with_error_col=True)
fresh_main.init_db()
fresh_main.init_db() # повторный запуск — no-op без ошибок
con = sqlite3.connect(fresh_main.DB_PATH)
n = con.execute("SELECT COUNT(*) FROM checks").fetchone()[0]
con.close()
assert n == 1 # данные не задвоились и не потерялись
def test_fresh_db_has_usage_columns(fresh_main):
fresh_main.init_db()
cols = _cols(fresh_main.DB_PATH)
for c in ("input_tokens", "output_tokens", "total_tokens", "usage_json"):
assert c in cols
+122
View File
@@ -0,0 +1,122 @@
"""Путь cron-пробы (probe.py): захват usage, NULL при неизвестном, сохранение при ошибках."""
import json
import sqlite3
def _rows(db_path):
con = sqlite3.connect(db_path)
con.row_factory = sqlite3.Row
rows = con.execute("SELECT * FROM checks ORDER BY id").fetchall()
out = [dict(r) for r in rows]
con.close()
return out
def test_probe_captures_usage(core, omni_server):
omni_server.set_completions(usage={"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
no_usage=False)
ok, lat, err, usage = core.probe.check_model("zai/test-model-a")
assert ok is True and err is None
assert usage["input_tokens"] == 11
assert usage["output_tokens"] == 4
assert usage["total_tokens"] == 15
assert json.loads(usage["usage_json"])["prompt_tokens"] == 11
def test_probe_missing_usage_stays_null(core, omni_server):
omni_server.set_completions(no_usage=True)
ok, lat, err, usage = core.probe.check_model("zai/test-model-a")
assert ok is True # контент есть — проверка успешна
# но usage неизвестен: NULL, не нули
assert usage == {"input_tokens": None, "output_tokens": None,
"total_tokens": None, "usage_json": None}
def test_probe_zero_usage_treated_as_unknown(core, omni_server):
omni_server.set_completions(usage={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
no_usage=False)
_, _, _, usage = core.probe.check_model("agy/test-model-b")
assert usage["total_tokens"] is None and usage["input_tokens"] is None
def test_probe_failed_content_preserves_usage(core, omni_server):
# ответ без choices: парсинг контента упадёт, но usage уже зафиксирован
omni_server.set_completions(usage={"prompt_tokens": 9, "completion_tokens": 2, "total_tokens": 11},
omit_choices=True)
ok, lat, err, usage = core.probe.check_model("zai/test-model-a")
assert ok is False
assert err and "KeyError" in err
assert usage["input_tokens"] == 9 # токены не потеряны
assert usage["total_tokens"] == 11
def test_probe_timeout_usage_null(core, omni_server):
omni_server.set_completions(delay=6) # CHECK_TIMEOUT=2 в тестовом env
ok, lat, err, usage = core.probe.check_model("zai/test-model-a")
assert ok is False
assert err and ("timed out" in err.lower() or "timeout" in err.lower())
assert usage["total_tokens"] is None
assert usage["usage_json"] is None
assert lat and lat >= 2000 # latency фиксируется всегда
omni_server.set_completions(delay=0)
def test_probe_http_error_usage_null(core, omni_server):
omni_server.set_completions(status=500)
ok, lat, err, usage = core.probe.check_model("agy/test-model-b")
assert ok is False and err
assert usage["input_tokens"] is None
def test_probe_run_end_to_end_both_models(core, omni_server):
"""Полный прогон run(): разные модели — разный исход по usage."""
omni_server.set_completions(usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15})
omni_server.set_per_model("agy/test-model-b", {"no_usage": True})
omni_server.set_per_model("zai/test-model-a",
{"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}})
core.probe.run()
rows = {r["model_id"]: r for r in _rows(core.db_path)}
assert set(rows) == {"zai/test-model-a", "agy/test-model-b"}
a = rows["zai/test-model-a"]
assert a["ok"] == 1 and (a["input_tokens"], a["output_tokens"], a["total_tokens"]) == (10, 5, 15)
assert json.loads(a["usage_json"])["total_tokens"] == 15
b = rows["agy/test-model-b"]
assert b["ok"] == 1 # без usage проверка всё равно успешна
assert b["input_tokens"] is None and b["total_tokens"] is None and b["usage_json"] is None
def test_probe_run_migrates_old_schema_standalone(tmp_path, omni_server):
"""Cron стартует раньше uvicorn: probe должен сам уметь мигрировать старую БД."""
from conftest import _set_env, restore_env, MODULES
import sys as _sys
db = str(tmp_path / "probe-only.db")
con = sqlite3.connect(db)
con.execute("CREATE TABLE checks (id INTEGER PRIMARY KEY AUTOINCREMENT, model_id TEXT NOT NULL, "
"checked_at TEXT NOT NULL, ok INTEGER NOT NULL, latency_ms INTEGER, error TEXT)")
con.execute("CREATE TABLE models (id TEXT PRIMARY KEY, provider TEXT NOT NULL, name TEXT NOT NULL, "
"raw_base TEXT NOT NULL, norm_base TEXT NOT NULL, clean_base TEXT NOT NULL, "
"text_model_match TEXT, agent_model_match TEXT, updated_at TEXT)")
con.execute("INSERT INTO models (id, provider, name, raw_base, norm_base, clean_base) "
"VALUES ('zai/test-model-a', 'zai', 'zai/test-model-a', 'test-model-a', 'testmodela', 'testmodela')")
con.commit()
con.close()
_set_env(db, omni_server.url)
try:
for m in MODULES:
_sys.modules.pop(m, None)
import probe as standalone_probe
omni_server.set_per_model("zai/test-model-a",
{"usage": {"input_tokens": 6, "output_tokens": 2}})
standalone_probe.run()
finally:
restore_env()
rows = _rows(db)
assert len(rows) == 1
assert rows[0]["input_tokens"] == 6
assert rows[0]["output_tokens"] == 2
assert rows[0]["total_tokens"] == 8 # вычислен из отчётных значений
+70
View File
@@ -0,0 +1,70 @@
"""extract_usage: разбор форматов usage разных провайдеров."""
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "app"))
from usage_utils import extract_usage
EMPTY = {"input_tokens": None, "output_tokens": None, "total_tokens": None, "usage_json": None}
def test_openai_shape():
u = extract_usage({"choices": [], "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}})
assert u["input_tokens"] == 10
assert u["output_tokens"] == 5
assert u["total_tokens"] == 15
assert u["usage_json"] == '{"completion_tokens": 5, "prompt_tokens": 10, "total_tokens": 15}'
def test_anthropic_shape_total_computed():
u = extract_usage({"usage": {"input_tokens": 7, "output_tokens": 3}})
assert u["input_tokens"] == 7
assert u["output_tokens"] == 3
assert u["total_tokens"] == 10 # точная сумма отчётных значений, не оценка
def test_camel_case_shape():
u = extract_usage({"usage": {"promptTokens": 4, "completionTokens": 2, "totalTokens": 6}})
assert (u["input_tokens"], u["output_tokens"], u["total_tokens"]) == (4, 2, 6)
def test_missing_usage_all_none():
assert extract_usage({"choices": []}) == EMPTY
assert extract_usage({"usage": None}) == EMPTY
assert extract_usage({}) == EMPTY
def test_non_dict_resp():
assert extract_usage(None) == EMPTY
assert extract_usage("error") == EMPTY
def test_zero_usage_treated_as_unknown():
# полностью нулевой usage = провайдер не отдал данные → NULL, не нули
u = extract_usage({"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}})
assert u == EMPTY
def test_string_encoded_usage():
u = extract_usage({"usage": '{"prompt_tokens": 8, "completion_tokens": 1, "total_tokens": 9}'})
assert (u["input_tokens"], u["output_tokens"], u["total_tokens"]) == (8, 1, 9)
def test_partial_total_only():
u = extract_usage({"usage": {"total_tokens": 42}})
assert u["input_tokens"] is None
assert u["output_tokens"] is None
assert u["total_tokens"] == 42
def test_garbage_values_ignored():
u = extract_usage({"usage": {"prompt_tokens": "много", "completion_tokens": True, "total_tokens": None}})
assert u == EMPTY
def test_usage_json_roundtrip():
raw = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "extra": {"x": 1}}
u = extract_usage({"usage": raw})
import json
assert json.loads(u["usage_json"]) == raw