Add src/selectel_ml_tui/storage.py

2026-09-02 16:43:39 +03:00
parent 4a6a188af1
commit 8ce0074784
1 changed files with 361 additions and 0 deletions
@@ -0,0 +1,361 @@
+"""Локальный кэш и история данных (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()