ylhan
20cfaa6ef7
这是进入 Pipecat Flows 改造前的回滚点。 提交时已核对本地与线上代码 md5 一致（engine.py 7e97bad8、 bot_server.py 1cf543ae），所以这个提交就是当前线上状态。 内容（此前一直在线上运行但未提交的改动）： - 姓名档案：用户自报"我叫林"同步写 confirmed 姓名，重连后恢复 - 疑问句防误写："谁吗"之类的疑问句不再被覆盖成姓名 - 新增 test_profile_memory.py（姓名档案专项回归） - 新增 test_profile_live.py（需要真实 LLM 的现场测试） - 搜索、身份、记忆测试的相应补充 背景：下一步要按 P1-Pipecat Flows Agent化与记忆改造交接.md 渐进引入 Pipecat Flows，交接文档明确要求先保护基线、不得 reset/checkout/覆盖现有改动。此提交即该基线。
354 lines
15 KiB
Python
354 lines
15 KiB
Python
"""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 set_user_display_name(self, user_id: str, display_name: str) -> None:
|
|
display_name = " ".join((display_name or "").split()).strip()
|
|
if not user_id or not display_name:
|
|
raise ValueError("user_id 和 display_name 不能为空")
|
|
|
|
def op(c: sqlite3.Connection):
|
|
_ensure_user_row(c, user_id)
|
|
c.execute("UPDATE users SET display_name=? WHERE user_id=?",
|
|
(display_name, user_id))
|
|
c.commit()
|
|
await self._run(op)
|
|
|
|
async def get_user_display_name(self, user_id: str) -> str:
|
|
def op(c: sqlite3.Connection):
|
|
row = c.execute("SELECT display_name FROM users WHERE user_id=?",
|
|
(user_id,)).fetchone()
|
|
return str(row["display_name"] or "") if row else ""
|
|
return 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"]
|