jewelry.py 11.5 KB
"""珠宝词库管理 + Prompt 组装。

DEFAULT_JEWELRY_LIBRARY  — 内置默认词库(9 类,启动期写到用户配置目录)
JewelryLibraryManager    — 加载 / 保存 / 增删 / 重置词库
PromptAssembler          — 表单数据拼成中文 prompt(处理空字段)
"""
import json
import logging
from pathlib import Path
from typing import Dict, List


DEFAULT_JEWELRY_LIBRARY = {
    "主石形状": [
        "圆形",
        "椭圆形",
        "梨形",
        "马眼形",
        "子弹形(Baguette 子弹刻面)",
        "垫形",
        "公主方形",
        "祖母绿形",
        "心形",
        "风筝形",
        "棺材形(Coffin Cut)",
        "菱形(Rhombus)",
        "正六边形",
        "四叶草形",
        "梯方形(Tapered Step)",
        "阿斯切形",
        "平底刻面风格"
    ],

    "主石材质": [
        "莫桑石",
        "钻石",
        "黑发晶",
        "蓝宝石",
        "红宝石",
        "粉蓝宝石",
        "绿碧玺",
        "黄水晶(天然包体)",
        "月光石",
        "摩根石",
        "海蓝宝",
        "天然白玉髓",
        "金绿宝石(猫眼)"
    ],

    "金属": [
        "14K黄金",
        "14K白金",
        "14K玫瑰金",
        "18K黄金",
        "18K白金",
        "18K玫瑰金",
        "双色金(白金+黄金)",
        "950铂金",
        "925银镀铑",
        "钛金属",
        "定制复古做旧金"
    ],

    "花头形式": [
        "全halo光环",
        "半halo光环",
        "双层halo",
        "花卉风格光环",
        "围圈雕刻光环",
        "围圈密钉镶",
        "经典圆形光环",
        "几何六边形光环",
        "非对称光环",
        "三石花头(cluster 结构)",
        "五石花头",
        "cluster堆砌花头(大小堆/不规则)",
        "双石结构(Two-stone)",
        "单石无光环(爪镶/包镶)",
        "花头侧面结构",
        "高耸花头(Cathedral halo)"
    ],

    "戒臂结构": [
        "直臂",
        "xox扭臂(交叉扭绞)",
        ">O< 戒臂结构",
        "<O> 戒臂结构",
        "V字戒臂",
        "交叉戒臂",
        "overlap重叠戒臂",
        "wave波浪戒臂",
        "刀锋臂",
        "大教堂戒臂(高肩设计)",
        "三股编织戒臂",
        "分裂戒臂(split shank)",
        "戒臂夹层",
        "小夹层戒臂(如莲花夹层设计)",
        "不对称戒臂"
    ],

    "戒臂处理": [
        "密钉镶戒臂",
        "微密钉镶",
        "镶石虎爪镶/逼镶",
        "抛光平滑戒臂",
        "珠边戒臂(milgrain)",
        "光金戒臂",
        "雕刻镂空花丝",
        "浮雕雕刻(凸雕)",
        "凹刻雕刻(内刻)",
        "几何雕刻纹理",
        "复古米粒边装饰(milgrain)",
        "编织纹理戒臂",
        "穿孔镂空细节",
        "锤纹处理"
    ],

    "辅石镶嵌": [
        "三石结构",
        "五石结构",
        "cluster自由堆砌侧石",
        "大小堆组合",
        "共爪镶侧钻",
        "包镶侧钻",
        "轨道镶",
        "槽镶"
    ],

    "特殊元素": [
        "花朵元素",
        "月亮元素",
        "星星元素",
        "日月星组合",
        "凯尔特结",
        "叶子图案",
        "自然植物藤蔓纹理",
        "蝴蝶结元素",
        "装饰艺术几何元素",
        "复古花纹",
        "哥特式结构元素"
    ]
}


class JewelryLibraryManager:
    """珠宝词库管理器"""

    def __init__(self, config_dir: Path):
        self.logger = logging.getLogger(__name__)
        self.config_dir = config_dir
        self.config_path = config_dir / "jewelry_library.json"
        self.library = self.load_library()

    def load_library(self) -> Dict[str, List[str]]:
        """加载词库,优先使用用户配置,不存在则使用默认词库"""
        if self.config_path.exists():
            try:
                with open(self.config_path, 'r', encoding='utf-8') as f:
                    library = json.load(f)

                # 验证数据完整性,补全缺失的类别
                needs_update = False
                for category, default_items in DEFAULT_JEWELRY_LIBRARY.items():
                    if category not in library:
                        self.logger.warning(f"检测到缺失类别: {category},从默认配置补全")
                        library[category] = default_items.copy()
                        needs_update = True
                    elif not library[category] or len(library[category]) == 0:
                        self.logger.warning(f"检测到空类别: {category},从默认配置补全")
                        library[category] = default_items.copy()
                        needs_update = True
                    elif len(library[category]) < len(default_items) * 0.5:
                        # 如果类别数据量少于默认值的50%,认为数据不完整,使用默认数据
                        self.logger.warning(
                            f"检测到类别 {category} 数据不完整(仅{len(library[category])}项,"
                            f"默认{len(default_items)}项),从默认配置补全"
                        )
                        library[category] = default_items.copy()
                        needs_update = True

                if needs_update:
                    self.logger.info("词库数据已补全,保存更新")
                    self.save_library(library)

                self.logger.info(f"珠宝词库加载成功: {self.config_path}")
                return library
            except Exception as e:
                self.logger.error(f"珠宝词库加载失败: {e},使用默认词库")
                library = DEFAULT_JEWELRY_LIBRARY.copy()
                try:
                    self.save_library(library)
                except Exception:
                    pass
                return library
        else:
            self.logger.info("首次使用,创建用户配置文件")
            library = DEFAULT_JEWELRY_LIBRARY.copy()
            self.save_library(library)
            return library

    def save_library(self, library: Dict[str, List[str]] = None):
        """保存词库到用户配置目录"""
        if library is None:
            library = self.library

        try:
            self.config_dir.mkdir(parents=True, exist_ok=True)
            with open(self.config_path, 'w', encoding='utf-8') as f:
                json.dump(library, f, ensure_ascii=False, indent=2)
            self.logger.info(f"珠宝词库保存成功: {self.config_path}")
        except Exception as e:
            self.logger.error(f"珠宝词库保存失败: {e}")
            raise

    def add_item(self, category: str, value: str):
        """添加词库项"""
        if category not in self.library:
            self.library[category] = []

        if value not in self.library[category]:
            self.library[category].append(value)
            self.save_library()
            self.logger.info(f"添加词库项: {category} -> {value}")
        else:
            self.logger.warning(f"词库项已存在: {category} -> {value}")

    def remove_item(self, category: str, value: str):
        """删除词库项"""
        if category in self.library and value in self.library[category]:
            self.library[category].remove(value)
            self.save_library()
            self.logger.info(f"删除词库项: {category} -> {value}")
        else:
            self.logger.warning(f"词库项不存在: {category} -> {value}")

    def reset_category(self, category: str):
        """恢复单个类别的默认词库"""
        if category in DEFAULT_JEWELRY_LIBRARY:
            self.library[category] = DEFAULT_JEWELRY_LIBRARY[category].copy()
            self.save_library()
            self.logger.info(f"恢复默认词库: {category}")
        else:
            self.logger.warning(f"未知的类别: {category}")

    def reset_all(self):
        """恢复所有类别的默认词库"""
        self.library = DEFAULT_JEWELRY_LIBRARY.copy()
        self.save_library()
        self.logger.info("恢复所有默认词库")


class PromptAssembler:
    """Prompt 组装器(纯中文)"""

    BASE_TEMPLATE = """一款高端精品珠宝戒指设计,主石为{主石},镶嵌于{金属}中。戒头采用{花头形式}围绕主石。戒臂采用{戒臂结构}设计,表面处理为{戒臂处理}。风格元素包括{特殊元素}。辅石采用{辅石镶嵌}增加光彩。高端珠宝渲染,干净的摄影棚光线,精准的金属抛光,强调宝石清晰度,不要出现手部。"""

    @staticmethod
    def assemble(form_data: Dict[str, str]) -> str:
        """智能组装 prompt,处理空值"""
        # 提取字段值(支持新旧两种格式)
        center_stone = form_data.get("主石", "").strip()  # 向后兼容

        # 新格式:拆分形状和材质
        center_stone_shape = form_data.get("主石形状", "").strip()
        center_stone_material = form_data.get("主石材质", "").strip()
        metal = form_data.get("金属", "").strip()
        halo_style = form_data.get("花头形式", "").strip()
        shank_structure = form_data.get("戒臂结构", "").strip()
        shank_treatment = form_data.get("戒臂处理", "").strip()
        special_motifs = form_data.get("特殊元素", "").strip()
        accent_setting = form_data.get("辅石镶嵌", "").strip()

        parts = []

        # 构建主石描述(优先使用新格式)
        if center_stone_shape and center_stone_material:
            main_stone_desc = f"{center_stone_shape}{center_stone_material}"
        elif center_stone:
            main_stone_desc = center_stone
        elif center_stone_shape:
            main_stone_desc = center_stone_shape
        elif center_stone_material:
            main_stone_desc = center_stone_material
        else:
            main_stone_desc = ""

        # 主体部分(主石和金属)
        if main_stone_desc and metal:
            parts.append(f"一款高端精品珠宝戒指设计,主石为{main_stone_desc},镶嵌于{metal}中。")
        elif metal:
            parts.append(f"一款高端精品珠宝戒指设计,镶嵌于{metal}中。")
        elif main_stone_desc:
            parts.append(f"一款高端精品珠宝戒指设计,主石为{main_stone_desc}。")
        else:
            parts.append("一款高端精品珠宝戒指设计。")

        # 戒头部分
        if halo_style:
            if main_stone_desc:
                parts.append(f"戒头采用{halo_style}围绕主石。")
            else:
                parts.append(f"戒头采用{halo_style}设计。")

        # 戒臂部分
        if shank_structure and shank_treatment:
            parts.append(f"戒臂采用{shank_structure}设计,表面处理为{shank_treatment}。")
        elif shank_structure:
            parts.append(f"戒臂采用{shank_structure}设计。")
        elif shank_treatment:
            parts.append(f"戒臂表面处理为{shank_treatment}。")

        # 特殊元素
        if special_motifs:
            parts.append(f"风格元素包括{special_motifs}。")

        # 辅石镶嵌
        if accent_setting:
            parts.append(f"辅石采用{accent_setting}增加光彩。")

        # 固定的渲染要求
        parts.append("高端珠宝渲染,干净的摄影棚光线,精准的金属抛光,强调宝石清晰度,不要出现手部。")

        return "".join(parts)