omnisvg-lite/history.py

"""SQLite-история генераций.

Один пользователь, без авторизации, одна таблица. WAL для устойчивости при
одновременной записи из UI и, например, фонового процесса.

Файл БД по умолчанию: `~/.omnisvg_lite/history.sqlite`. Можно переопределить
через `OMNISVG_DB_PATH` (env) или явно в конструкторе.

Схема расширена по сравнению с design.md: добавлены `n_requested`, `n_returned`,
`previews` (JSON list[str]), `error_reason`, `best_index`, `temperature` —
чтобы UI мог показать полную картину без дополнительных запросов.
"""

from __future__ import annotations

import json
import logging
import os
import sqlite3
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Iterator


log = logging.getLogger(__name__)


# Дефолтный путь к БД — в домашнем каталоге пользователя. Это не абсолютный
# путь к проекту: пользователь может переопределить через env.
DEFAULT_DB_DIR = Path.home() / ".omnisvg_lite"
DEFAULT_DB_PATH = DEFAULT_DB_DIR / "history.sqlite"


def _resolve_db_path(db_path: str | Path | None) -> Path:
    if db_path is not None:
        return Path(db_path)
    env = os.environ.get("OMNISVG_DB_PATH")
    if env:
        return Path(env)
    return DEFAULT_DB_PATH


_SCHEMA = """
CREATE TABLE IF NOT EXISTS generations (
    id                INTEGER PRIMARY KEY AUTOINCREMENT,
    created_at        REAL    NOT NULL,
    prompt            TEXT    NOT NULL,
    mode              TEXT    NOT NULL,
    model             TEXT    NOT NULL,
    n_requested       INTEGER NOT NULL,
    n_returned        INTEGER NOT NULL,
    temperature       REAL    NOT NULL,
    status            TEXT    NOT NULL,
    error_reason      TEXT,
    raw_outputs       TEXT    NOT NULL,
    validated_outputs TEXT    NOT NULL,
    previews          TEXT    NOT NULL,
    best_index        INTEGER
);
CREATE INDEX IF NOT EXISTS idx_generations_created_at
    ON generations(created_at DESC);
"""


@dataclass
class Record:
    """Запись о генерации.

    Attributes:
        prompt: пользовательский промпт.
        mode: "icon" | "illustration".
        model: имя модели.
        n_requested: сколько кандидатов запросили.
        n_returned: сколько LM Studio реально вернула.
        temperature: использованная температура.
        status: "ok" | "partial" | "failed".
        error_reason: строка с ошибкой или None.
        raw_outputs: list[str] — N сырых ответов модели.
        validated_outputs: list[str] — прошедшие валидатор SVG.
        previews: list[str] — пути к PNG-превью (rel или abs).
        best_index: индекс лучшего кандидата в validated_outputs (или None).
        created_at: time.time() (если None — заполняется в `add`).
    """

    prompt: str
    mode: str
    model: str
    n_requested: int
    n_returned: int
    temperature: float
    status: str
    error_reason: str | None
    raw_outputs: list[str] = field(default_factory=list)
    validated_outputs: list[str] = field(default_factory=list)
    previews: list[str] = field(default_factory=list)
    best_index: int | None = None
    created_at: float | None = None


class History:
    """Контекст-менеджер для работы с историей генераций.

    Используется так:
        with History() as h:
            row_id = h.add(Record(...))
            for r in h.list_recent(20):
                ...

    На каждый enter/exit открывается/закрывается соединение. Для долгоживущих
    демонов Gradio это нормально: событий мало.
    """

    def __init__(self, db_path: str | Path | None = None):
        self.db_path = _resolve_db_path(db_path)
        self._conn: sqlite3.Connection | None = None

    # -- context manager ------------------------------------------------

    def __enter__(self) -> "History":
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._conn = sqlite3.connect(str(self.db_path))
        self._conn.row_factory = sqlite3.Row
        # WAL: пишущие транзакции не блокируют читателей, читатели — писателей.
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._conn.execute("PRAGMA synchronous=NORMAL")
        self._conn.executescript(_SCHEMA)
        self._conn.commit()
        log.info("history: БД открыта по %s", self.db_path)
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        if self._conn is not None:
            try:
                if exc_type is None:
                    self._conn.commit()
                else:
                    self._conn.rollback()
            finally:
                self._conn.close()
                self._conn = None

    @property
    def conn(self) -> sqlite3.Connection:
        if self._conn is None:
            raise RuntimeError("History используется вне контекст-менеджера")
        return self._conn

    # -- CRUD -----------------------------------------------------------

    def add(self, record: Record) -> int:
        """Добавляет запись и возвращает её id."""
        created_at = record.created_at if record.created_at is not None else time.time()
        cur = self.conn.execute(
            """
            INSERT INTO generations
              (created_at, prompt, mode, model, n_requested, n_returned,
               temperature, status, error_reason,
               raw_outputs, validated_outputs, previews, best_index)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                created_at,
                record.prompt,
                record.mode,
                record.model,
                record.n_requested,
                record.n_returned,
                record.temperature,
                record.status,
                record.error_reason,
                json.dumps(record.raw_outputs, ensure_ascii=False),
                json.dumps(record.validated_outputs, ensure_ascii=False),
                json.dumps(record.previews, ensure_ascii=False),
                record.best_index,
            ),
        )
        self.conn.commit()
        return int(cur.lastrowid)

    def list_recent(self, limit: int = 50) -> list[dict[str, Any]]:
        """Возвращает последние `limit` записей, новые сверху."""
        cur = self.conn.execute(
            """
            SELECT id, created_at, prompt, mode, model, n_requested, n_returned,
                   temperature, status, error_reason,
                   raw_outputs, validated_outputs, previews, best_index
              FROM generations
             ORDER BY created_at DESC, id DESC
             LIMIT ?
            """,
            (limit,),
        )
        return [_row_to_dict(row) for row in cur.fetchall()]

    def get(self, record_id: int) -> dict[str, Any] | None:
        """Возвращает одну запись по id или None."""
        cur = self.conn.execute(
            """
            SELECT id, created_at, prompt, mode, model, n_requested, n_returned,
                   temperature, status, error_reason,
                   raw_outputs, validated_outputs, previews, best_index
              FROM generations
             WHERE id = ?
            """,
            (record_id,),
        )
        row = cur.fetchone()
        if row is None:
            return None
        return _row_to_dict(row)

    def count(self) -> int:
        cur = self.conn.execute("SELECT COUNT(*) AS c FROM generations")
        return int(cur.fetchone()["c"])


def _row_to_dict(row: sqlite3.Row) -> dict[str, Any]:
    d = dict(row)
    # JSON-поля превращаем в Python-структуры.
    for key in ("raw_outputs", "validated_outputs", "previews"):
        try:
            d[key] = json.loads(d[key]) if d.get(key) else []
        except (TypeError, ValueError):
            log.warning("history: не удалось разобрать JSON в поле %s", key)
            d[key] = []
    return d


__all__ = ["DEFAULT_DB_PATH", "History", "Record"]