feat: учёт token usage в checks с агрегатами и UI
Build and deploy / build-deploy (push) Has been cancelled
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
@@ -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)
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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 # вычислен из отчётных значений
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user