313 lines
12 KiB
Python
313 lines
12 KiB
Python
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())
|