diff --git a/AITrain/main_controller.py b/AITrain/main_controller.py new file mode 100644 index 0000000..cdf3117 --- /dev/null +++ b/AITrain/main_controller.py @@ -0,0 +1,337 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +''' +双AI角色对话系统主控制程序 +完整的工作流程：PDF处理 -> 角色加载 -> RAG对话 -> 历史记录 +''' + +import os +import sys +import shutil +from typing import List, Dict +import json + +def check_dependencies(): + """检查依赖库""" + missing_deps = [] + + try: + import PyPDF2 + except ImportError: + missing_deps.append("PyPDF2") + + try: + import pymupdf + print("✓ pymupdf 可用") + except ImportError: + print("⚠ pymupdf 不可用，将使用 PyPDF2") + + try: + import sentence_transformers + import faiss + print("✓ 向量化功能可用") + except ImportError: + print("⚠ 向量化功能不可用，将使用文本匹配") + + if missing_deps: + print(f"✗ 缺少依赖库: {', '.join(missing_deps)}") + print("请运行: pip install PyPDF2 sentence-transformers faiss-cpu") + return False + + return True + +def setup_directories(): + """设置项目目录结构""" + directories = [ + "./knowledge_base", + "./characters", + "./worldview", + "./rag_knowledge", + "./conversation_data" + ] + + for dir_path in directories: + os.makedirs(dir_path, exist_ok=True) + print(f"✓ 目录就绪: {dir_path}") + +def copy_demo_files(): + """复制演示文档到知识库目录""" + file_mappings = [ + ("./worldview/worldview_template_coc.json", "./knowledge_base/worldview_template_coc.json"), + ("./characters/character_template_detective.json", "./knowledge_base/character_template_detective.json"), + ("./characters/character_template_professor.json", "./knowledge_base/character_template_professor.json") + ] + + for source, target in file_mappings: + if os.path.exists(source): + shutil.copy2(source, target) + print(f"✓ 复制文档: {os.path.basename(target)}") + +def process_pdf_workflow(): + """PDF处理工作流""" + print("\n" + "="*60) + print("PDF世界观文档处理") + print("="*60) + + from pdf_to_rag_processor import PDFToRAGProcessor + + pdf_path = input("请输入PDF文件路径 (例: ./coc.pdf): ").strip() + + if not os.path.exists(pdf_path): + print(f"✗ 文件不存在: {pdf_path}") + return False + + try: + processor = PDFToRAGProcessor() + result = processor.process_pdf_to_rag(pdf_path, "./rag_knowledge") + + print(f"\n✓ PDF处理完成!") + print(f" - 文档块数: {result['chunks_count']}") + print(f" - 概念数: {result['concepts_count']}") + print(f" - 向量索引: {'启用' if result['vector_enabled'] else '未启用'}") + + return True + + except Exception as e: + print(f"✗ PDF处理失败: {e}") + return False + +def show_character_info(): + """显示角色信息""" + print("\n" + "="*60) + print("角色设定信息") + print("="*60) + + knowledge_dir = "./knowledge_base" + character_files = [f for f in os.listdir(knowledge_dir) if f.startswith('character') and f.endswith('.json')] + + for char_file in character_files: + try: + with open(os.path.join(knowledge_dir, char_file), 'r', encoding='utf-8') as f: + char_data = json.load(f) + + name = char_data.get('character_name', '未知') + occupation = char_data.get('basic_info', {}).get('occupation', '未知') + traits = char_data.get('personality', {}).get('core_traits', []) + + print(f"\n角色: {name}") + print(f" 职业: {occupation}") + print(f" 特点: {', '.join(traits[:3])}") + + except Exception as e: + print(f"✗ 读取角色文件失败: {char_file} - {e}") + +def run_dialogue_system(): + """运行对话系统""" + print("\n" + "="*60) + print("启动双AI角色对话系统") + print("="*60) + + try: + from dual_ai_dialogue_system import main as dialogue_main + dialogue_main() + except Exception as e: + print(f"✗ 对话系统启动失败: {e}") + import traceback + traceback.print_exc() + +def create_demo_scenario(): + """创建演示场景""" + print("\n创建演示对话场景...") + + try: + from dual_ai_dialogue_system import RAGKnowledgeBase, ConversationManager, DualAIDialogueEngine + from npc_dialogue_generator import NPCDialogueGenerator + + # 初始化组件 + kb = RAGKnowledgeBase("./knowledge_base") + conv_mgr = ConversationManager("./conversation_data/demo_conversations.db") + + # 检查模型路径 + base_model_path = '/mnt/g/Project02/AITrain/Qwen/Qwen3-4B' + lora_model_path = './output/NPC_Dialogue_LoRA/final_model' + + if not os.path.exists(base_model_path): + print(f"✗ 基础模型路径不存在: {base_model_path}") + print("请修改 main_controller.py 中的模型路径") + return + + if not os.path.exists(lora_model_path): + lora_model_path = None + print("⚠ LoRA模型不存在，使用基础模型") + + llm_generator = NPCDialogueGenerator(base_model_path, lora_model_path, kb.character_data) + dialogue_engine = DualAIDialogueEngine(kb, conv_mgr, llm_generator) + + # 创建演示对话 + characters = ["维多利亚·布莱克伍德", "阿奇博尔德·韦恩"] + worldview = "克苏鲁的呼唤" + + session_id = conv_mgr.create_session(characters, worldview) + print(f"✓ 创建演示会话: {session_id}") + + # 运行几轮对话 + topic = "最近发生的神秘事件" + print(f"\n开始演示对话 - 主题: {topic}") + print("-" * 40) + + # 演示不同的历史上下文设置 + # print("演示1: 使用默认上下文设置（历史3轮，信息2个）") + # dialogue_engine.run_conversation_turn(session_id, characters, 6, topic) + + + session_id = conv_mgr.create_session(characters, worldview) + print(f"✓ 创建演示会话: {session_id}") + print("\n演示3: 使用最少历史上下文（历史1轮，信息1个）") + dialogue_engine.run_conversation_turn(session_id, characters, 6, topic, 1, 10) + + session_id = conv_mgr.create_session(characters, worldview) + print(f"✓ 创建演示会话: {session_id}") + print("\n演示2: 使用更多历史上下文（历史10轮，信息10个）") + dialogue_engine.run_conversation_turn(session_id, characters, 6, topic, 5, 10) + + print(f"\n✓ 演示完成！会话ID: {session_id}") + print("你可以通过主对话系统继续这个对话") + + except Exception as e: + print(f"✗ 演示场景创建失败: {e}") + import traceback + traceback.print_exc() + +def show_system_status(): + """显示系统状态""" + print("\n" + "="*60) + print("系统状态检查") + print("="*60) + + # 检查文件 + files_to_check = [ + ("./knowledge_base/worldview_template_coc.json", "世界观模板"), + ("./knowledge_base/character_template_detective.json", "侦探角色"), + ("./knowledge_base/character_template_professor.json", "教授角色"), + ("./pdf_to_rag_processor.py", "PDF处理器"), + ("./dual_ai_dialogue_system.py", "对话系统"), + ("./npc_dialogue_generator.py", "NPC生成器") + ] + + print("\n文件检查:") + for file_path, description in files_to_check: + if os.path.exists(file_path): + print(f"✓ {description}: {file_path}") + else: + print(f"✗ {description}: {file_path} (不存在)") + + # 检查目录 + print("\n目录检查:") + directories = ["./knowledge_base", "./rag_knowledge", "./conversation_data"] + for dir_path in directories: + if os.path.exists(dir_path): + file_count = len([f for f in os.listdir(dir_path) if os.path.isfile(os.path.join(dir_path, f))]) + print(f"✓ {dir_path}: {file_count} 个文件") + else: + print(f"✗ {dir_path}: 不存在") + + # 检查对话会话 + try: + from dual_ai_dialogue_system import ConversationManager + conv_mgr = ConversationManager("./conversation_data/conversations.db") + sessions = conv_mgr.list_sessions() + print(f"\n✓ 对话会话: {len(sessions)} 个") + except Exception as e: + print(f"\n✗ 对话会话检查失败: {e}") + +def main(): + """主控制程序""" + print("="*70) + print(" 双AI角色对话系统 - 主控制程序") + print(" 基于RAG的世界观增强对话引擎") + print("="*70) + + # 检查依赖 + if not check_dependencies(): + return + + # 设置目录 + # setup_directories() + # copy_demo_files() + + while True: + print("\n" + "="*50) + print("主菜单 - 请选择操作:") + print("1. 处理PDF世界观文档 (转换为RAG格式)") + print("2. 查看角色设定信息") + print("3. 启动双AI对话系统") + print("4. 创建演示对话场景") + print("5. 系统状态检查") + print("6. 查看使用说明") + print("0. 退出") + print("="*50) + + choice = input("请输入选择 (0-6): ").strip() + + if choice == '0': + print("\n感谢使用双AI角色对话系统！") + break + + elif choice == '1': + process_pdf_workflow() + + elif choice == '2': + show_character_info() + + elif choice == '3': + run_dialogue_system() + + elif choice == '4': + create_demo_scenario() + + elif choice == '5': + show_system_status() + + elif choice == '6': + show_usage_guide() + + else: + print("❌ 无效选择，请重新输入") + +def show_usage_guide(): + """显示使用说明""" + print("\n" + "="*60) + print("系统使用说明") + print("="*60) + + guide = """ +🚀 快速开始: +1. 首次使用建议先运行"创建演示对话场景" +2. 如有PDF世界观文档，选择"处理PDF世界观文档" +3. 通过"启动双AI对话系统"开始角色对话 + +📁 文档格式说明: +- 世界观文档: worldview_template_coc.json (参考COC设定) +- 角色设定: character_template_*.json (包含详细人设) + +🔧 系统功能: +- PDF自动转换为RAG知识库 +- 基于向量相似度的上下文检索 +- 持久化对话历史存储 +- 角色设定一致性保持 + +📝 自定义角色: +1. 参考 character_template_*.json 格式 +2. 保存到 knowledge_base/ 目录 +3. 重启对话系统加载新角色 + +💾 对话数据: +- 历史对话保存在 conversation_data/ 目录 +- 支持会话恢复和历史查看 +- 自动记录使用的上下文信息 + +⚠️ 注意事项: +- 确保模型路径正确设置 +- 首次运行需要下载向量化模型 +- PDF处理需要足够内存 +""" + print(guide) + +if __name__ == '__main__': + main() \ No newline at end of file