hermes_core/db.py

"""数据库层 — SQLite schema 创建与 CRUD 操作."""

import json
import sqlite3
from pathlib import Path
from typing import Optional
from hermes_core.types import (
    Record, Scope, TrainingRun, RecordState, TrainingStatus, HERMES_DATA_DIR,
    Dimension,
)


def _db_path(user_id: str) -> Path:
    return HERMES_DATA_DIR / "users" / user_id / "hermes.db"


def _ensure_dir(path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)


def _row_to_record(row: sqlite3.Row) -> Record:
    dims_raw = json.loads(row["dimensions"])
    dims = [Dimension(**d) for d in dims_raw]
    return Record(
        id=row["id"],
        user_id=row["user_id"],
        scope_id=row["scope_id"],
        scope_label=row["scope_label"],
        dimensions=dims,
        confidence=row["confidence"],
        occurrences=row["occurrences"],
        state=RecordState(row["state"]),
        source_conv=row["source_conv"] or "",
        source_agent=row["source_agent"] or "",
        created_at=row["created_at"],
        updated_at=row["updated_at"],
    )


def _row_to_scope(row: sqlite3.Row) -> Scope:
    centroid = json.loads(row["centroid"]) if row["centroid"] else None
    return Scope(
        id=row["id"],
        label=row["label"],
        centroid=centroid,
        record_count=row["record_count"],
        coherence=row["coherence"],
        status=row["status"],
        needs_training=bool(row["needs_training"]),
        created_at=row["created_at"],
        last_activity=row["last_activity"],
    )


def _row_to_training_run(row: sqlite3.Row) -> TrainingRun:
    progress = json.loads(row["progress"]) if row["progress"] else {}
    return TrainingRun(
        id=row["id"],
        scope_id=row["scope_id"],
        version=row["version"],
        status=TrainingStatus(row["status"]),
        progress=progress,
        content_hash=row["content_hash"] or "",
        checkpoint_path=row["checkpoint_path"] or "",
        started_at=row["started_at"] or "",
        finished_at=row["finished_at"] or "",
        error_msg=row["error_msg"] or "",
    )


def init_db(user_id: str | sqlite3.Connection) -> sqlite3.Connection:
    """初始化数据库，建表后返回连接。幂等：表已存在则跳过。

    Args:
        user_id: 用户 ID 字符串，或已有的 sqlite3.Connection
    """
    if isinstance(user_id, sqlite3.Connection):
        conn = user_id
    else:
        path = _db_path(user_id)
        _ensure_dir(path)
        conn = sqlite3.connect(str(path))
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")

    conn.executescript("""
        CREATE TABLE IF NOT EXISTS records (
            id          TEXT PRIMARY KEY,
            user_id     TEXT NOT NULL,
            scope_id    TEXT NOT NULL,
            scope_label TEXT NOT NULL,
            dimensions  TEXT NOT NULL DEFAULT '[]',
            confidence  REAL DEFAULT 0.5,
            occurrences INTEGER DEFAULT 1,
            state       TEXT DEFAULT 'active',
            source_conv TEXT DEFAULT '',
            source_agent TEXT DEFAULT '',
            created_at  TEXT NOT NULL DEFAULT (datetime('now')),
            updated_at  TEXT NOT NULL DEFAULT (datetime('now'))
        );

        CREATE TABLE IF NOT EXISTS scopes (
            id              TEXT PRIMARY KEY,
            label           TEXT NOT NULL,
            centroid        TEXT,
            record_count    INTEGER DEFAULT 0,
            coherence       REAL DEFAULT 1.0,
            status          TEXT DEFAULT 'active',
            needs_training  INTEGER DEFAULT 0,
            created_at      TEXT NOT NULL DEFAULT (datetime('now')),
            last_activity   TEXT NOT NULL DEFAULT (datetime('now'))
        );

        CREATE TABLE IF NOT EXISTS training_runs (
            id              TEXT PRIMARY KEY,
            scope_id        TEXT NOT NULL,
            version         INTEGER NOT NULL,
            status          TEXT NOT NULL DEFAULT 'pending',
            progress        TEXT DEFAULT '{}',
            content_hash    TEXT DEFAULT '',
            checkpoint_path TEXT DEFAULT '',
            started_at      TEXT DEFAULT '',
            finished_at     TEXT DEFAULT '',
            error_msg       TEXT DEFAULT ''
        );

        CREATE TABLE IF NOT EXISTS dimension_constraints (
            scope_id        TEXT PRIMARY KEY,
            user_id         TEXT NOT NULL,
            max_dims        INTEGER DEFAULT 1,
            last_conversation_id TEXT DEFAULT '',
            updated_at      TEXT NOT NULL DEFAULT (datetime('now'))
        );
    """)
    conn.commit()
    return conn


# ── Record CRUD ──

def insert_record(conn: sqlite3.Connection, record: Record) -> Record:
    """插入或 upsert 一条记录。"""
    dims_json = json.dumps([{"key": d.key, "value": d.value, "context": d.context}
                            for d in record.dimensions], ensure_ascii=False)
    conn.execute("""
        INSERT INTO records (id, user_id, scope_id, scope_label, dimensions,
                             confidence, occurrences, state, source_conv,
                             source_agent)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ON CONFLICT(id) DO UPDATE SET
            scope_label=excluded.scope_label,
            dimensions=excluded.dimensions,
            confidence=excluded.confidence,
            occurrences=excluded.occurrences,
            state=excluded.state,
            updated_at=datetime('now')
    """, (record.id, record.user_id, record.scope_id, record.scope_label,
          dims_json, record.confidence, record.occurrences, record.state.value,
          record.source_conv, record.source_agent))
    conn.commit()
    return record


def get_record(conn: sqlite3.Connection, record_id: str) -> Optional[Record]:
    row = conn.execute("SELECT * FROM records WHERE id=?", (record_id,)).fetchone()
    return _row_to_record(row) if row else None


def get_records_by_scope(conn: sqlite3.Connection, scope_id: str) -> list[Record]:
    rows = conn.execute(
        "SELECT * FROM records WHERE scope_id=? ORDER BY created_at", (scope_id,)
    ).fetchall()
    return [_row_to_record(r) for r in rows]


def get_active_records(conn: sqlite3.Connection, scope_id: str) -> list[Record]:
    rows = conn.execute(
        "SELECT * FROM records WHERE scope_id=? AND state='active' ORDER BY created_at",
        (scope_id,)
    ).fetchall()
    return [_row_to_record(r) for r in rows]


def update_record_state(conn: sqlite3.Connection, record_id: str, state: RecordState) -> None:
    conn.execute("UPDATE records SET state=?, updated_at=datetime('now') WHERE id=?",
                 (state.value, record_id))
    conn.commit()


def update_record_label(conn: sqlite3.Connection, record_id: str,
                        scope_id: str, scope_label: str) -> None:
    conn.execute("""
        UPDATE records SET scope_id=?, scope_label=?, updated_at=datetime('now')
        WHERE id=?
    """, (scope_id, scope_label, record_id))
    conn.commit()


# ── Scope CRUD ──

def upsert_scope(conn: sqlite3.Connection, scope: Scope) -> Scope:
    centroid_json = json.dumps(scope.centroid) if scope.centroid else None
    conn.execute("""
        INSERT INTO scopes (id, label, centroid, record_count, coherence, status,
                            needs_training, last_activity)
        VALUES (?, ?, ?, ?, ?, ?, ?, datetime('now'))
        ON CONFLICT(id) DO UPDATE SET
            label=excluded.label,
            centroid=COALESCE(excluded.centroid, scopes.centroid),
            record_count=excluded.record_count,
            coherence=excluded.coherence,
            status=excluded.status,
            needs_training=excluded.needs_training,
            last_activity=datetime('now')
    """, (scope.id, scope.label, centroid_json, scope.record_count,
          scope.coherence, scope.status, int(scope.needs_training)))
    conn.commit()
    return scope


def get_scope(conn: sqlite3.Connection, scope_id: str) -> Optional[Scope]:
    row = conn.execute("SELECT * FROM scopes WHERE id=?", (scope_id,)).fetchone()
    return _row_to_scope(row) if row else None


def get_active_scopes(conn: sqlite3.Connection) -> list[Scope]:
    rows = conn.execute(
        "SELECT * FROM scopes WHERE status='active' ORDER BY last_activity DESC"
    ).fetchall()
    return [_row_to_scope(r) for r in rows]


def get_scopes_needing_training(conn: sqlite3.Connection) -> list[Scope]:
    rows = conn.execute(
        "SELECT * FROM scopes WHERE needs_training=1 AND status='active' ORDER BY last_activity DESC"
    ).fetchall()
    return [_row_to_scope(r) for r in rows]


# ── TrainingRun CRUD ──

def insert_training_run(conn: sqlite3.Connection, run: TrainingRun) -> TrainingRun:
    progress_json = json.dumps(run.progress)
    conn.execute("""
        INSERT INTO training_runs (id, scope_id, version, status, progress,
                                   content_hash, checkpoint_path)
        VALUES (?, ?, ?, ?, ?, ?, ?)
    """, (run.id, run.scope_id, run.version, run.status.value,
          progress_json, run.content_hash, run.checkpoint_path))
    conn.commit()
    return run


def update_training_run(conn: sqlite3.Connection, run_id: str, **kwargs) -> None:
    allowed = {"status", "progress", "content_hash", "checkpoint_path",
               "started_at", "finished_at", "error_msg"}
    updates = {}
    for k, v in kwargs.items():
        if k in allowed:
            if k == "status" and isinstance(v, TrainingStatus):
                updates[k] = v.value
            elif k == "progress" and isinstance(v, dict):
                updates[k] = json.dumps(v)
            else:
                updates[k] = v
    if not updates:
        return
    set_clause = ", ".join(f"{k}=?" for k in updates)
    values = list(updates.values()) + [run_id]
    conn.execute(f"UPDATE training_runs SET {set_clause} WHERE id=?", values)
    conn.commit()


def get_latest_checkpoint(conn: sqlite3.Connection, scope_id: str) -> Optional[TrainingRun]:
    row = conn.execute("""
        SELECT * FROM training_runs
        WHERE scope_id=? AND status='done'
        ORDER BY version DESC LIMIT 1
    """, (scope_id,)).fetchone()
    return _row_to_training_run(row) if row else None


# ── Dimension Constraints ──

def get_max_dims(conn: sqlite3.Connection, scope_id: str) -> Optional[int]:
    row = conn.execute(
        "SELECT max_dims FROM dimension_constraints WHERE scope_id=?",
        (scope_id,)
    ).fetchone()
    return row["max_dims"] if row else None


def set_max_dims(conn: sqlite3.Connection, scope_id: str, user_id: str, max_dims: int) -> None:
    conn.execute("""
        INSERT INTO dimension_constraints (scope_id, user_id, max_dims)
        VALUES (?, ?, ?)
        ON CONFLICT(scope_id) DO UPDATE SET
            max_dims=excluded.max_dims,
            updated_at=datetime('now')
    """, (scope_id, user_id, max_dims))
    conn.commit()