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:
@@ -0,0 +1,29 @@
|
||||
"""Unified-UID chat engine package."""
|
||||
|
||||
from .engine import (
|
||||
active_push_worker,
|
||||
compact_history,
|
||||
delete_last_history,
|
||||
handle_chat_command,
|
||||
handle_group_message,
|
||||
handle_memory_command,
|
||||
handle_private_message,
|
||||
init_chat,
|
||||
set_delivery,
|
||||
set_websocket,
|
||||
start_active_push,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"active_push_worker",
|
||||
"compact_history",
|
||||
"delete_last_history",
|
||||
"handle_chat_command",
|
||||
"handle_group_message",
|
||||
"handle_memory_command",
|
||||
"handle_private_message",
|
||||
"init_chat",
|
||||
"set_delivery",
|
||||
"set_websocket",
|
||||
"start_active_push",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
@@ -0,0 +1,368 @@
|
||||
"""Thread-safe, scope-isolated local sticker storage.
|
||||
|
||||
The module deliberately performs no network access. ``source['url']`` may be a
|
||||
local filename or a ``file://`` URI when a preview should be copied; remote URLs
|
||||
are retained as metadata but are not downloaded.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from datetime import datetime, timezone
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import random
|
||||
import re
|
||||
import secrets
|
||||
import shutil
|
||||
import tempfile
|
||||
import threading
|
||||
from typing import Any
|
||||
from urllib.parse import unquote, urlparse
|
||||
|
||||
|
||||
SCENES: dict[str, str] = {
|
||||
"daily_cute": "日常卖萌、轻松可爱的互动",
|
||||
"happy": "开心、愉快或表达喜悦",
|
||||
"sad": "难过、失落或情绪低落",
|
||||
"very_sad": "非常难过、哭泣或极度沮丧",
|
||||
"comfort": "安慰、鼓励或陪伴他人",
|
||||
"awkward": "尴尬、无语或不知所措",
|
||||
"angry": "生气、不满或抗议",
|
||||
"surprised": "惊讶、意外或震惊",
|
||||
"thinking": "思考、疑惑或认真琢磨",
|
||||
"greeting": "打招呼、问候或道别",
|
||||
"thanks": "感谢、感激或礼貌回应",
|
||||
"tease": "调侃、逗弄或开玩笑",
|
||||
"celebrate": "庆祝、祝贺或欢呼",
|
||||
"other": "不属于以上场景的其他情绪或用途",
|
||||
}
|
||||
|
||||
ALLOWED_SUFFIXES = frozenset({".webp", ".tgs", ".webm", ".png", ".jpg"})
|
||||
|
||||
|
||||
def scene_prompt() -> str:
|
||||
"""Return model-readable instructions for scene classification."""
|
||||
|
||||
lines = [
|
||||
"请为贴纸选择 1 至 3 个最匹配的场景;只能使用以下场景标识:"
|
||||
]
|
||||
lines.extend(f"- {scene}:{description}" for scene, description in SCENES.items())
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _copy_atomic(source: Path, destination: Path) -> None:
|
||||
"""Copy a file without exposing a partially written destination."""
|
||||
|
||||
fd, temporary_name = tempfile.mkstemp(
|
||||
prefix=f".{destination.name}.", suffix=".tmp", dir=destination.parent
|
||||
)
|
||||
os.close(fd)
|
||||
temporary = Path(temporary_name)
|
||||
try:
|
||||
shutil.copy2(source, temporary)
|
||||
os.replace(temporary, destination)
|
||||
finally:
|
||||
temporary.unlink(missing_ok=True)
|
||||
|
||||
|
||||
def _local_preview(value: Any) -> Path | None:
|
||||
"""Resolve a local preview source without ever performing network I/O."""
|
||||
|
||||
if isinstance(value, os.PathLike):
|
||||
return Path(value)
|
||||
if not isinstance(value, str) or not value:
|
||||
return None
|
||||
parsed = urlparse(value)
|
||||
if parsed.scheme == "file":
|
||||
if parsed.netloc not in {"", "localhost"}:
|
||||
return None
|
||||
return Path(unquote(parsed.path))
|
||||
if parsed.scheme:
|
||||
return None
|
||||
return Path(value)
|
||||
|
||||
|
||||
class StickerStore:
|
||||
"""Persist stickers below ``root/<scope>`` with atomic index updates."""
|
||||
|
||||
def __init__(self, root: str | os.PathLike[str]) -> None:
|
||||
self.root = Path(root).expanduser().resolve()
|
||||
self.root.mkdir(parents=True, exist_ok=True)
|
||||
self._lock = threading.RLock()
|
||||
|
||||
@staticmethod
|
||||
def _scope_name(scope: str | int) -> str:
|
||||
if isinstance(scope, bool):
|
||||
raise ValueError("scope must be a UID or group ID")
|
||||
name = str(scope)
|
||||
if (
|
||||
not name
|
||||
or name in {".", ".."}
|
||||
or "/" in name
|
||||
or "\\" in name
|
||||
or "\x00" in name
|
||||
or Path(name).is_absolute()
|
||||
):
|
||||
raise ValueError("scope must be one safe path component")
|
||||
return name
|
||||
|
||||
@staticmethod
|
||||
def _require_within(path: Path, parent: Path) -> None:
|
||||
try:
|
||||
path.resolve().relative_to(parent.resolve())
|
||||
except ValueError as exc:
|
||||
raise ValueError("sticker storage path escapes its scope") from exc
|
||||
|
||||
def _layout(self, scope: str | int, *, create: bool) -> tuple[Path, Path, Path]:
|
||||
name = self._scope_name(scope)
|
||||
scope_dir = self.root / name
|
||||
objects = scope_dir / "objects"
|
||||
previews = scope_dir / "previews"
|
||||
if create:
|
||||
scope_dir.mkdir(parents=True, exist_ok=True)
|
||||
objects.mkdir(exist_ok=True)
|
||||
previews.mkdir(exist_ok=True)
|
||||
self._require_within(scope_dir, self.root)
|
||||
self._require_within(objects, scope_dir)
|
||||
self._require_within(previews, scope_dir)
|
||||
return scope_dir, objects, previews
|
||||
|
||||
@staticmethod
|
||||
def _empty_index() -> dict[str, Any]:
|
||||
return {"version": 1, "stickers": {}}
|
||||
|
||||
@staticmethod
|
||||
def _load_index(scope_dir: Path) -> dict[str, Any]:
|
||||
index_path = scope_dir / "index.json"
|
||||
if not index_path.exists():
|
||||
return StickerStore._empty_index()
|
||||
with index_path.open("r", encoding="utf-8") as handle:
|
||||
index = json.load(handle)
|
||||
if not isinstance(index, dict) or not isinstance(index.get("stickers"), dict):
|
||||
raise ValueError(f"invalid sticker index: {index_path}")
|
||||
return index
|
||||
|
||||
@staticmethod
|
||||
def _write_index(scope_dir: Path, index: dict[str, Any]) -> None:
|
||||
index_path = scope_dir / "index.json"
|
||||
fd, temporary_name = tempfile.mkstemp(
|
||||
prefix=".index.", suffix=".tmp", dir=scope_dir
|
||||
)
|
||||
temporary = Path(temporary_name)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as handle:
|
||||
json.dump(index, handle, ensure_ascii=False, indent=2, sort_keys=True)
|
||||
handle.write("\n")
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
os.replace(temporary, index_path)
|
||||
finally:
|
||||
temporary.unlink(missing_ok=True)
|
||||
|
||||
@staticmethod
|
||||
def _valid_scenes(classification: dict[str, Any]) -> list[str] | None:
|
||||
if not isinstance(classification, dict) or classification.get("accept") is False:
|
||||
return None
|
||||
raw_scenes = classification.get("scenes")
|
||||
if not isinstance(raw_scenes, (list, tuple)) or not 1 <= len(raw_scenes) <= 3:
|
||||
return None
|
||||
if any(not isinstance(scene, str) or scene not in SCENES for scene in raw_scenes):
|
||||
return None
|
||||
# Preserve classifier order while avoiding duplicate entries.
|
||||
return list(dict.fromkeys(raw_scenes))
|
||||
|
||||
def learn(
|
||||
self,
|
||||
scope: str | int,
|
||||
source: dict[str, Any],
|
||||
classification: dict[str, Any],
|
||||
) -> dict[str, Any] | None:
|
||||
"""Learn a sticker, or merge an existing ``file_unique_id``.
|
||||
|
||||
Returns a copy of the stored metadata. Invalid classifications, missing
|
||||
source data, unsupported formats, and missing new source files return
|
||||
``None``. Duplicate learning updates the index only and never copies
|
||||
either the object or its preview again.
|
||||
"""
|
||||
|
||||
scenes = self._valid_scenes(classification)
|
||||
if scenes is None or not isinstance(source, dict):
|
||||
return None
|
||||
|
||||
original_value = source.get("original_path")
|
||||
unique_value = source.get("file_unique_id")
|
||||
if original_value is None or unique_value is None:
|
||||
return None
|
||||
try:
|
||||
original = Path(original_value).expanduser()
|
||||
except TypeError:
|
||||
return None
|
||||
file_unique_id = str(unique_value).strip()
|
||||
if not file_unique_id:
|
||||
return None
|
||||
suffix = original.suffix.lower()
|
||||
if suffix not in ALLOWED_SUFFIXES:
|
||||
return None
|
||||
|
||||
description_value = classification.get("description", "")
|
||||
description = "" if description_value is None else str(description_value)
|
||||
|
||||
with self._lock:
|
||||
scope_dir, objects, previews = self._layout(scope, create=True)
|
||||
index = self._load_index(scope_dir)
|
||||
stickers: dict[str, dict[str, Any]] = index["stickers"]
|
||||
|
||||
for existing in stickers.values():
|
||||
if existing.get("file_unique_id") != file_unique_id:
|
||||
continue
|
||||
merged = list(dict.fromkeys([*existing.get("scenes", []), *scenes]))
|
||||
existing["scenes"] = merged[:3]
|
||||
if "description" in classification:
|
||||
existing["description"] = description
|
||||
self._write_index(scope_dir, index)
|
||||
return copy.deepcopy(existing)
|
||||
|
||||
if not original.is_file():
|
||||
return None
|
||||
|
||||
sticker_id = secrets.token_hex(16)
|
||||
while sticker_id in stickers:
|
||||
sticker_id = secrets.token_hex(16)
|
||||
object_path = objects / f"{sticker_id}{suffix}"
|
||||
|
||||
preview_source = _local_preview(source.get("url"))
|
||||
preview_path: Path | None = None
|
||||
if preview_source is not None and preview_source.is_file():
|
||||
preview_suffix = preview_source.suffix.lower()
|
||||
if not re.fullmatch(r"\.[a-z0-9]{1,10}", preview_suffix):
|
||||
preview_suffix = suffix
|
||||
preview_path = previews / f"{sticker_id}{preview_suffix}"
|
||||
|
||||
created_paths: list[Path] = []
|
||||
try:
|
||||
_copy_atomic(original, object_path)
|
||||
created_paths.append(object_path)
|
||||
if preview_path is not None:
|
||||
_copy_atomic(preview_source, preview_path)
|
||||
created_paths.append(preview_path)
|
||||
|
||||
record = {
|
||||
"id": sticker_id,
|
||||
"scenes": scenes,
|
||||
"description": description,
|
||||
"emoji": source.get("emoji"),
|
||||
"set_name": source.get("set_name"),
|
||||
"source_uid": source.get("source_uid"),
|
||||
"file_unique_id": file_unique_id,
|
||||
"format": suffix.removeprefix("."),
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
"path": str(object_path.resolve()),
|
||||
"preview_path": (
|
||||
str(preview_path.resolve()) if preview_path is not None else None
|
||||
),
|
||||
"url": source.get("url"),
|
||||
}
|
||||
stickers[sticker_id] = record
|
||||
self._write_index(scope_dir, index)
|
||||
except Exception:
|
||||
for created_path in created_paths:
|
||||
created_path.unlink(missing_ok=True)
|
||||
raise
|
||||
return copy.deepcopy(record)
|
||||
|
||||
def random_for_scene(self, scope: str | int, scene: str) -> dict[str, Any] | None:
|
||||
"""Return a uniformly selected random sticker belonging to ``scene``."""
|
||||
|
||||
with self._lock:
|
||||
scope_dir, _, _ = self._layout(scope, create=False)
|
||||
index = self._load_index(scope_dir)
|
||||
candidates = [
|
||||
record
|
||||
for record in index["stickers"].values()
|
||||
if scene in record.get("scenes", [])
|
||||
]
|
||||
return copy.deepcopy(random.choice(candidates)) if candidates else None
|
||||
|
||||
def get(self, scope: str | int, sticker_id: str) -> dict[str, Any] | None:
|
||||
"""Return one metadata record, or ``None`` when it does not exist."""
|
||||
|
||||
with self._lock:
|
||||
scope_dir, _, _ = self._layout(scope, create=False)
|
||||
record = self._load_index(scope_dir)["stickers"].get(sticker_id)
|
||||
return copy.deepcopy(record) if record is not None else None
|
||||
|
||||
def migrate_all_to_global(self, global_scope: str = "global") -> int:
|
||||
"""Merge legacy per-chat libraries into one shared library and archive them.
|
||||
|
||||
Existing IDs may change, which is safe because models select scenes rather
|
||||
than file IDs. ``file_unique_id`` deduplication makes retries idempotent.
|
||||
"""
|
||||
migrated = 0
|
||||
with self._lock:
|
||||
archive_root = self.root / ".migrated"
|
||||
for scope_dir in sorted(self.root.iterdir()):
|
||||
if (not scope_dir.is_dir() or scope_dir.name.startswith(".")
|
||||
or scope_dir.name == global_scope):
|
||||
continue
|
||||
index = self._load_index(scope_dir)
|
||||
complete = True
|
||||
for record in list(index["stickers"].values()):
|
||||
source = {
|
||||
"original_path": record.get("path"),
|
||||
"url": record.get("preview_path"),
|
||||
"file_unique_id": record.get("file_unique_id"),
|
||||
"emoji": record.get("emoji"),
|
||||
"set_name": record.get("set_name"),
|
||||
"source_uid": record.get("source_uid"),
|
||||
}
|
||||
classification = {
|
||||
"accept": True,
|
||||
"scenes": record.get("scenes") or ["other"],
|
||||
"description": record.get("description", ""),
|
||||
}
|
||||
if self.learn(global_scope, source, classification) is not None:
|
||||
migrated += 1
|
||||
else:
|
||||
complete = False
|
||||
if not complete:
|
||||
# Keep the legacy scope available for repair/retry. Any
|
||||
# records already copied are deduplicated by file_unique_id.
|
||||
continue
|
||||
archive_root.mkdir(exist_ok=True)
|
||||
destination = archive_root / scope_dir.name
|
||||
if not destination.exists():
|
||||
os.replace(scope_dir, destination)
|
||||
return migrated
|
||||
|
||||
def prompt_summary(self, scope: str | int) -> str:
|
||||
"""Summarize only scene names, counts, and descriptions for a model."""
|
||||
|
||||
with self._lock:
|
||||
scope_dir, _, _ = self._layout(scope, create=False)
|
||||
records = list(self._load_index(scope_dir)["stickers"].values())
|
||||
|
||||
if not records:
|
||||
return ""
|
||||
|
||||
lines = ["贴纸库存:"]
|
||||
for scene, scene_description in SCENES.items():
|
||||
matching = [record for record in records if scene in record.get("scenes", [])]
|
||||
if not matching:
|
||||
continue
|
||||
descriptions = list(
|
||||
dict.fromkeys(
|
||||
str(record.get("description", "")).strip()
|
||||
for record in matching
|
||||
if str(record.get("description", "")).strip()
|
||||
)
|
||||
)
|
||||
brief = ";".join(descriptions) if descriptions else "暂无"
|
||||
lines.append(
|
||||
f"- {scene}({scene_description}):{len(matching)} 个;简述:{brief}"
|
||||
)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
__all__ = ["ALLOWED_SUFFIXES", "SCENES", "StickerStore", "scene_prompt"]
|
||||
Reference in New Issue
Block a user