"""SQLite persistence for history + memory (SPEC-009). One database per deployment, stdlib sqlite3 only (D-010). Callers run the sync methods in a worker thread (`asyncio.to_thread`) so the event loop never blocks on disk (PER-05 / D6). """ import logging import os import pickle import sqlite3 from contextlib import closing from pathlib import Path from typing import Any, Dict, List, Optional SCHEMA_VERSION = 3 class PersistentStore: def __init__(self, db_path: Path) -> None: self.db_path = Path(db_path) self._init_db() def _connect(self) -> sqlite3.Connection: conn = sqlite3.connect(self.db_path) conn.execute("PRAGMA journal_mode=WAL") return conn def _init_db(self) -> None: # Forward-only migrations keyed on user_version (PER-06) self.db_path.parent.mkdir(parents=True, exist_ok=True) with closing(self._connect()) as conn, conn: version = conn.execute("PRAGMA user_version").fetchone()[0] if version < 1: conn.execute( "CREATE TABLE IF NOT EXISTS history (id INTEGER PRIMARY KEY, channel TEXT NOT NULL, role TEXT NOT NULL, content TEXT NOT NULL)" ) conn.execute("CREATE INDEX IF NOT EXISTS history_channel ON history (channel)") conn.execute("CREATE TABLE IF NOT EXISTS memory (channel TEXT PRIMARY KEY, content TEXT NOT NULL)") if version < 2: conn.execute( "CREATE TABLE IF NOT EXISTS usage (day TEXT NOT NULL, key TEXT NOT NULL, value REAL NOT NULL, PRIMARY KEY (day, key))" ) if version < 3: conn.execute( "CREATE TABLE IF NOT EXISTS user_facts (id INTEGER PRIMARY KEY, user TEXT NOT NULL, fact TEXT NOT NULL," " source TEXT NOT NULL DEFAULT 'self', updated_at TEXT NOT NULL DEFAULT (datetime('now')))" ) conn.execute( "CREATE TABLE IF NOT EXISTS pinned_facts (id INTEGER PRIMARY KEY, channel TEXT, fact TEXT NOT NULL," " created_at TEXT NOT NULL DEFAULT (datetime('now')))" ) conn.execute( "CREATE TABLE IF NOT EXISTS episodes (id INTEGER PRIMARY KEY, channel TEXT NOT NULL, summary TEXT NOT NULL," " created_at TEXT NOT NULL DEFAULT (datetime('now')))" ) conn.execute( "CREATE TABLE IF NOT EXISTS observations (id INTEGER PRIMARY KEY, channel TEXT NOT NULL, user TEXT NOT NULL," " kind TEXT NOT NULL, content TEXT NOT NULL, created_at TEXT NOT NULL DEFAULT (datetime('now')))" ) # Legacy single-string memories carry over as one episode each (MEM-08) conn.execute("INSERT INTO episodes (channel, summary) SELECT channel, content FROM memory") if version < SCHEMA_VERSION: conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") os.chmod(self.db_path, 0o600) # conversation data (PER-04) def load_history(self, channel: str) -> List[Dict[str, Any]]: with closing(self._connect()) as conn: rows = conn.execute("SELECT role, content FROM history WHERE channel = ? ORDER BY id", (channel,)).fetchall() return [{"role": role, "content": content} for role, content in rows] def save_history(self, channel: str, history: List[Dict[str, Any]]) -> None: # Full replace per save: histories are small (<= history-limit) # and trims must be reflected (D-011) with closing(self._connect()) as conn, conn: conn.execute("DELETE FROM history WHERE channel = ?", (channel,)) conn.executemany( "INSERT INTO history (channel, role, content) VALUES (?, ?, ?)", [(channel, str(entry.get("role", "user")), str(entry.get("content", ""))) for entry in history], ) def load_memory(self, channel: str) -> Optional[str]: with closing(self._connect()) as conn: row = conn.execute("SELECT content FROM memory WHERE channel = ?", (channel,)).fetchone() return row[0] if row else None def save_memory(self, channel: str, content: str) -> None: with closing(self._connect()) as conn, conn: conn.execute( "INSERT INTO memory (channel, content) VALUES (?, ?) ON CONFLICT(channel) DO UPDATE SET content = excluded.content", (channel, content), ) def usage_add(self, day: str, key: str, amount: float) -> None: with closing(self._connect()) as conn, conn: conn.execute( "INSERT INTO usage (day, key, value) VALUES (?, ?, ?) ON CONFLICT(day, key) DO UPDATE SET value = value + excluded.value", (day, key, amount), ) def usage_get(self, day: str, key: str) -> float: with closing(self._connect()) as conn: row = conn.execute("SELECT value FROM usage WHERE day = ? AND key = ?", (day, key)).fetchone() return float(row[0]) if row else 0.0 def delete_history_of_user(self, user: str) -> int: """Remove persisted rows carrying this user's messages (SAF-08).""" with closing(self._connect()) as conn, conn: cursor = conn.execute("DELETE FROM history WHERE content LIKE ?", (f'%"user": "{user}"%',)) return cursor.rowcount # --- structured memory (SPEC-002) --- def add_observation(self, channel: str, user: str, kind: str, content: str) -> None: with closing(self._connect()) as conn, conn: conn.execute("INSERT INTO observations (channel, user, kind, content) VALUES (?, ?, ?, ?)", (channel, user, kind, content)) def unconsumed_observations(self, channel: str) -> int: with closing(self._connect()) as conn: return int(conn.execute("SELECT COUNT(*) FROM observations WHERE channel = ?", (channel,)).fetchone()[0]) def peek_observations(self, channel: str) -> List[Dict[str, Any]]: with closing(self._connect()) as conn: rows = conn.execute("SELECT id, user, kind, content FROM observations WHERE channel = ? ORDER BY id", (channel,)).fetchall() return [{"id": row[0], "user": row[1], "kind": row[2], "content": row[3]} for row in rows] def consume_observations(self, channel: str, up_to_id: int) -> None: with closing(self._connect()) as conn, conn: conn.execute("DELETE FROM observations WHERE channel = ? AND id <= ?", (channel, up_to_id)) def add_user_fact(self, user: str, fact: str, source: str = "self") -> None: with closing(self._connect()) as conn, conn: conn.execute("INSERT INTO user_facts (user, fact, source) VALUES (?, ?, ?)", (user, fact, source)) def facts_for(self, users: List[str]) -> List[Dict[str, Any]]: if not users: return [] marks = ",".join("?" for _ in users) with closing(self._connect()) as conn: # marks is only "?" placeholders; user values stay parameterized rows = conn.execute( f"SELECT id, user, fact FROM user_facts WHERE user IN ({marks}) ORDER BY id", tuple(users) # nosec B608 ).fetchall() return [{"id": row[0], "user": row[1], "fact": row[2]} for row in rows] def delete_fact(self, fact_id: int) -> int: with closing(self._connect()) as conn, conn: return conn.execute("DELETE FROM user_facts WHERE id = ?", (fact_id,)).rowcount def purge_old_facts(self, retention_days: int) -> int: with closing(self._connect()) as conn, conn: cursor = conn.execute("DELETE FROM user_facts WHERE updated_at < datetime('now', ?)", (f"-{int(retention_days)} days",)) return cursor.rowcount def add_pinned(self, channel: Optional[str], fact: str) -> None: with closing(self._connect()) as conn, conn: conn.execute("INSERT INTO pinned_facts (channel, fact) VALUES (?, ?)", (channel, fact)) def pinned_for(self, channel: str) -> List[Dict[str, Any]]: with closing(self._connect()) as conn: rows = conn.execute( "SELECT id, channel, fact FROM pinned_facts WHERE channel IS NULL OR channel = ? ORDER BY id", (channel,) ).fetchall() return [{"id": row[0], "channel": row[1], "fact": row[2]} for row in rows] def pinned_all(self) -> List[Dict[str, Any]]: with closing(self._connect()) as conn: rows = conn.execute("SELECT id, channel, fact FROM pinned_facts ORDER BY id").fetchall() return [{"id": row[0], "channel": row[1], "fact": row[2]} for row in rows] def delete_pinned(self, pin_id: int) -> int: with closing(self._connect()) as conn, conn: return conn.execute("DELETE FROM pinned_facts WHERE id = ?", (pin_id,)).rowcount def add_episode(self, channel: str, summary: str) -> None: with closing(self._connect()) as conn, conn: conn.execute("INSERT INTO episodes (channel, summary) VALUES (?, ?)", (channel, summary)) def recent_episodes(self, channel: str, count: int) -> List[str]: with closing(self._connect()) as conn: rows = conn.execute("SELECT summary FROM episodes WHERE channel = ? ORDER BY id DESC LIMIT ?", (channel, count)).fetchall() return [row[0] for row in reversed(rows)] def trim_episodes(self, channel: str, keep: int) -> int: with closing(self._connect()) as conn, conn: cursor = conn.execute( "DELETE FROM episodes WHERE channel = ? AND id NOT IN (SELECT id FROM episodes WHERE channel = ? ORDER BY id DESC LIMIT ?)", (channel, channel, keep), ) return cursor.rowcount def purge_user_memory(self, user: str) -> int: """Facts, observations and episode traces of one user (MEM-09).""" removed = 0 with closing(self._connect()) as conn, conn: removed += conn.execute("DELETE FROM user_facts WHERE user = ?", (user,)).rowcount removed += conn.execute("DELETE FROM observations WHERE user = ?", (user,)).rowcount removed += conn.execute("DELETE FROM episodes WHERE summary LIKE ?", (f"%{user}%",)).rowcount return removed def migrate_pickles(self, channel: str, history_file: Path, memory_file: Path) -> None: """Import legacy pickles once; rename them *.migrated (PER-03).""" if self.load_history(channel) or self.load_memory(channel) is not None: return if history_file.exists(): try: with open(history_file, "rb") as fd: self.save_history(channel, pickle.load(fd)) history_file.rename(history_file.with_name(history_file.name + ".migrated")) logging.info(f"migrated legacy history pickle for {channel}") except Exception as err: logging.error(f"failed to migrate history pickle {history_file}: {err!r}") if memory_file.exists(): try: with open(memory_file, "rb") as fd: legacy_memory = str(pickle.load(fd)) self.save_memory(channel, legacy_memory) self.add_episode(channel, legacy_memory) # MEM-08 memory_file.rename(memory_file.with_name(memory_file.name + ".migrated")) logging.info(f"migrated legacy memory pickle for {channel}") except Exception as err: logging.error(f"failed to migrate memory pickle {memory_file}: {err!r}")