Initial import of Shirobana bot backend
Include utility modules, shared channels, prompts, read-only skills, wardrobe memory and regression tests. Externalize legacy credentials and exclude private runtime state.
This commit is contained in:
RRRosmontis
@@ -0,0 +1,370 @@
|
||||
# -*- 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
|
||||
Reference in New Issue
Block a user