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,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