import asyncio import os import sys import yaml from pathlib import Path from sqlalchemy import select from sqlalchemy.dialects.mysql import insert project_root = Path(__file__).parent.parent sys.path.insert(0, str(project_root)) from database.connection import get_db_manager from database.models import Role, RoleLanguage from config import settings from utils.logger import session_logger class RoleImporter: def __init__(self): self.db_manager = None self.roles_dir = Path(settings.assets_dir) / "roles_definitions" async def initialize(self): """初始化数据库连接""" self.db_manager = await get_db_manager() await self.db_manager.initialize() async def load_yaml_files(self): """扫描并加载所有YAML角色定义文件""" if not self.roles_dir.exists(): print(f"角色定义目录不存在: {self.roles_dir}") return [] yaml_files = list(self.roles_dir.glob("*.yaml")) + list(self.roles_dir.glob("*.yml")) roles_data = [] for yaml_file in yaml_files: try: with open(yaml_file, 'r', encoding='utf-8') as f: data = yaml.safe_load(f) if data: role_key = yaml_file.stem data['role_key'] = role_key data['source_file'] = str(yaml_file) roles_data.append(data) print(f"加载角色配置: {role_key} <- {yaml_file.name}") except Exception as e: print(f"加载YAML文件失败 {yaml_file}: {e}") return roles_data def validate_role_config(self, role_data): """验证角色配置的完整性""" errors = [] required_fields = ['name', 'role_key'] for field in required_fields: if field not in role_data: errors.append(f"缺少必需字段: {field}") if 'multilingual' in role_data: for lang_code, lang_config in role_data['multilingual'].items(): if not isinstance(lang_config, dict): errors.append(f"语言配置 {lang_code} 必须是字典格式") continue if 'content' not in lang_config: errors.append(f"语言 {lang_code} 缺少content字段") return errors def extract_role_data(self, role_config): """从配置中提取主角色数据""" return { 'role_key': role_config['role_key'], 'name': role_config.get('name', ''), 'description': role_config.get('description', ''), 'content': role_config.get('content', ''), 'default_language': role_config.get('default_language'), # 'asr_provider': role_config.get('asr_provider'), # 'llm_provider': role_config.get('llm_provider'), # 'tts_provider': role_config.get('tts_provider'), # 'aws_language_code': role_config.get('aws_language_code'), 'volcano_model_id': role_config.get('volcano_model_id'), # 'volcano_voice_type': role_config.get('volcano_voice_type'), # 'tencent_voice_type': role_config.get('tencent_voice_type'), # 'aliyun_voice_name': role_config.get('aliyun_voice_name'), 'minimax_voice_id': role_config.get('minimax_voice_id'), 'url': role_config.get('url'), 'homophones': role_config.get('homophones'), 'enabled': True } def extract_language_data(self, role_id, lang_code, lang_config): """从配置中提取语言特定数据""" return { 'role_id': role_id, 'language_code': lang_code, 'name': lang_config.get('name'), 'content': lang_config.get('content'), # 'asr_provider': lang_config.get('asr_provider'), # 'llm_provider': lang_config.get('llm_provider'), # 'tts_provider': lang_config.get('tts_provider'), # 'aws_language_code': lang_config.get('aws_language_code'), # 'volcano_voice_type': lang_config.get('volcano_voice_type'), # 'tencent_voice_type': lang_config.get('tencent_voice_type'), # 'aliyun_voice_name': lang_config.get('aliyun_voice_name'), 'minimax_voice_id': lang_config.get('minimax_voice_id'), 'url': lang_config.get('url') } async def import_role(self, role_config, update_existing=False): """导入单个角色到数据库""" session = await self.db_manager.get_session() try: errors = self.validate_role_config(role_config) if errors: print(f"角色 {role_config.get('role_key', 'unknown')} 验证失败:") for error in errors: print(f" - {error}") return False role_key = role_config['role_key'] existing_role = await session.execute( select(Role).where(Role.role_key == role_key) ) existing_role = existing_role.scalar_one_or_none() if existing_role and not update_existing: print(f"角色 {role_key} 已存在，跳过导入（使用 --update 强制更新）") return True role_data = self.extract_role_data(role_config) if existing_role: for key, value in role_data.items(): if key != 'role_key': # 不更新主键 setattr(existing_role, key, value) role_id = existing_role.id print(f"更新角色: {role_key}") else: stmt = insert(Role).values(**role_data) result = await session.execute(stmt) role_id = result.lastrowid print(f"创建角色: {role_key}") if 'multilingual' in role_config: if existing_role: await session.execute( RoleLanguage.__table__.delete().where( RoleLanguage.role_id == role_id ) ) for lang_code, lang_config in role_config['multilingual'].items(): lang_data = self.extract_language_data(role_id, lang_code, lang_config) lang_stmt = insert(RoleLanguage).values(**lang_data) await session.execute(lang_stmt) print(f" 添加语言配置: {lang_code}") await session.commit() print(f"✓ 角色 {role_key} 导入成功") return True except Exception as e: await session.rollback() print(f"✗ 导入角色 {role_config.get('role_key', 'unknown')} 失败: {e}") return False finally: await session.close() async def import_all_roles(self, update_existing=False): """导入所有角色配置""" print("开始导入角色配置到数据库...") print(f"角色定义目录: {self.roles_dir}") roles_data = await self.load_yaml_files() if not roles_data: print("未找到任何角色配置文件") return print(f"找到 {len(roles_data)} 个角色配置文件") print("-" * 50) success_count = 0 failed_count = 0 for role_config in roles_data: success = await self.import_role(role_config, update_existing) if success: success_count += 1 else: failed_count += 1 print() # 空行分隔 print("-" * 50) print(f"导入完成:") print(f" 成功: {success_count}") print(f" 失败: {failed_count}") print(f" 总计: {len(roles_data)}") async def list_roles(self): """列出数据库中的所有角色""" session = await self.db_manager.get_session() try: result = await session.execute( select(Role.role_key, Role.name, Role.enabled) .order_by(Role.role_key) ) roles = result.fetchall() if not roles: print("数据库中没有角色配置") return print(f"数据库中的角色配置 (共 {len(roles)} 个):") print("-" * 60) print(f"{'角色Key':<20} {'角色名称':<25} {'状态':<10}") print("-" * 60) for role in roles: status = "启用" if role.enabled else "禁用" print(f"{role.role_key:<20} {role.name:<25} {status:<10}") except Exception as e: print(f"查询角色列表失败: {e}") finally: await session.close() async def delete_role(self, role_key): """删除指定角色""" session = await self.db_manager.get_session() try: result = await session.execute( select(Role).where(Role.role_key == role_key) ) role = result.scalar_one_or_none() if not role: print(f"角色 {role_key} 不存在") return False await session.delete(role) await session.commit() print(f"✓ 角色 {role_key} 删除成功") return True except Exception as e: await session.rollback() print(f"✗ 删除角色 {role_key} 失败: {e}") return False finally: await session.close() async def main(): """主函数""" import argparse parser = argparse.ArgumentParser(description="角色配置导入工具") parser.add_argument('--update', action='store_true', help='更新已存在的角色') parser.add_argument('--list', action='store_true', help='列出数据库中的角色') parser.add_argument('--delete', type=str, help='删除指定的角色') parser.add_argument('--role', type=str, help='只导入指定的角色文件') args = parser.parse_args() importer = RoleImporter() try: await importer.initialize() if args.list: await importer.list_roles() elif args.delete: await importer.delete_role(args.delete) elif args.role: role_file = importer.roles_dir / f"{args.role}.yaml" if not role_file.exists(): role_file = importer.roles_dir / f"{args.role}.yml" if not role_file.exists(): print(f"角色配置文件不存在: {args.role}") return try: with open(role_file, 'r', encoding='utf-8') as f: role_config = yaml.safe_load(f) role_config['role_key'] = args.role role_config['source_file'] = str(role_file) await importer.import_role(role_config, args.update) except Exception as e: print(f"导入角色 {args.role} 失败: {e}") else: await importer.import_all_roles(args.update) except Exception as e: print(f"操作失败: {e}") finally: if importer.db_manager: await importer.db_manager.close() if __name__ == "__main__": asyncio.run(main())