RRRosmontis
34089d0042
Include utility modules, shared channels, prompts, read-only skills, wardrobe memory and regression tests. Externalize legacy credentials and exclude private runtime state.
369 lines
14 KiB
Python
369 lines
14 KiB
Python
"""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"]
|