#!/usr/bin/env python3 """记忆层测试。 重点不是"能存能取"，而是方案 §8 反复强调的隔离性质： - 未知身份不碰私人记忆 - 删别人的记忆删不掉 - 自动抽取的候选不参与检索 - 删用户要连会话和摘要一起删干净 这些是"错了很难发现、发现了很严重"的类型，所以单独测。 """ from __future__ import annotations import asyncio import shutil import sys import tempfile import time from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) from app.memory import ( # noqa: E402 ConfirmStatus, Memory, MemoryKind, Scope, SQLiteMemoryStore, Turn, ) from app.memory.sqlite_store import FALLBACK_MAX_CORPUS # noqa: E402 OK, FAIL = [], [] def check(name: str, cond: bool, detail: str = "") -> None: (OK if cond else FAIL).append(name) print(f" [{'PASS' if cond else 'FAIL'}] {name}" + (f" — {detail}" if detail else "")) async def run(store: SQLiteMemoryStore) -> None: await store.open() print("\n=== 1. 基本读写 ===") await store.ensure_user("u_ylhan", "ylhan") await store.ensure_user("u_guest", "访客") mid = await store.remember(Memory( owner_user_id="u_ylhan", kind=MemoryKind.PREFERENCE, key="回复风格", value="喜欢简短直接，不要长篇解释")) check("写入返回 memory_id", isinstance(mid, int) and mid > 0, str(mid)) got = await store.retrieve("u_ylhan", "回复风格") check("能检索到", len(got) == 1 and got[0].key == "回复风格") check("字段完整", got[0].kind == MemoryKind.PREFERENCE and got[0].value.startswith("喜欢")) print("\n=== 2. 身份隔离（最关键）===") other = await store.retrieve("u_guest", "回复风格") check("别人检索不到", len(other) == 0, f"拿到 {len(other)} 条") listed = await store.list_memories("u_guest") check("别人列不出", len(listed) == 0) # 用 memory_id 越权删 deleted = await store.forget("u_guest", mid) check("越权删除失败", deleted is False) still = await store.retrieve("u_ylhan", "回复风格") check("原主记忆还在", len(still) == 1) print("\n=== 3. 同键新值替换旧值 ===") await store.remember(Memory( owner_user_id="u_ylhan", kind=MemoryKind.PREFERENCE, key="回复风格", value="喜欢更活泼一点的说法")) rows = await store.retrieve("u_ylhan", "回复风格") check("同键只剩一条", len(rows) == 1, f"{len(rows)} 条") check("值已更新", "活泼" in rows[0].value, rows[0].value) # 不同 key 不应互相覆盖 await store.remember(Memory(owner_user_id="u_ylhan", kind=MemoryKind.PREFERENCE, key="称呼", value="叫他老韩")) rows2 = await store.list_memories("u_ylhan") check("不同 key 各自保留", len(rows2) == 2, f"{len(rows2)} 条") print("\n=== 4. 候选记忆默认不参与检索 ===") # 判据是"这条候选不在结果里"，**不是"结果为空"**。 # 后者是实现细节：检索在关键词不命中时会退回最近记忆 # （见 sqlite_store.FALLBACK_MAX_CORPUS），结果本来就可能非空。 CAND = "用户好像养了一只叫豆豆的猫" await store.remember(Memory( owner_user_id="u_ylhan", kind=MemoryKind.FACT, key="", value=CAND, confirm_status=ConfirmStatus.CANDIDATE)) hit = await store.retrieve("u_ylhan", "豆豆") check("候选不出现在检索结果里", all(m.value != CAND for m in hit), f"拿到 {len(hit)} 条：{[m.value[:12] for m in hit]}") hit2 = await store.retrieve("u_ylhan", "豆豆", include_candidates=True) check("显式要求时可检索到", any(m.value == CAND for m in hit2), f"拿到 {len(hit2)} 条") allm = await store.list_memories("u_ylhan") check("候选确实已落库", any(m.confirm_status == ConfirmStatus.CANDIDATE for m in allm)) print("\n=== 5. 过期 ===") EXPIRED = "十分钟后提醒我关火" await store.remember(Memory( owner_user_id="u_ylhan", kind=MemoryKind.TASK, key="", value=EXPIRED, expires_at=time.time() - 1)) # 已过期 expired = await store.retrieve("u_ylhan", "关火") check("过期记忆不出现在结果里", all(m.value != EXPIRED for m in expired), f"拿到 {len(expired)} 条") check("过期记忆也不在列表里", all(m.value != EXPIRED for m in await store.list_memories("u_ylhan"))) print("\n=== 6. 会话与轮次 ===") await store.start_session("sess_1", "u_ylhan") for i, (role, text) in enumerate([("user", "今天天气怎么样"), ("assistant", "今天晴，十八到二十六度。")]): await store.add_turn(Turn(turn_id=f"t{i}", session_id="sess_1", user_id="u_ylhan", role=role, text=text)) check("会话与轮次写入无异常", True) print("\n=== 7. 中文检索打分 ===") r = SQLiteMemoryStore.relevance("客厅的灯", "客厅有一盏米家吸顶灯") check("相关文本有分", r > 0.3, f"{r:.2f}") r2 = SQLiteMemoryStore.relevance("客厅的灯", "喜欢听爵士乐") check("无关文本为 0", r2 == 0.0, f"{r2:.2f}") r3 = SQLiteMemoryStore.relevance("回复风格", "回复风格：喜欢简短") check("子串直接命中", r3 == 1.0, f"{r3:.2f}") print("\n=== 8. 删除用户要清干净 ===") n = await store.purge_user("u_ylhan") check("返回删除的记忆数", n >= 3, f"{n} 条") check("记忆已清空", len(await store.list_memories("u_ylhan")) == 0) check("检索也拿不到", len(await store.retrieve("u_ylhan", "回复风格")) == 0) # 会话和轮次靠外键级联删除 rows = await store._run(lambda c: c.execute( "SELECT COUNT(*) AS n FROM turns WHERE user_id=?", ("u_ylhan",)).fetchone()["n"]) check("轮次一并删除", rows == 0, f"残留 {rows} 条") rows2 = await store._run(lambda c: c.execute( "SELECT COUNT(*) AS n FROM sessions WHERE user_id=?", ("u_ylhan",)).fetchone()["n"]) check("会话一并删除", rows2 == 0, f"残留 {rows2} 条") print("\n=== 9. 重复建用户不报错、不覆盖 ===") await store.ensure_user("u_x", "第一次") await store.ensure_user("u_x", "第二次") name = await store._run(lambda c: c.execute( "SELECT display_name FROM users WHERE user_id=?", ("u_x",)).fetchone()["display_name"]) check("已存在则不改", name == "第一次", name) await store.set_user_display_name("u_x", "林") check("可显式更新用户称呼", await store.get_user_display_name("u_x") == "林") check("不存在的用户称呼为空", await store.get_user_display_name("u_missing") == "") print("\n=== 10. 兜底检索要限量（否则就成了每轮全量注入）===") await store.ensure_user("u_big") await store.remember(Memory(owner_user_id="u_big", kind=MemoryKind.FACT, value="常喝拿铁")) # "推荐个饮品"和"常喝拿铁"没有任何共同字符，字面打分为 0 got = await store.retrieve("u_big", "推荐个饮品") check("记忆少时兜底返回最近一条", len(got) == 1, f"{len(got)} 条") for i in range(FALLBACK_MAX_CORPUS + 2): await store.remember(Memory(owner_user_id="u_big", kind=MemoryKind.FACT, value=f"无关记录编号{i}")) got2 = await store.retrieve("u_big", "推荐个饮品") check("记忆多时不再无条件兜底", len(got2) == 0, f"{len(got2)} 条") got3 = await store.retrieve("u_big", "拿铁") check("关键词命中仍生效", any("拿铁" in m.value for m in got3), f"{len(got3)} 条") await store.close() def main() -> int: tmp = Path(tempfile.mkdtemp(prefix="memtest_")) try: store = SQLiteMemoryStore(tmp / "memory.db") asyncio.run(run(store)) finally: shutil.rmtree(tmp, ignore_errors=True) print("\n" + "=" * 62) print(f"通过 {len(OK)} / {len(OK) + len(FAIL)}") for n in FAIL: print(f" 未通过: {n}") return 0 if not FAIL else 1 if __name__ == "__main__": sys.exit(main())