diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d64c88e --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +.venv/ +.pytest_cache/ +__pycache__/ +.smoke/ diff --git a/app/main.py b/app/main.py index 169e9d9..7f1095f 100644 --- a/app/main.py +++ b/app/main.py @@ -17,6 +17,8 @@ from fastapi import FastAPI, Query from fastapi.responses import HTMLResponse, JSONResponse from fastapi.staticfiles import StaticFiles +from usage_utils import extract_usage + 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", "") @@ -25,6 +27,10 @@ ARENA_CACHE_AGENT = os.environ.get("ARENA_CACHE_AGENT", "/data/arena_agent.json" # Сколько последних проверок показывать сегментами в колонке Uptime. 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 ─────────────────────────────────────────────────────────────── @@ -98,7 +104,11 @@ def init_db(): model_id TEXT NOT NULL, checked_at TEXT 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); @@ -112,6 +122,17 @@ def init_db(): if "error" not in cols: 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.close() @@ -360,7 +381,7 @@ def model_stats(db: sqlite3.Connection, model_id: str, hours: int = 24) -> dict: ).fetchone() 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", (model_id, cutoff) ).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 # История последних проверок (старые → новые) для сегментной витрины. - # Каждый элемент — детали проверки: ok, latency, время, текст ошибки. + # Каждый элемент — детали проверки: ok, latency, время, текст ошибки, токены. history = [ { "ok": bool(r["ok"]), "latency_ms": r["latency_ms"], "at": r["checked_at"], "error": r["error"], + "input_tokens": r["input_tokens"], + "output_tokens": r["output_tokens"], + "total_tokens": r["total_tokens"], } 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 ──────────────────────────────────────────────────────────────── @asynccontextmanager @@ -493,6 +604,8 @@ def api_check(model_id: str): ok = False latency = None error = None + usage = {"input_tokens": None, "output_tokens": None, + "total_tokens": None, "usage_json": None} try: body = json.dumps({ "model": model_id, @@ -509,6 +622,9 @@ def api_check(model_id: str): ) with urllib.request.urlopen(req, timeout=30) as r: resp = json.loads(r.read()) + # Usage фиксируем ДО валидации контента: даже если парсинг choices + # упадёт, реально потраченные токены сохраняются в checks. + usage = extract_usage(resp) content = resp["choices"][0]["message"]["content"] ok = bool(content.strip()) except Exception as e: @@ -520,8 +636,12 @@ def api_check(model_id: str): now = datetime.now(timezone.utc).isoformat() db = get_db() db.execute( - "INSERT INTO checks (model_id, checked_at, ok, latency_ms, error) VALUES (?,?,?,?,?)", - (model_id, now, int(ok), latency, error), + "INSERT INTO checks (model_id, checked_at, ok, latency_ms, 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() stats = model_stats(db, model_id) @@ -531,6 +651,9 @@ def api_check(model_id: str): "ok": ok, "latency_ms": latency, "checked_at": now, + "input_tokens": usage["input_tokens"], + "output_tokens": usage["output_tokens"], + "total_tokens": usage["total_tokens"], "history": stats["history"], "uptime_pct": stats["uptime_pct"], } @@ -542,8 +665,8 @@ def api_checks(model_id: str): limit = int(os.environ.get("CHECKS_LIMIT", "100")) db = get_db() rows = db.execute( - "SELECT ok, latency_ms, checked_at, error FROM checks " - "WHERE model_id=? ORDER BY checked_at DESC LIMIT ?", + "SELECT ok, latency_ms, checked_at, error, input_tokens, output_tokens, total_tokens, usage_json " + "FROM checks WHERE model_id=? ORDER BY checked_at DESC LIMIT ?", (model_id, limit), ).fetchall() db.close() @@ -555,17 +678,31 @@ def api_checks(model_id: str): "latency_ms": r["latency_ms"], "checked_at": r["checked_at"], "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 ], } +@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) 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() -app.mount("/static", StaticFiles(directory="/app/app/static"), name="static") +app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static") diff --git a/app/probe.py b/app/probe.py index b7f3ce2..4762360 100644 --- a/app/probe.py +++ b/app/probe.py @@ -12,6 +12,8 @@ import time import urllib.request from datetime import datetime, timezone +from usage_utils import extract_usage + 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", "") @@ -19,9 +21,17 @@ TIMEOUT = int(os.environ.get("CHECK_TIMEOUT", "30")) 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() error = None + usage = {"input_tokens": None, "output_tokens": None, + "total_tokens": None, "usage_json": None} try: body = json.dumps({ "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: resp = json.loads(r.read()) + # Usage — до любой валидации контента. + usage = extract_usage(resp) content = resp["choices"][0]["message"]["content"] ok = bool(content.strip()) 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) ok = False 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(): 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()] db.close() @@ -58,8 +89,8 @@ def run(): from concurrent.futures import ThreadPoolExecutor, as_completed def do(mid): - ok, lat, err = check_model(mid) - return mid, ok, lat, err + ok, lat, err, usage = check_model(mid) + return mid, ok, lat, err, usage db = sqlite3.connect(DB_PATH, check_same_thread=False) now = datetime.now(timezone.utc).isoformat() @@ -68,10 +99,14 @@ def run(): futs = {ex.submit(do, m): m for m in models} done = ok_count = 0 for f in as_completed(futs): - mid, ok, lat, err = f.result() + mid, ok, lat, err, usage = f.result() db.execute( - "INSERT INTO checks (model_id, checked_at, ok, latency_ms, error) VALUES (?,?,?,?,?)", - (mid, now, int(ok), lat, err), + "INSERT INTO checks (model_id, checked_at, ok, latency_ms, error, " + "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 if ok: diff --git a/app/static/index.html b/app/static/index.html index 48078a8..eaa0068 100644 --- a/app/static/index.html +++ b/app/static/index.html @@ -245,6 +245,8 @@
—
В рейтинге
—
Живых
—
Проверено сегодня
+
—
Токены 24ч
+
—
Покрытие usage
@@ -351,9 +353,10 @@ function showSegTip(el) { const tip = segTip(); const ok = data.ok; const lat = data.latency_ms !== null && data.latency_ms !== undefined ? `${data.latency_ms} мс` : '—'; + const toks = usageStr(data); tip.innerHTML = `
${ok ? '✓ Успех' : '✗ Провал'}
-
${fmtTime(data.at)} · ${lat}
+
${fmtTime(data.at)} · ${lat}${toks ? ` · ${escapeHtml(toks)}` : ''}
${data.error ? `
${escapeHtml(data.error)}
` : ''} `; tip.style.display = 'block'; @@ -413,6 +416,37 @@ function renderRow(m) { `; } +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) { document.getElementById('s-total').textContent = data.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')}`; applyFilters(); + loadTokenStats(); } async function loadProviders() { @@ -527,6 +562,7 @@ async function checkOne(modelId, btn) { if (data.uptime_pct !== undefined) allData[idx].uptime_pct = data.uptime_pct; } applyFilters(); + loadTokenStats(); } finally { btn.textContent = '▶ Тест'; btn.classList.remove('checking'); @@ -562,15 +598,17 @@ async function openCheckHistory(modelId, btn) { if (!checks.length) { body.innerHTML = '
Нет данных о проверках
'; } else { - body.innerHTML = checks.map(c => ` + body.innerHTML = checks.map(c => { + const toks = usageStr(c); + return `
${c.ok ? '✓ Успех' : '✗ Провал'} - ${fmtTime(c.checked_at)} · ${c.latency_ms !== null && c.latency_ms !== undefined ? c.latency_ms + ' мс' : '—'} + ${fmtTime(c.checked_at)} · ${c.latency_ms !== null && c.latency_ms !== undefined ? c.latency_ms + ' мс' : '—'}${toks ? ` · ${escapeHtml(toks)} т.` : ''}
${c.error ? `
${escapeHtml(c.error)}
` : ''} -
- `).join(''); +
`; + }).join(''); } modal.classList.add('open'); diff --git a/app/usage_utils.py b/app/usage_utils.py new file mode 100644 index 0000000..a99f339 --- /dev/null +++ b/app/usage_utils.py @@ -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, + } diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..8c77ed3 --- /dev/null +++ b/tests/conftest.py @@ -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) diff --git a/tests/smoke_e2e.py b/tests/smoke_e2e.py new file mode 100644 index 0000000..35f4b93 --- /dev/null +++ b/tests/smoke_e2e.py @@ -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() diff --git a/tests/test_api.py b/tests/test_api.py new file mode 100644 index 0000000..ae69481 --- /dev/null +++ b/tests/test_api.py @@ -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 diff --git a/tests/test_migration.py b/tests/test_migration.py new file mode 100644 index 0000000..d2d67bd --- /dev/null +++ b/tests/test_migration.py @@ -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 diff --git a/tests/test_probe.py b/tests/test_probe.py new file mode 100644 index 0000000..49a4dec --- /dev/null +++ b/tests/test_probe.py @@ -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 # вычислен из отчётных значений diff --git a/tests/test_usage_utils.py b/tests/test_usage_utils.py new file mode 100644 index 0000000..dc722a4 --- /dev/null +++ b/tests/test_usage_utils.py @@ -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