Oleksandr Kozachuk
113 lines
5.3 KiB
Python
113 lines
5.3 KiB
Python
"""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 = 2
|
|
|
|
|
|
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 < 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
|
|
|
|
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:
|
|
self.save_memory(channel, str(pickle.load(fd)))
|
|
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}")
|