"""SQLite 记忆存储。 几个刻意的选择： * **列名用 mem_key / mem_value**，不用 key / value——后者在 SQL 里 容易和关键字撞，每个语句都得加引号，早晚漏一处。 * **WAL 模式**：读写不互相阻塞，将来多进程访问同一份库也不容易锁死。 * **外键级联删除**：删用户时记忆/会话/轮次一起走。方案 §8 要求 "删除同时覆盖会话上下文、摘要和未来索引"，靠数据库约束比靠代码记得住。 * **所有 SQL 操作走 asyncio.to_thread**。本地 SQLite 其实微秒级， 但接口承诺了 async，就得真的不阻塞事件循环——否则将来换成远程 实现时，调用方那些"同步假设"会全部暴露成 bug。 * **每条查询都带 owner_user_id**。没有"按 memory_id 单独取"的口子： 猜到 id 就能读别人的记忆，是这类系统最容易犯的错。 """ from __future__ import annotations import asyncio import sqlite3 import threading import time from pathlib import Path from typing import Sequence from .base import ConfirmStatus, Memory, MemoryKind, MemoryStore, Scope, Turn #: 关键词全不命中时，允许"退回最近记忆"的记忆总量上限。 #: 超过这个数就退回严格匹配——否则每轮都会把记忆全塞进 prompt， #: 既不再是"按需检索"，也会持续推高首 token 延迟。 FALLBACK_MAX_CORPUS = 10 SCHEMA = """ PRAGMA journal_mode = WAL; PRAGMA foreign_keys = ON; CREATE TABLE IF NOT EXISTS users ( user_id TEXT PRIMARY KEY, display_name TEXT NOT NULL DEFAULT '', trust_level TEXT NOT NULL DEFAULT 'normal', created_at REAL NOT NULL ); CREATE TABLE IF NOT EXISTS sessions ( session_id TEXT PRIMARY KEY, user_id TEXT NOT NULL REFERENCES users(user_id) ON DELETE CASCADE, started_at REAL NOT NULL, ended_at REAL ); CREATE TABLE IF NOT EXISTS turns ( turn_id TEXT PRIMARY KEY, session_id TEXT NOT NULL REFERENCES sessions(session_id) ON DELETE CASCADE, user_id TEXT NOT NULL REFERENCES users(user_id) ON DELETE CASCADE, role TEXT NOT NULL, text TEXT NOT NULL, created_at REAL NOT NULL ); CREATE TABLE IF NOT EXISTS memories ( memory_id INTEGER PRIMARY KEY AUTOINCREMENT, owner_user_id TEXT NOT NULL REFERENCES users(user_id) ON DELETE CASCADE, scope TEXT NOT NULL DEFAULT 'private', kind TEXT NOT NULL DEFAULT 'fact', mem_key TEXT NOT NULL DEFAULT '', mem_value TEXT NOT NULL, confirm_status TEXT NOT NULL DEFAULT 'confirmed', source_session_id TEXT, source_turn_id TEXT, created_at REAL NOT NULL, updated_at REAL NOT NULL, expires_at REAL ); CREATE INDEX IF NOT EXISTS idx_mem_owner ON memories(owner_user_id, confirm_status); CREATE INDEX IF NOT EXISTS idx_mem_key ON memories(owner_user_id, kind, mem_key); CREATE INDEX IF NOT EXISTS idx_turns_sess ON turns(session_id, created_at); CREATE INDEX IF NOT EXISTS idx_sess_user ON sessions(user_id); """ def _ensure_user_row(c: sqlite3.Connection, user_id: str, display_name: str = "") -> None: """保证 users 里有这一行。 为什么放在存储层而不是让调用方自己 ensure： memories / turns / sessions 都有指向 users 的外键， 「被引用的 user 必须存在」是这个 schema 的不变量。 让每个调用点各自记得先建档，早晚漏一处——而且漏了之后 报的是 "FOREIGN KEY constraint failed"，跟业务逻辑对不上号， 排查起来很绕。由存储层自己兜住，调用方就不需要知道这个约束。 实测踩过：会话中途切到一个未建档的 user_id， 第一轮写 turns 就炸在这里。 """ c.execute("INSERT INTO users(user_id, display_name, created_at) VALUES(?,?,?) " "ON CONFLICT(user_id) DO NOTHING", (user_id, display_name, time.time())) def _row_to_memory(row: sqlite3.Row) -> Memory: return Memory( memory_id=row["memory_id"], owner_user_id=row["owner_user_id"], scope=Scope(row["scope"]), kind=MemoryKind(row["kind"]), key=row["mem_key"], value=row["mem_value"], confirm_status=ConfirmStatus(row["confirm_status"]), source_session_id=row["source_session_id"], source_turn_id=row["source_turn_id"], created_at=row["created_at"], updated_at=row["updated_at"], expires_at=row["expires_at"], ) class SQLiteMemoryStore(MemoryStore): def __init__(self, db_path: Path | str): self.db_path = Path(db_path) self._conn: sqlite3.Connection | None = None self._lock = threading.Lock() # ---------------------------------------------------------- 内部 def _db(self) -> sqlite3.Connection: if self._conn is None: raise RuntimeError("MemoryStore 未 open()") return self._conn def _exec(self, fn): """在受控线程里跑同步 SQL，串行化以免共用连接出问题。""" with self._lock: return fn(self._db()) async def _run(self, fn): return await asyncio.to_thread(self._exec, fn) # ---------------------------------------------------------- 生命周期 async def open(self) -> None: def _open() -> sqlite3.Connection: self.db_path.parent.mkdir(parents=True, exist_ok=True) conn = sqlite3.connect(self.db_path, check_same_thread=False) conn.row_factory = sqlite3.Row conn.executescript(SCHEMA) conn.commit() return conn # 此处还没有并发使用者，不需要持锁跨 await self._conn = await asyncio.to_thread(_open) async def close(self) -> None: if self._conn is None: return def _close(c: sqlite3.Connection): c.commit() c.close() self._conn = None await self._run(_close) # ---------------------------------------------------------- 用户与会话 async def ensure_user(self, user_id: str, display_name: str = "", trust_level: str = "normal") -> None: def op(c: sqlite3.Connection): c.execute( "INSERT INTO users(user_id, display_name, trust_level, created_at) " "VALUES(?,?,?,?) ON CONFLICT(user_id) DO NOTHING", (user_id, display_name, trust_level, time.time())) c.commit() await self._run(op) async def start_session(self, session_id: str, user_id: str) -> None: def op(c: sqlite3.Connection): c.execute("INSERT INTO users(user_id, created_at) VALUES(?,?) " "ON CONFLICT(user_id) DO NOTHING", (user_id, time.time())) c.execute("INSERT INTO sessions(session_id, user_id, started_at) " "VALUES(?,?,?) ON CONFLICT(session_id) DO NOTHING", (session_id, user_id, time.time())) c.commit() await self._run(op) async def end_session(self, session_id: str) -> None: def op(c: sqlite3.Connection): c.execute("UPDATE sessions SET ended_at=? WHERE session_id=?", (time.time(), session_id)) c.commit() await self._run(op) async def add_turn(self, turn: Turn) -> None: def op(c: sqlite3.Connection): _ensure_user_row(c, turn.user_id) c.execute("INSERT INTO turns(turn_id, session_id, user_id, role, text, created_at) " "VALUES(?,?,?,?,?,?) ON CONFLICT(turn_id) DO NOTHING", (turn.turn_id, turn.session_id, turn.user_id, turn.role, turn.text, turn.created_at)) c.commit() await self._run(op) # ---------------------------------------------------------- 记忆读写 async def remember(self, mem: Memory) -> int: now = time.time() def op(c: sqlite3.Connection): _ensure_user_row(c, mem.owner_user_id) # 已确认的同类同键记忆：新值替换旧值（方案 §8） if mem.confirm_status == ConfirmStatus.CONFIRMED and mem.key: row = c.execute( "SELECT memory_id FROM memories WHERE owner_user_id=? AND kind=? " "AND mem_key=? AND confirm_status=?", (mem.owner_user_id, mem.kind.value, mem.key, ConfirmStatus.CONFIRMED.value)).fetchone() if row: c.execute( "UPDATE memories SET mem_value=?, scope=?, updated_at=?, " "expires_at=?, source_session_id=?, source_turn_id=? " "WHERE memory_id=? AND owner_user_id=?", (mem.value, mem.scope.value, now, mem.expires_at, mem.source_session_id, mem.source_turn_id, row["memory_id"], mem.owner_user_id)) c.commit() return row["memory_id"] cur = c.execute( "INSERT INTO memories(owner_user_id, scope, kind, mem_key, mem_value, " "confirm_status, source_session_id, source_turn_id, created_at, " "updated_at, expires_at) VALUES(?,?,?,?,?,?,?,?,?,?,?)", (mem.owner_user_id, mem.scope.value, mem.kind.value, mem.key, mem.value, mem.confirm_status.value, mem.source_session_id, mem.source_turn_id, mem.created_at or now, now, mem.expires_at)) c.commit() return int(cur.lastrowid) return await self._run(op) async def forget(self, owner_user_id: str, memory_id: int) -> bool: def op(c: sqlite3.Connection): # 必须同时匹配 owner——不能只凭 memory_id 删 cur = c.execute("DELETE FROM memories WHERE memory_id=? AND owner_user_id=?", (memory_id, owner_user_id)) c.commit() return cur.rowcount > 0 return await self._run(op) async def retrieve(self, owner_user_id: str, query: str, *, limit: int = 5, include_candidates: bool = False, scopes: Sequence[Scope] | None = None) -> list[Memory]: allowed = [s.value for s in (scopes or [Scope.PRIVATE, Scope.SHARED])] statuses = [ConfirmStatus.CONFIRMED.value] if include_candidates: statuses.append(ConfirmStatus.CANDIDATE.value) placeholders = ",".join("?" * len(allowed)) status_ph = ",".join("?" * len(statuses)) def op(c: sqlite3.Connection): rows = c.execute( f"SELECT * FROM memories WHERE owner_user_id=? " f"AND scope IN ({placeholders}) AND confirm_status IN ({status_ph})", (owner_user_id, *allowed, *statuses)).fetchall() now = time.time() scored = [] for r in rows: m = _row_to_memory(r) if m.is_expired(now): continue score = self.relevance(query, m.text) if query.strip() else 0.0 scored.append((score, m.updated_at, m)) hits = [x for x in scored if x[0] > 0] if not hits and 0 < len(scored) <= FALLBACK_MAX_CORPUS: # 关键词一条都没命中时，退回"最近更新的若干条"。 # # 理由是实测出来的：查"我想喝点什么"匹配不上"喜欢喝美式咖啡"—— # 两句共享"喝"字，但没有共同的二元组，按字面打分是 0。 # 中文口语短句和记忆文本之间"意思相关但用词不同"极其常见， # 纯字面匹配漏得厉害。 # # 但兜底必须**限量**：不加限制就变成"每轮无条件注入记忆"， # 那就不是方案 §8 说的"按需检索"了，而且记忆一多会持续 # 拉长 prompt、推高首 token 延迟。 # # 所以只在记忆总量很小时兜底——这个阶段"全给"和"按需给" # 结果差不多，而漏给的代价（助手像不认识你）更大。 # 超过阈值就退回严格匹配，接受漏召。 # # 这是权宜之计，不是终局：真正的解法是向量检索 # （方案 §8 把它列在"真实使用中检索不足再评估"）。 # 到那时换掉打分和这段逻辑即可，接口不动。 hits = scored hits.sort(key=lambda x: (x[0], x[1]), reverse=True) return [m for _, _, m in hits[:limit]] return await self._run(op) async def list_memories(self, owner_user_id: str, *, include_candidates: bool = True) -> list[Memory]: statuses = [ConfirmStatus.CONFIRMED.value] if include_candidates: statuses.append(ConfirmStatus.CANDIDATE.value) status_ph = ",".join("?" * len(statuses)) def op(c: sqlite3.Connection): rows = c.execute( f"SELECT * FROM memories WHERE owner_user_id=? " f"AND confirm_status IN ({status_ph}) ORDER BY updated_at DESC", (owner_user_id, *statuses)).fetchall() now = time.time() return [m for m in (_row_to_memory(r) for r in rows) if not m.is_expired(now)] return await self._run(op) async def purge_user(self, owner_user_id: str) -> int: def op(c: sqlite3.Connection): cur = c.execute("DELETE FROM memories WHERE owner_user_id=?", (owner_user_id,)) n = cur.rowcount c.execute("DELETE FROM turns WHERE user_id=?", (owner_user_id,)) c.execute("DELETE FROM sessions WHERE user_id=?", (owner_user_id,)) # users 放最后：外键级联会把残留一并带走 c.execute("DELETE FROM users WHERE user_id=?", (owner_user_id,)) c.commit() return n return await self._run(op) __all__ = ["SQLiteMemoryStore", "SCHEMA"]