RRRosmontis
34089d0042
Include utility modules, shared channels, prompts, read-only skills, wardrobe memory and regression tests. Externalize legacy credentials and exclude private runtime state.
371 lines
14 KiB
Python
371 lines
14 KiB
Python
# -*- coding: utf-8 -*-
|
||
# memory_store.py — 结构化长期记忆存储（SQLite + FTS5，零外部依赖）
|
||
#
|
||
# 事实库: facts(uid, content, importance, keywords, created_at, archived, source)
|
||
# 召回: FTS5 trigram（≥3 字子串）+ LIKE 兜底（短查询/中文双字）
|
||
# 迁移: 首次接触某 uid 时把 memory/<uid>.txt 的旧文本记忆导入 facts（meta 表幂等）
|
||
import json
|
||
import os
|
||
import re
|
||
import sqlite3
|
||
import threading
|
||
import time
|
||
|
||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
DB_PATH = os.environ.get("CHAT_MEMORY_DB", os.path.join(BASE_DIR, "memory.db"))
|
||
LEGACY_DIR = os.environ.get("CHAT_MEMORY_LEGACY_DIR", os.path.join(BASE_DIR, "memory"))
|
||
|
||
_lock = threading.RLock()
|
||
|
||
|
||
def _connect():
|
||
con = sqlite3.connect(DB_PATH, timeout=30)
|
||
con.execute("PRAGMA journal_mode=WAL")
|
||
return con
|
||
|
||
|
||
def init_db():
|
||
with _lock, _connect() as con:
|
||
con.executescript("""
|
||
CREATE TABLE IF NOT EXISTS facts(
|
||
id INTEGER PRIMARY KEY,
|
||
uid INTEGER NOT NULL,
|
||
content TEXT NOT NULL,
|
||
importance INTEGER NOT NULL DEFAULT 3,
|
||
keywords TEXT NOT NULL DEFAULT '',
|
||
created_at REAL NOT NULL,
|
||
archived INTEGER NOT NULL DEFAULT 0,
|
||
source TEXT NOT NULL DEFAULT ''
|
||
);
|
||
CREATE VIRTUAL TABLE IF NOT EXISTS facts_fts USING fts5(
|
||
content, keywords, content='facts', content_rowid='id', tokenize='trigram'
|
||
);
|
||
CREATE TRIGGER IF NOT EXISTS facts_ai AFTER INSERT ON facts BEGIN
|
||
INSERT INTO facts_fts(rowid, content, keywords) VALUES (new.id, new.content, new.keywords);
|
||
END;
|
||
CREATE TRIGGER IF NOT EXISTS facts_ad AFTER DELETE ON facts BEGIN
|
||
INSERT INTO facts_fts(facts_fts, rowid, content, keywords)
|
||
VALUES ('delete', old.id, old.content, old.keywords);
|
||
END;
|
||
CREATE TRIGGER IF NOT EXISTS facts_au AFTER UPDATE ON facts BEGIN
|
||
INSERT INTO facts_fts(facts_fts, rowid, content, keywords)
|
||
VALUES ('delete', old.id, old.content, old.keywords);
|
||
INSERT INTO facts_fts(rowid, content, keywords) VALUES (new.id, new.content, new.keywords);
|
||
END;
|
||
CREATE TABLE IF NOT EXISTS meta(key TEXT PRIMARY KEY, value TEXT);
|
||
CREATE TABLE IF NOT EXISTS triples(
|
||
id INTEGER PRIMARY KEY,
|
||
uid INTEGER NOT NULL,
|
||
subject TEXT NOT NULL,
|
||
predicate TEXT NOT NULL,
|
||
object TEXT NOT NULL,
|
||
importance INTEGER NOT NULL DEFAULT 3,
|
||
valid_from REAL NOT NULL,
|
||
valid_until REAL,
|
||
created_at REAL NOT NULL,
|
||
source TEXT NOT NULL DEFAULT 'model'
|
||
);
|
||
CREATE INDEX IF NOT EXISTS triples_uid ON triples(uid);
|
||
CREATE INDEX IF NOT EXISTS triples_subj ON triples(uid, subject);
|
||
CREATE TABLE IF NOT EXISTS relations(
|
||
uid INTEGER PRIMARY KEY,
|
||
first_met REAL,
|
||
msg_count INTEGER NOT NULL DEFAULT 0,
|
||
affinity INTEGER NOT NULL DEFAULT 50,
|
||
last_seen REAL,
|
||
notes TEXT NOT NULL DEFAULT '[]'
|
||
);
|
||
""")
|
||
|
||
|
||
def _now():
|
||
return time.time()
|
||
|
||
|
||
def _norm_content(s):
|
||
return re.sub(r"\s+", " ", str(s or "")).strip()
|
||
|
||
|
||
def _norm_keywords(kw):
|
||
if isinstance(kw, str):
|
||
kw = [k.strip() for k in re.split(r"[,，、/ ]+", kw) if k.strip()]
|
||
elif isinstance(kw, (list, tuple)):
|
||
kw = [str(k).strip() for k in kw if str(k).strip()]
|
||
else:
|
||
kw = []
|
||
return ",".join(kw[:8])
|
||
|
||
|
||
def add_facts(uid, facts, source="model"):
|
||
"""写入若干条事实；同 uid 完全相同（content 一致）的活跃条目去重跳过。
|
||
返回新增的 id 列表。"""
|
||
init_db()
|
||
uid = int(uid)
|
||
added = []
|
||
with _lock, _connect() as con:
|
||
for f in facts or []:
|
||
if isinstance(f, str):
|
||
f = {"content": f}
|
||
if not isinstance(f, dict):
|
||
continue
|
||
content = _norm_content(f.get("content"))
|
||
if not content or len(content) < 2:
|
||
continue
|
||
importance = f.get("importance", 3)
|
||
try:
|
||
importance = int(importance)
|
||
except (TypeError, ValueError):
|
||
importance = 3
|
||
importance = max(1, min(5, importance))
|
||
keywords = _norm_keywords(f.get("keywords"))
|
||
if con.execute(
|
||
"SELECT 1 FROM facts WHERE uid=? AND content=? AND archived=0",
|
||
(uid, content)).fetchone():
|
||
continue
|
||
cur = con.execute(
|
||
"INSERT INTO facts(uid, content, importance, keywords, created_at, archived, source)"
|
||
" VALUES (?,?,?,?,?,0,?)",
|
||
(uid, content, importance, keywords, _now(), source))
|
||
added.append(cur.lastrowid)
|
||
return added
|
||
|
||
|
||
def list_facts(uid, include_archived=False, limit=50):
|
||
init_db()
|
||
q = ("SELECT id, content, importance, keywords, created_at, archived "
|
||
"FROM facts WHERE uid=?")
|
||
if not include_archived:
|
||
q += " AND archived=0"
|
||
q += " ORDER BY created_at DESC LIMIT ?"
|
||
with _lock, _connect() as con:
|
||
return con.execute(q, (int(uid), limit)).fetchall()
|
||
|
||
|
||
def forget_fact(uid, fact_id):
|
||
"""归档一条事实（不物理删除，可审计）。返回是否命中。"""
|
||
init_db()
|
||
with _lock, _connect() as con:
|
||
cur = con.execute(
|
||
"UPDATE facts SET archived=1 WHERE uid=? AND id=? AND archived=0",
|
||
(int(uid), int(fact_id)))
|
||
return cur.rowcount > 0
|
||
|
||
|
||
def count_facts(uid, active_only=True):
|
||
init_db()
|
||
q = "SELECT COUNT(*) FROM facts WHERE uid=?"
|
||
if active_only:
|
||
q += " AND archived=0"
|
||
with _lock, _connect() as con:
|
||
return con.execute(q, (int(uid),)).fetchone()[0]
|
||
|
||
|
||
def search_facts(uid, query, limit=5):
|
||
"""FTS5 trigram 召回（≥3 字）；短查询/无命中时 LIKE 兜底。"""
|
||
init_db()
|
||
uid = int(uid)
|
||
query = _norm_content(query)
|
||
if not query:
|
||
return []
|
||
rows = []
|
||
with _lock, _connect() as con:
|
||
if len(query) >= 3:
|
||
try:
|
||
rows = con.execute(
|
||
"SELECT f.id, f.content, f.importance, f.created_at "
|
||
"FROM facts_fts JOIN facts f ON f.id = facts_fts.rowid "
|
||
"WHERE facts_fts MATCH ? AND f.uid=? AND f.archived=0 "
|
||
"ORDER BY rank, f.created_at DESC LIMIT ?",
|
||
(query, uid, limit)).fetchall()
|
||
except sqlite3.OperationalError:
|
||
rows = []
|
||
if not rows:
|
||
like = f"%{query}%"
|
||
rows = con.execute(
|
||
"SELECT id, content, importance, created_at FROM facts "
|
||
"WHERE uid=? AND archived=0 AND (content LIKE ? OR keywords LIKE ?) "
|
||
"ORDER BY created_at DESC LIMIT ?",
|
||
(uid, like, like, limit)).fetchall()
|
||
return rows
|
||
|
||
|
||
def touch_relation(uid):
|
||
"""每轮对话调用：自动登记/更新关系统计（相识时间、对话次数、最近活跃）。
|
||
每 50 轮顺手做一次程序式整理（归档低重要度老记忆，不花 token）。"""
|
||
init_db()
|
||
uid = int(uid)
|
||
now = _now()
|
||
with _lock, _connect() as con:
|
||
row = con.execute("SELECT msg_count FROM relations WHERE uid=?", (uid,)).fetchone()
|
||
if row:
|
||
new_count = row[0] + 1
|
||
con.execute("UPDATE relations SET msg_count=?, last_seen=? WHERE uid=?",
|
||
(new_count, now, uid))
|
||
else:
|
||
new_count = 1
|
||
con.execute(
|
||
"INSERT INTO relations(uid, first_met, msg_count, affinity, last_seen, notes)"
|
||
" VALUES (?,?,1,50,?,'[]')", (uid, now, now))
|
||
if new_count % 50 == 0:
|
||
maybe_archive_stale(uid)
|
||
|
||
|
||
def get_relation(uid):
|
||
init_db()
|
||
with _lock, _connect() as con:
|
||
row = con.execute(
|
||
"SELECT first_met, msg_count, affinity, last_seen, notes FROM relations WHERE uid=?",
|
||
(int(uid),)).fetchone()
|
||
if not row:
|
||
return None
|
||
try:
|
||
notes = json.loads(row[4] or "[]")
|
||
except Exception:
|
||
notes = []
|
||
return {"first_met": row[0], "msg_count": row[1], "affinity": row[2],
|
||
"last_seen": row[3], "notes": notes}
|
||
|
||
|
||
def add_relation_note(uid, note):
|
||
"""模型写的「关系上的重要变化/约定/里程碑」（一句话）。notes 保留最近 5 条。"""
|
||
init_db()
|
||
uid = int(uid)
|
||
note = _norm_content(note)
|
||
if not note:
|
||
return
|
||
with _lock, _connect() as con:
|
||
row = con.execute("SELECT notes FROM relations WHERE uid=?", (uid,)).fetchone()
|
||
if not row:
|
||
touch_relation(uid)
|
||
row = ("[]",)
|
||
try:
|
||
notes = json.loads(row[0] or "[]")
|
||
except Exception:
|
||
notes = []
|
||
notes.append({"ts": _now(), "text": note[:200]})
|
||
notes = notes[-5:]
|
||
con.execute("UPDATE relations SET notes=? WHERE uid=?",
|
||
(json.dumps(notes, ensure_ascii=False), uid))
|
||
|
||
|
||
def maybe_archive_stale(uid, older_than_days=30, max_importance=2):
|
||
"""轻量整理：归档「老且不重要」的记忆（程序式，不花 token）。返回归档条数。"""
|
||
init_db()
|
||
cutoff = _now() - older_than_days * 86400
|
||
with _lock, _connect() as con:
|
||
cur = con.execute(
|
||
"UPDATE facts SET archived=1 WHERE uid=? AND archived=0 AND importance<=? AND created_at<?",
|
||
(int(uid), max_importance, cutoff))
|
||
return cur.rowcount
|
||
|
||
|
||
def add_links(uid, links, source="model"):
|
||
"""写入 SPO 三元组。同 (uid, subject, predicate) 的旧边若 object 不同，
|
||
置 valid_until（时效窗：旧事实保留但不再生效）。返回新增 id 列表。"""
|
||
init_db()
|
||
uid = int(uid)
|
||
now = _now()
|
||
added = []
|
||
with _lock, _connect() as con:
|
||
for lk in links or []:
|
||
if not isinstance(lk, dict):
|
||
continue
|
||
s = _norm_content(lk.get("subject"))
|
||
pr = _norm_content(lk.get("predicate"))
|
||
o = _norm_content(lk.get("object"))
|
||
if not s or not pr or not o:
|
||
continue
|
||
try:
|
||
importance = max(1, min(5, int(lk.get("importance", 3))))
|
||
except (TypeError, ValueError):
|
||
importance = 3
|
||
# 同主同谓的旧边：值相同则跳过；值不同则置时效窗
|
||
old = con.execute(
|
||
"SELECT id, object FROM triples WHERE uid=? AND subject=? AND predicate=?"
|
||
" AND valid_until IS NULL", (uid, s, pr)).fetchall()
|
||
if any(row[1] == o for row in old):
|
||
continue
|
||
for row in old:
|
||
con.execute("UPDATE triples SET valid_until=? WHERE id=?", (now, row[0]))
|
||
cur = con.execute(
|
||
"INSERT INTO triples(uid, subject, predicate, object, importance,"
|
||
" valid_from, valid_until, created_at, source) VALUES (?,?,?,?,?,?,NULL,?,?)",
|
||
(uid, s, pr, o, importance, now, now, source))
|
||
added.append(cur.lastrowid)
|
||
return added
|
||
|
||
|
||
def entity_card(uid, name, limit=10):
|
||
"""某实体的卡片：以它为主语或宾语的当前有效三元组。"""
|
||
init_db()
|
||
uid = int(uid)
|
||
name = _norm_content(name)
|
||
if not name:
|
||
return []
|
||
with _lock, _connect() as con:
|
||
return con.execute(
|
||
"SELECT subject, predicate, object, importance FROM triples "
|
||
"WHERE uid=? AND valid_until IS NULL AND (subject=? OR object=?) "
|
||
"ORDER BY importance DESC, created_at DESC LIMIT ?",
|
||
(uid, name, name, limit)).fetchall()
|
||
|
||
|
||
def list_entities(uid, limit=50):
|
||
"""全部实体（主语/宾语出现过的名字）及各自有效边数。"""
|
||
init_db()
|
||
with _lock, _connect() as con:
|
||
rows = con.execute(
|
||
"SELECT name, COUNT(*) FROM ("
|
||
" SELECT subject AS name FROM triples WHERE uid=? AND valid_until IS NULL"
|
||
" UNION ALL SELECT object FROM triples WHERE uid=? AND valid_until IS NULL"
|
||
") GROUP BY name ORDER BY COUNT(*) DESC LIMIT ?",
|
||
(int(uid), int(uid), limit)).fetchall()
|
||
return rows
|
||
|
||
|
||
def recall_facts(uid, query, top_k=5, pinned_k=3):
|
||
"""召回式读取：高重要度置顶 pinned_k 条 + 按当前输入 FTS5 召回 top_k 条，按 id 去重。
|
||
返回 [(id, content, importance), ...]"""
|
||
init_db()
|
||
uid = int(uid)
|
||
with _lock, _connect() as con:
|
||
pinned = con.execute(
|
||
"SELECT id, content, importance FROM facts WHERE uid=? AND archived=0 "
|
||
"ORDER BY importance DESC, created_at DESC LIMIT ?",
|
||
(uid, pinned_k)).fetchall()
|
||
seen = {r[0] for r in pinned}
|
||
out = list(pinned)
|
||
for r in search_facts(uid, query, limit=top_k):
|
||
if r[0] not in seen:
|
||
out.append((r[0], r[1], r[2]))
|
||
seen.add(r[0])
|
||
return out
|
||
|
||
|
||
def migrate_legacy(uid, legacy_dir=None):
|
||
"""首次接触某 uid 时把 memory/<uid>.txt 的旧文本记忆导入 facts（meta 幂等）。"""
|
||
init_db()
|
||
uid = int(uid)
|
||
legacy_dir = legacy_dir or LEGACY_DIR
|
||
marker = f"imported:{uid}"
|
||
with _lock, _connect() as con:
|
||
if con.execute("SELECT 1 FROM meta WHERE key=?", (marker,)).fetchone():
|
||
return 0
|
||
path = os.path.join(legacy_dir, f"{uid}.txt")
|
||
if not os.path.exists(path):
|
||
return 0
|
||
facts = []
|
||
with open(path, encoding="utf-8") as f:
|
||
for ln in f:
|
||
ln = ln.strip()
|
||
if not ln:
|
||
continue
|
||
m = re.match(r"^\[\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\]\s*(.*)$", ln)
|
||
facts.append({"content": m.group(1) if m else ln,
|
||
"importance": 3, "keywords": []})
|
||
added = len(add_facts(uid, facts, source="legacy"))
|
||
with _lock, _connect() as con:
|
||
con.execute("INSERT OR REPLACE INTO meta(key, value) VALUES (?,?)",
|
||
(marker, str(_now())))
|
||
return added
|