362 lines
13 KiB
Python
362 lines
13 KiB
Python
"""Локальный кэш и история данных (SQLite).
|
|
|
|
База данных в каталоге данных платформы (Linux: ~/.local/share,
|
|
macOS: ~/Library/Application Support). Хранит снимки потребления и метрик
|
|
при каждом успешном обновлении, а также историю по периодам для графиков.
|
|
При недоступности API TUI может показать последние данные из кэша.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import csv
|
|
import datetime as dt
|
|
import json
|
|
import sqlite3
|
|
from decimal import Decimal
|
|
from pathlib import Path
|
|
from typing import Any, cast
|
|
|
|
from selectel_ml_tui.config import data_dir
|
|
from selectel_ml_tui.models import (
|
|
Consumption,
|
|
HistoryPoint,
|
|
Metrics,
|
|
ModelMetrics,
|
|
TokenUsage,
|
|
parse_iso_dt,
|
|
)
|
|
|
|
SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS consumption_snapshots (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
created_at TEXT NOT NULL,
|
|
period TEXT NOT NULL,
|
|
amount_rub TEXT NOT NULL,
|
|
input_tokens INTEGER NOT NULL DEFAULT 0,
|
|
output_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cache_read_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cache_write_tokens INTEGER NOT NULL DEFAULT 0
|
|
);
|
|
CREATE TABLE IF NOT EXISTS metrics_snapshots (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
created_at TEXT NOT NULL,
|
|
availability_pct REAL NOT NULL DEFAULT 100,
|
|
avg_ttft_ms REAL NOT NULL DEFAULT 0,
|
|
input_tokens INTEGER NOT NULL DEFAULT 0,
|
|
output_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cache_read_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cache_write_tokens INTEGER NOT NULL DEFAULT 0
|
|
);
|
|
CREATE TABLE IF NOT EXISTS model_metrics (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
created_at TEXT NOT NULL,
|
|
model_id TEXT NOT NULL,
|
|
model_name TEXT NOT NULL DEFAULT '',
|
|
availability_pct REAL NOT NULL DEFAULT 100,
|
|
avg_ttft_ms REAL NOT NULL DEFAULT 0,
|
|
input_tokens INTEGER NOT NULL DEFAULT 0,
|
|
output_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cache_read_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cache_write_tokens INTEGER NOT NULL DEFAULT 0
|
|
);
|
|
CREATE TABLE IF NOT EXISTS history (
|
|
timestamp TEXT PRIMARY KEY,
|
|
amount_rub TEXT NOT NULL,
|
|
input_tokens INTEGER NOT NULL DEFAULT 0,
|
|
output_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cache_read_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cache_write_tokens INTEGER NOT NULL DEFAULT 0
|
|
);
|
|
"""
|
|
|
|
|
|
def default_db_path() -> Path:
|
|
"""Путь к базе данных по умолчанию (каталог данных платформы)."""
|
|
return data_dir() / "cache.db"
|
|
|
|
|
|
def _utc_now() -> str:
|
|
return dt.datetime.now(dt.timezone.utc).isoformat()
|
|
|
|
|
|
def _history_columns(point: HistoryPoint) -> tuple[Any, ...]:
|
|
return (
|
|
point.timestamp.isoformat(),
|
|
str(point.amount_rub),
|
|
point.tokens.input_tokens,
|
|
point.tokens.output_tokens,
|
|
point.tokens.cache_read_tokens,
|
|
point.tokens.cache_write_tokens,
|
|
)
|
|
|
|
|
|
def _history_from_row(row: sqlite3.Row) -> HistoryPoint:
|
|
return HistoryPoint(
|
|
timestamp=parse_iso_dt(cast(str, row["timestamp"])),
|
|
amount_rub=Decimal(str(row["amount_rub"])),
|
|
tokens=TokenUsage(
|
|
input_tokens=cast(int, row["input_tokens"]),
|
|
output_tokens=cast(int, row["output_tokens"]),
|
|
cache_read_tokens=cast(int, row["cache_read_tokens"]),
|
|
cache_write_tokens=cast(int, row["cache_write_tokens"]),
|
|
),
|
|
)
|
|
|
|
|
|
class Storage:
|
|
"""SQLite-кэш: снимки потребления/метрик и история для графиков."""
|
|
|
|
def __init__(self, path: Path | None = None) -> None:
|
|
self._path = path if path is not None else default_db_path()
|
|
self._path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._conn = sqlite3.connect(self._path)
|
|
self._conn.row_factory = sqlite3.Row
|
|
self._init_schema()
|
|
|
|
@property
|
|
def path(self) -> Path:
|
|
return self._path
|
|
|
|
def _init_schema(self) -> None:
|
|
self._conn.executescript(SCHEMA)
|
|
self._conn.commit()
|
|
|
|
def save_consumption(self, item: Consumption) -> None:
|
|
"""Сохранить снимок потребления."""
|
|
self._conn.execute(
|
|
"""
|
|
INSERT INTO consumption_snapshots
|
|
(created_at, period, amount_rub, input_tokens, output_tokens,
|
|
cache_read_tokens, cache_write_tokens)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
_utc_now(),
|
|
item.period.isoformat(),
|
|
str(item.amount_rub),
|
|
item.input_tokens,
|
|
item.output_tokens,
|
|
item.cache_read_tokens,
|
|
item.cache_write_tokens,
|
|
),
|
|
)
|
|
self._conn.commit()
|
|
|
|
def save_metrics(self, metrics: Metrics) -> None:
|
|
"""Сохранить снимок общих метрик."""
|
|
self._conn.execute(
|
|
"""
|
|
INSERT INTO metrics_snapshots
|
|
(created_at, availability_pct, avg_ttft_ms, input_tokens,
|
|
output_tokens, cache_read_tokens, cache_write_tokens)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
_utc_now(),
|
|
metrics.availability_pct,
|
|
metrics.avg_ttft_ms,
|
|
metrics.input_tokens,
|
|
metrics.output_tokens,
|
|
metrics.cache_read_tokens,
|
|
metrics.cache_write_tokens,
|
|
),
|
|
)
|
|
self._conn.commit()
|
|
|
|
def save_model_metrics(self, models: list[ModelMetrics]) -> None:
|
|
"""Сохранить снимок метрик по моделям (один пакет = один created_at)."""
|
|
batch = _utc_now()
|
|
for model in models:
|
|
self._conn.execute(
|
|
"""
|
|
INSERT INTO model_metrics
|
|
(created_at, model_id, model_name, availability_pct,
|
|
avg_ttft_ms, input_tokens, output_tokens,
|
|
cache_read_tokens, cache_write_tokens)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
batch,
|
|
model.model_id,
|
|
model.model_name,
|
|
model.availability_pct,
|
|
model.avg_ttft_ms,
|
|
model.input_tokens,
|
|
model.output_tokens,
|
|
model.cache_read_tokens,
|
|
model.cache_write_tokens,
|
|
),
|
|
)
|
|
self._conn.commit()
|
|
|
|
def save_history(self, point: HistoryPoint) -> None:
|
|
"""Сохранить/обновить точку истории (upsert по timestamp)."""
|
|
self._conn.execute(
|
|
"""
|
|
INSERT OR REPLACE INTO history
|
|
(timestamp, amount_rub, input_tokens, output_tokens,
|
|
cache_read_tokens, cache_write_tokens)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
""",
|
|
_history_columns(point),
|
|
)
|
|
self._conn.commit()
|
|
|
|
def save_snapshot(
|
|
self,
|
|
*,
|
|
consumption: Consumption | None = None,
|
|
metrics: Metrics | None = None,
|
|
models: list[ModelMetrics] | None = None,
|
|
history_point: HistoryPoint | None = None,
|
|
) -> None:
|
|
"""Сохранить полный снимок данных после успешного обновления."""
|
|
if consumption is not None:
|
|
self.save_consumption(consumption)
|
|
if metrics is not None:
|
|
self.save_metrics(metrics)
|
|
if models:
|
|
self.save_model_metrics(models)
|
|
if history_point is not None:
|
|
self.save_history(history_point)
|
|
|
|
def get_latest_consumption(self) -> Consumption | None:
|
|
"""Последний снимок потребления (для офлайн-режима)."""
|
|
row = self._conn.execute(
|
|
"SELECT * FROM consumption_snapshots ORDER BY id DESC LIMIT 1"
|
|
).fetchone()
|
|
if row is None:
|
|
return None
|
|
return Consumption(
|
|
period=parse_iso_dt(cast(str, row["period"])),
|
|
amount_rub=Decimal(str(row["amount_rub"])),
|
|
input_tokens=cast(int, row["input_tokens"]),
|
|
output_tokens=cast(int, row["output_tokens"]),
|
|
cache_read_tokens=cast(int, row["cache_read_tokens"]),
|
|
cache_write_tokens=cast(int, row["cache_write_tokens"]),
|
|
)
|
|
|
|
def get_latest_metrics(self) -> Metrics | None:
|
|
"""Последний снимок общих метрик."""
|
|
row = self._conn.execute(
|
|
"SELECT * FROM metrics_snapshots ORDER BY id DESC LIMIT 1"
|
|
).fetchone()
|
|
if row is None:
|
|
return None
|
|
return Metrics(
|
|
availability_pct=cast(float, row["availability_pct"]),
|
|
avg_ttft_ms=cast(float, row["avg_ttft_ms"]),
|
|
input_tokens=cast(int, row["input_tokens"]),
|
|
output_tokens=cast(int, row["output_tokens"]),
|
|
cache_read_tokens=cast(int, row["cache_read_tokens"]),
|
|
cache_write_tokens=cast(int, row["cache_write_tokens"]),
|
|
)
|
|
|
|
def get_latest_model_metrics(self) -> list[ModelMetrics]:
|
|
"""Снимок метрик по моделям за последнее обновление."""
|
|
rows = self._conn.execute(
|
|
"""
|
|
SELECT * FROM model_metrics
|
|
WHERE created_at = (SELECT MAX(created_at) FROM model_metrics)
|
|
ORDER BY model_id
|
|
"""
|
|
).fetchall()
|
|
models: list[ModelMetrics] = []
|
|
for row in rows:
|
|
models.append(
|
|
ModelMetrics(
|
|
model_id=cast(str, row["model_id"]),
|
|
model_name=cast(str, row["model_name"]),
|
|
availability_pct=cast(float, row["availability_pct"]),
|
|
avg_ttft_ms=cast(float, row["avg_ttft_ms"]),
|
|
input_tokens=cast(int, row["input_tokens"]),
|
|
output_tokens=cast(int, row["output_tokens"]),
|
|
cache_read_tokens=cast(int, row["cache_read_tokens"]),
|
|
cache_write_tokens=cast(int, row["cache_write_tokens"]),
|
|
)
|
|
)
|
|
return models
|
|
|
|
def get_history(self, days: int = 7) -> list[HistoryPoint]:
|
|
"""Серия точек истории за последние ``days`` дней для графика."""
|
|
cutoff = (
|
|
dt.datetime.now(dt.timezone.utc) - dt.timedelta(days=days)
|
|
).isoformat()
|
|
rows = self._conn.execute(
|
|
"SELECT * FROM history WHERE timestamp >= ? ORDER BY timestamp ASC",
|
|
(cutoff,),
|
|
).fetchall()
|
|
return [_history_from_row(row) for row in rows]
|
|
|
|
def cleanup(self, max_age_days: int = 180) -> int:
|
|
"""Удалить записи старше ``max_age_days`` дней. Возвращает число удалённых."""
|
|
cutoff = (
|
|
dt.datetime.now(dt.timezone.utc) - dt.timedelta(days=max_age_days)
|
|
).isoformat()
|
|
deleted = 0
|
|
for table, column in (
|
|
("consumption_snapshots", "created_at"),
|
|
("metrics_snapshots", "created_at"),
|
|
("model_metrics", "created_at"),
|
|
("history", "timestamp"),
|
|
):
|
|
cursor = self._conn.execute(
|
|
f"DELETE FROM {table} WHERE {column} < ?", (cutoff,)
|
|
)
|
|
deleted += cursor.rowcount
|
|
self._conn.commit()
|
|
return deleted
|
|
|
|
def export_json(self, path: Path) -> None:
|
|
"""Экспортировать всю историю в JSON."""
|
|
points = self.get_history(days=3650)
|
|
data = [
|
|
{
|
|
"timestamp": point.timestamp.isoformat(),
|
|
"amount_rub": str(point.amount_rub),
|
|
"input_tokens": point.tokens.input_tokens,
|
|
"output_tokens": point.tokens.output_tokens,
|
|
"cache_read_tokens": point.tokens.cache_read_tokens,
|
|
"cache_write_tokens": point.tokens.cache_write_tokens,
|
|
}
|
|
for point in points
|
|
]
|
|
path.write_text(
|
|
json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8"
|
|
)
|
|
|
|
def export_csv(self, path: Path) -> None:
|
|
"""Экспортировать всю историю в CSV."""
|
|
points = self.get_history(days=3650)
|
|
with path.open("w", newline="", encoding="utf-8") as fh:
|
|
writer = csv.writer(fh)
|
|
writer.writerow(
|
|
[
|
|
"timestamp",
|
|
"amount_rub",
|
|
"input_tokens",
|
|
"output_tokens",
|
|
"cache_read_tokens",
|
|
"cache_write_tokens",
|
|
]
|
|
)
|
|
for point in points:
|
|
writer.writerow(_history_columns(point))
|
|
|
|
def clear(self) -> int:
|
|
"""Полностью очистить кэш. Возвращает число удалённых записей."""
|
|
deleted = 0
|
|
for table in (
|
|
"consumption_snapshots",
|
|
"metrics_snapshots",
|
|
"model_metrics",
|
|
"history",
|
|
):
|
|
cursor = self._conn.execute(f"DELETE FROM {table}")
|
|
deleted += cursor.rowcount
|
|
self._conn.commit()
|
|
return deleted
|
|
|
|
def close(self) -> None:
|
|
self._conn.close()
|