""" 通用数据加载模块 提供项目中常用的数据加载函数 """ import os import sys from typing import List from lib.utils import read_json def load_persona_data(persona_dir: str) -> dict: """加载人设数据 Args: persona_dir: 人设目录路径 Returns: 人设数据字典 Raises: SystemExit: 文件不存在时退出 """ persona_data_path = os.path.join(persona_dir, "人设.json") try: return read_json(persona_data_path) except FileNotFoundError: print(f"❌ 找不到人设数据文件: {persona_data_path}") print(f"请检查路径是否正确: {persona_dir}") sys.exit(1) def load_inspiration_list(persona_dir: str) -> List[str]: """加载灵感点列表 Args: persona_dir: 人设目录路径 Returns: 灵感点文本列表 Raises: SystemExit: 文件不存在或格式错误时退出 """ inspiration_list_path = os.path.join(persona_dir, "灵感点.json") try: inspiration_data = read_json(inspiration_list_path) if not isinstance(inspiration_data, list) or len(inspiration_data) == 0: print(f"❌ 灵感文件格式错误或为空: {inspiration_list_path}") sys.exit(1) return [item["灵感点"] for item in inspiration_data] except FileNotFoundError: print(f"❌ 找不到灵感文件: {inspiration_list_path}") print("请先运行 extract_inspirations.py 生成灵感点文件") sys.exit(1) def load_inspiration_data(persona_dir: str) -> List[dict]: """加载完整的灵感点数据(包含 meta 信息) Args: persona_dir: 人设目录路径 Returns: 灵感点数据列表,每项包含 {"灵感点": str, "meta": dict} Raises: SystemExit: 文件不存在或格式错误时退出 """ inspiration_list_path = os.path.join(persona_dir, "灵感点.json") try: inspiration_data = read_json(inspiration_list_path) if not isinstance(inspiration_data, list) or len(inspiration_data) == 0: print(f"❌ 灵感文件格式错误或为空: {inspiration_list_path}") sys.exit(1) return inspiration_data except FileNotFoundError: print(f"❌ 找不到灵感文件: {inspiration_list_path}") print("请先运行 extract_inspirations.py 生成灵感点文件") sys.exit(1) def select_inspiration(inspiration_arg: str, inspiration_list: List[str]) -> str: """根据参数选择灵感 Args: inspiration_arg: 灵感参数(数字索引或灵感名称) inspiration_list: 灵感点文本列表 Returns: 选中的灵感点文本 Raises: SystemExit: 选择失败时退出 """ try: # 尝试作为索引解析 inspiration_index = int(inspiration_arg) if 0 <= inspiration_index < len(inspiration_list): inspiration = inspiration_list[inspiration_index] print(f"使用灵感[{inspiration_index}]: {inspiration}") return inspiration else: print(f"❌ 灵感索引超出范围: {inspiration_index} (有效范围: 0-{len(inspiration_list)-1})") except ValueError: # 不是数字,当作灵感名称 if inspiration_arg in inspiration_list: print(f"使用灵感: {inspiration_arg}") return inspiration_arg else: print(f"❌ 找不到灵感: {inspiration_arg}") # 显示可用灵感列表后退出 print(f"可用灵感列表:") for i, insp in enumerate(inspiration_list[:10]): print(f" {i}: {insp}") if len(inspiration_list) > 10: print(f" ... 还有 {len(inspiration_list) - 10} 个") sys.exit(1) def load_step1_result(persona_dir: str, inspiration: str, model_name: str, scope: str = "all") -> dict: """加载 step1 匹配结果 Args: persona_dir: 人设目录路径 inspiration: 灵感点名称 model_name: 模型名称(如 "google/gemini-2.5-pro") scope: 范围标识("all" 或 "top10" 等) Returns: step1 结果字典 Raises: SystemExit: 文件不存在时退出 """ # 提取模型简称 model_name_short = model_name.replace("google/", "").replace("/", "_") # 构建文件路径 step1_file = os.path.join( persona_dir, "how", "灵感点", inspiration, f"{scope}_step1_灵感人设匹配_{model_name_short}.json" ) try: return read_json(step1_file) except FileNotFoundError: print(f"❌ 找不到 step1 结果文件: {step1_file}") print(f"请先运行 step1_inspiration_match.py 生成结果") sys.exit(1)