Files
ai_coder 5c8c13891f
Build and deploy / build-deploy (push) Has been cancelled
feat: учёт token usage в checks с агрегатами и UI
- миграция 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)
2026-10-25 17:27:38 +00:00

172 lines
8.1 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Ручной путь (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