208 lines
8.0 KiB
Python
Executable File
208 lines
8.0 KiB
Python
Executable File
import sqlite3
|
|
from datetime import datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
class StatsStore:
|
|
DAILY_RETENTION_DAYS = 7
|
|
|
|
def __init__(self, db_path: str | Path):
|
|
self.db_path = Path(db_path)
|
|
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._init_db()
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
conn = sqlite3.connect(self.db_path)
|
|
conn.row_factory = sqlite3.Row
|
|
return conn
|
|
|
|
def _init_db(self):
|
|
with self._connect() as conn:
|
|
conn.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS latest_summary (
|
|
id INTEGER PRIMARY KEY CHECK (id = 1),
|
|
updated_at INTEGER NOT NULL,
|
|
total_keys INTEGER NOT NULL,
|
|
active_keys INTEGER NOT NULL,
|
|
cooldown_keys INTEGER NOT NULL,
|
|
disabled_keys INTEGER NOT NULL,
|
|
total_concurrency INTEGER NOT NULL,
|
|
total_rpm INTEGER NOT NULL,
|
|
total_tpm INTEGER NOT NULL,
|
|
total_rpd INTEGER NOT NULL,
|
|
total_yesterday_rpd INTEGER NOT NULL
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS latest_key_stats (
|
|
key_id TEXT PRIMARY KEY,
|
|
updated_at INTEGER NOT NULL,
|
|
model_name TEXT NOT NULL,
|
|
provider TEXT NOT NULL,
|
|
status TEXT NOT NULL,
|
|
enabled INTEGER NOT NULL,
|
|
config_id TEXT,
|
|
owner TEXT,
|
|
current_concurrency INTEGER NOT NULL,
|
|
current_rpm INTEGER NOT NULL,
|
|
current_tpm INTEGER NOT NULL,
|
|
current_rpd INTEGER NOT NULL,
|
|
cooldown_remaining REAL NOT NULL,
|
|
rpm_limit INTEGER NOT NULL,
|
|
tpm_limit INTEGER NOT NULL,
|
|
rpd_limit INTEGER NOT NULL,
|
|
max_concurrency INTEGER NOT NULL
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS daily_usage (
|
|
day TEXT PRIMARY KEY,
|
|
label TEXT NOT NULL,
|
|
total_rpd INTEGER NOT NULL,
|
|
updated_at INTEGER NOT NULL
|
|
);
|
|
"""
|
|
)
|
|
|
|
def save_export(self, stats: list[dict[str, Any]], usage_summary: dict[str, Any], daily_usage: list[dict[str, Any]]):
|
|
now = int(datetime.now(timezone.utc).timestamp())
|
|
cutoff = (datetime.now(timezone.utc) - timedelta(days=self.DAILY_RETENTION_DAYS)).strftime("%Y%m%d")
|
|
|
|
with self._connect() as conn:
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO latest_summary (
|
|
id, updated_at, total_keys, active_keys, cooldown_keys, disabled_keys,
|
|
total_concurrency, total_rpm, total_tpm, total_rpd, total_yesterday_rpd
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
ON CONFLICT(id) DO UPDATE SET
|
|
updated_at = excluded.updated_at,
|
|
total_keys = excluded.total_keys,
|
|
active_keys = excluded.active_keys,
|
|
cooldown_keys = excluded.cooldown_keys,
|
|
disabled_keys = excluded.disabled_keys,
|
|
total_concurrency = excluded.total_concurrency,
|
|
total_rpm = excluded.total_rpm,
|
|
total_tpm = excluded.total_tpm,
|
|
total_rpd = excluded.total_rpd,
|
|
total_yesterday_rpd = excluded.total_yesterday_rpd
|
|
""",
|
|
(
|
|
1,
|
|
now,
|
|
usage_summary["total_keys"],
|
|
usage_summary["active_keys"],
|
|
usage_summary["cooldown_keys"],
|
|
usage_summary["disabled_keys"],
|
|
usage_summary["total_concurrency"],
|
|
usage_summary["total_rpm"],
|
|
usage_summary["total_tpm"],
|
|
usage_summary["total_rpd"],
|
|
usage_summary.get("total_yesterday_rpd", 0),
|
|
),
|
|
)
|
|
|
|
conn.execute("DELETE FROM latest_key_stats")
|
|
conn.executemany(
|
|
"""
|
|
INSERT INTO latest_key_stats (
|
|
key_id, updated_at, model_name, provider, status, enabled, config_id, owner,
|
|
current_concurrency, current_rpm, current_tpm, current_rpd, cooldown_remaining,
|
|
rpm_limit, tpm_limit, rpd_limit, max_concurrency
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
[
|
|
(
|
|
item["id"],
|
|
now,
|
|
item["model_name"],
|
|
item["provider"],
|
|
item["status"],
|
|
1 if item["enabled"] else 0,
|
|
item.get("config_id"),
|
|
item.get("owner"),
|
|
item["usage"]["current_concurrency"],
|
|
item["usage"]["current_rpm"],
|
|
item["usage"]["current_tpm"],
|
|
item["usage"]["current_rpd"],
|
|
item["cooldown_remaining"],
|
|
item["limits"]["rpm"],
|
|
item["limits"]["tpm"],
|
|
item["limits"]["rpd"],
|
|
item["limits"]["max_concurrency"],
|
|
)
|
|
for item in stats
|
|
],
|
|
)
|
|
|
|
conn.executemany(
|
|
"""
|
|
INSERT INTO daily_usage (day, label, total_rpd, updated_at)
|
|
VALUES (?, ?, ?, ?)
|
|
ON CONFLICT(day) DO UPDATE SET
|
|
label = excluded.label,
|
|
total_rpd = MAX(daily_usage.total_rpd, excluded.total_rpd),
|
|
updated_at = excluded.updated_at
|
|
""",
|
|
[
|
|
(
|
|
item["day"],
|
|
item["label"],
|
|
item["total_rpd"],
|
|
now,
|
|
)
|
|
for item in daily_usage
|
|
],
|
|
)
|
|
conn.execute("DELETE FROM daily_usage WHERE day < ?", (cutoff,))
|
|
|
|
def get_daily_usage(self, days: int = 7) -> list[dict[str, Any]]:
|
|
if days <= 0:
|
|
return []
|
|
|
|
with self._connect() as conn:
|
|
rows = conn.execute(
|
|
"""
|
|
SELECT day, label, total_rpd
|
|
FROM daily_usage
|
|
ORDER BY day DESC
|
|
LIMIT ?
|
|
""",
|
|
(days,),
|
|
).fetchall()
|
|
|
|
rows = list(reversed(rows))
|
|
return [
|
|
{
|
|
"day": row["day"],
|
|
"label": row["label"],
|
|
"total_rpd": row["total_rpd"],
|
|
}
|
|
for row in rows
|
|
]
|
|
|
|
def get_latest_summary(self) -> dict[str, Any] | None:
|
|
with self._connect() as conn:
|
|
row = conn.execute("SELECT * FROM latest_summary WHERE id = 1").fetchone()
|
|
|
|
if not row:
|
|
return None
|
|
|
|
return dict(row)
|
|
|
|
def get_latest_key_stats(self) -> list[dict[str, Any]]:
|
|
with self._connect() as conn:
|
|
rows = conn.execute(
|
|
"""
|
|
SELECT *
|
|
FROM latest_key_stats
|
|
ORDER BY provider, model_name, key_id
|
|
"""
|
|
).fetchall()
|
|
|
|
return [dict(row) for row in rows]
|
|
|
|
def delete_day(self, day: str):
|
|
with self._connect() as conn:
|
|
conn.execute("DELETE FROM daily_usage WHERE day = ?", (day,))
|