"""数据库层 — 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()