121 lines
4.4 KiB
Python
121 lines
4.4 KiB
Python
"""酒馆角色卡解析器。
|
||
|
||
从 PNG（``tEXt`` chunk 的 ``chara``/``ccv3``）/ JSON / CHARX（zip 含 ``card.json``）
|
||
解析角色卡，归一化为统一的 ``CharacterCard``（V2/V3 的 ``data`` 字段拍平）。
|
||
|
||
纯字节 / 文本解析，无文件 IO——调用方（service 层）负责读文件后传入字节。
|
||
V1 扁平结构自动归一化；``character_book`` 原样保留，由映射层转世谱条目。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import io
|
||
import json
|
||
import zipfile
|
||
from dataclasses import dataclass, field
|
||
|
||
_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
|
||
_ZIP_SIGNATURE = b"PK\x03\x04"
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class CharacterCard:
|
||
"""归一化后的角色卡（V1/V2/V3 统一为拍平字段）。"""
|
||
|
||
spec_version: str # "1" / "2.0" / "3.0"
|
||
name: str = ""
|
||
description: str = ""
|
||
personality: str = ""
|
||
scenario: str = ""
|
||
first_mes: str = ""
|
||
mes_example: str = ""
|
||
creator_notes: str = ""
|
||
system_prompt: str = ""
|
||
post_history_instructions: str = ""
|
||
alternate_greetings: list[str] = field(default_factory=list)
|
||
tags: list[str] = field(default_factory=list)
|
||
creator: str = ""
|
||
character_version: str = ""
|
||
character_book: dict | None = None # 原始 lorebook，由映射层转 LorebookBook
|
||
nickname: str | None = None # V3
|
||
extensions: dict = field(default_factory=dict) # 保留未知键（酒馆规范：不得删除）
|
||
|
||
|
||
def parse_character_card(raw: bytes) -> CharacterCard:
|
||
"""从原始字节解析角色卡（自动识别 PNG / JSON / CHARX）。"""
|
||
if raw[:8] == _PNG_SIGNATURE:
|
||
data = _parse_png(raw)
|
||
elif raw[:4] == _ZIP_SIGNATURE:
|
||
data = _parse_charx(raw)
|
||
else:
|
||
data = json.loads(raw.decode("utf-8-sig"))
|
||
return _normalize(data)
|
||
|
||
|
||
def _parse_png(raw: bytes) -> dict:
|
||
"""从 PNG tEXt chunk 提取 chara / ccv3（base64 → JSON）。优先 ccv3。"""
|
||
chunks: list[tuple[str, str]] = []
|
||
offset = len(_PNG_SIGNATURE)
|
||
while offset + 8 <= len(raw):
|
||
length = int.from_bytes(raw[offset : offset + 4], "big")
|
||
chunk_type = raw[offset + 4 : offset + 8]
|
||
data_start = offset + 8
|
||
data_end = data_start + length
|
||
if data_end > len(raw):
|
||
break
|
||
if chunk_type == b"tEXt":
|
||
null = raw[data_start:data_end].find(b"\x00")
|
||
if null >= 0:
|
||
keyword = raw[data_start : data_start + null].decode("latin-1")
|
||
text = raw[data_start + null + 1 : data_end].decode("latin-1")
|
||
chunks.append((keyword, text))
|
||
offset = data_end + 4 # 跳过 CRC
|
||
|
||
for want in ("ccv3", "chara"):
|
||
for keyword, text in chunks:
|
||
if keyword == want:
|
||
decoded = base64.b64decode(text)
|
||
return json.loads(decoded.decode("utf-8"))
|
||
raise ValueError("PNG 中未找到角色卡元数据（chara / ccv3 chunk）")
|
||
|
||
|
||
def _parse_charx(raw: bytes) -> dict:
|
||
"""从 CHARX（zip）提取 card.json。"""
|
||
with zipfile.ZipFile(io.BytesIO(raw)) as zf:
|
||
card_name = next((n for n in zf.namelist() if n.endswith("card.json")), None)
|
||
if card_name is None:
|
||
raise ValueError("CHARX 中未找到 card.json")
|
||
with zf.open(card_name) as f:
|
||
return json.loads(f.read().decode("utf-8"))
|
||
|
||
|
||
def _normalize(data: dict) -> CharacterCard:
|
||
"""V1 扁平结构归一化为 V2/V3（data 包裹）。"""
|
||
if isinstance(data.get("data"), dict):
|
||
inner = data["data"]
|
||
spec_version = str(data.get("spec_version", "2.0"))
|
||
else:
|
||
inner = data
|
||
spec_version = "1"
|
||
|
||
return CharacterCard(
|
||
spec_version=spec_version,
|
||
name=inner.get("name", ""),
|
||
description=inner.get("description", ""),
|
||
personality=inner.get("personality", ""),
|
||
scenario=inner.get("scenario", ""),
|
||
first_mes=inner.get("first_mes", ""),
|
||
mes_example=inner.get("mes_example", ""),
|
||
creator_notes=inner.get("creator_notes", ""),
|
||
system_prompt=inner.get("system_prompt", ""),
|
||
post_history_instructions=inner.get("post_history_instructions", ""),
|
||
alternate_greetings=list(inner.get("alternate_greetings", []) or []),
|
||
tags=list(inner.get("tags", []) or []),
|
||
creator=inner.get("creator", ""),
|
||
character_version=inner.get("character_version", ""),
|
||
character_book=inner.get("character_book"),
|
||
nickname=inner.get("nickname"),
|
||
extensions=dict(inner.get("extensions", {}) or {}),
|
||
)
|