调整部分结构

This commit is contained in:
LuiKlee
2025-10-05 18:42:28 +08:00
committed by GitHub
parent fbe35d1d01
commit 34495c07cd

675
bot.py
View File

@@ -1,315 +1,504 @@
# import asyncio
import asyncio
import os
import platform
import sys
import time
import platform
import traceback
from pathlib import Path
from contextlib import asynccontextmanager
import hashlib
from typing import Optional, Dict, Any
from colorama import Fore, init
from dotenv import load_dotenv # 处理.env文件
# 初始化基础工具
from colorama import init, Fore
from dotenv import load_dotenv
from rich.traceback import install
# maim_message imports for console input
# 最早期初始化日志系统,确保所有后续模块都使用正确的日志格式
from src.common.logger import get_logger, initialize_logging, shutdown_logging
# 初始化日志系统
from src.common.logger import initialize_logging, get_logger, shutdown_logging
# UI日志适配器
# 初始化日志和错误显示
initialize_logging()
from src.main import MainSystem # noqa
from src import BaseMain
from src.manager.async_task_manager import async_task_manager
from src.chat.knowledge.knowledge_lib import initialize_lpmm_knowledge
from src.config.config import global_config
from src.common.database.database import initialize_sql_database
from src.common.database.sqlalchemy_models import initialize_database as init_db
logger = get_logger("main")
install(extra_lines=3)
# 常量定义
SUPPORTED_DATABASES = ['sqlite', 'mysql', 'postgresql']
SHUTDOWN_TIMEOUT = 10.0
EULA_CHECK_INTERVAL = 2
MAX_EULA_CHECK_ATTEMPTS = 30
MAX_ENV_FILE_SIZE = 1024 * 1024 # 1MB限制
# 设置工作目录为脚本所在目录
script_dir = os.path.dirname(os.path.abspath(__file__))
os.chdir(script_dir)
logger.info(f"已设置工作目录为: {script_dir}")
logger.info("工作目录已设置")
class ConfigManager:
"""配置管理器"""
@staticmethod
def ensure_env_file():
"""确保.env文件存在如果不存在则从模板创建"""
env_file = Path(".env")
template_env = Path("template/template.env")
if not env_file.exists():
if template_env.exists():
logger.info("未找到.env文件正在从模板创建...")
try:
env_file.write_text(template_env.read_text(encoding='utf-8'), encoding='utf-8')
logger.info("已从template/template.env创建.env文件")
logger.warning("请编辑.env文件将EULA_CONFIRMED设置为true并配置其他必要参数")
except Exception as e:
logger.error(f"创建.env文件失败: {e}")
sys.exit(1)
else:
logger.error("未找到.env文件和template.env模板文件")
sys.exit(1)
# 检查并创建.env文件
def ensure_env_file():
"""确保.env文件存在,如果不存在则从模板创建"""
env_file = Path(".env")
template_env = Path("template/template.env")
@staticmethod
def verify_env_file_integrity():
"""验证.env文件完整性"""
env_file = Path(".env")
if not env_file.exists():
return False
# 检查文件大小
file_size = env_file.stat().st_size
if file_size == 0 or file_size > MAX_ENV_FILE_SIZE:
logger.error(f".env文件大小异常: {file_size}字节")
return False
# 检查文件内容是否包含必要字段
try:
content = env_file.read_text(encoding='utf-8')
if 'EULA_CONFIRMED' not in content:
logger.error(".env文件缺少EULA_CONFIRMED字段")
return False
except Exception as e:
logger.error(f"读取.env文件失败: {e}")
return False
return True
if not env_file.exists():
if template_env.exists():
logger.info("未找到.env文件正在从模板创建...")
import shutil
@staticmethod
def safe_load_dotenv():
"""安全加载环境变量"""
try:
if not ConfigManager.verify_env_file_integrity():
logger.error(".env文件完整性验证失败")
return False
load_dotenv()
logger.info("环境变量加载成功")
return True
except Exception as e:
logger.error(f"加载环境变量失败: {e}")
return False
shutil.copy(template_env, env_file)
logger.info("已从template/template.env创建.env文件")
logger.warning("请编辑.env文件将EULA_CONFIRMED设置为true并配置其他必要参数")
else:
logger.error("未找到.env文件和template.env模板文件")
class EULAManager:
"""EULA管理类"""
@staticmethod
async def check_eula():
"""检查EULA和隐私条款确认状态"""
confirm_logger = get_logger("confirm")
if not ConfigManager.safe_load_dotenv():
confirm_logger.error("无法加载环境变量EULA检查失败")
sys.exit(1)
eula_confirmed = os.getenv('EULA_CONFIRMED', '').lower()
if eula_confirmed == 'true':
logger.info("EULA已通过环境变量确认")
return
# 提示用户确认EULA
confirm_logger.critical("您需要同意EULA和隐私条款才能使用MoFox_Bot")
confirm_logger.critical("请阅读以下文件:")
confirm_logger.critical(" - EULA.md (用户许可协议)")
confirm_logger.critical(" - PRIVACY.md (隐私条款)")
confirm_logger.critical("然后编辑 .env 文件,将 'EULA_CONFIRMED=false' 改为 'EULA_CONFIRMED=true'")
attempts = 0
while attempts < MAX_EULA_CHECK_ATTEMPTS:
try:
await asyncio.sleep(EULA_CHECK_INTERVAL)
attempts += 1
# 重新加载环境变量
ConfigManager.safe_load_dotenv()
eula_confirmed = os.getenv('EULA_CONFIRMED', '').lower()
if eula_confirmed == 'true':
confirm_logger.info("EULA确认成功感谢您的同意")
return
if attempts % 5 == 0:
confirm_logger.critical(f"请修改 .env 文件中的 EULA_CONFIRMED=true (尝试 {attempts}/{MAX_EULA_CHECK_ATTEMPTS})")
except KeyboardInterrupt:
confirm_logger.info("用户取消,程序退出")
sys.exit(0)
except Exception as e:
confirm_logger.error(f"检查EULA状态失败: {e}")
if attempts >= MAX_EULA_CHECK_ATTEMPTS:
confirm_logger.error("达到最大检查次数,程序退出")
sys.exit(1)
confirm_logger.error("EULA确认超时程序退出")
sys.exit(1)
class TaskManager:
"""任务管理器"""
@staticmethod
async def cancel_pending_tasks(loop, timeout=SHUTDOWN_TIMEOUT):
"""取消所有待处理的任务"""
remaining_tasks = [
t for t in asyncio.all_tasks(loop)
if t is not asyncio.current_task(loop) and not t.done()
]
if not remaining_tasks:
logger.info("没有待取消的任务")
return True
logger.info(f"正在取消 {len(remaining_tasks)} 个剩余任务...")
# 取消任务
for task in remaining_tasks:
task.cancel()
# 等待任务完成
try:
results = await asyncio.wait_for(
asyncio.gather(*remaining_tasks, return_exceptions=True),
timeout=timeout
)
# 检查任务结果
for i, result in enumerate(results):
if isinstance(result, Exception):
logger.warning(f"任务 {i} 取消时发生异常: {result}")
logger.info("所有剩余任务已成功取消")
return True
except asyncio.TimeoutError:
logger.warning("等待任务取消超时,强制继续关闭")
return False
except Exception as e:
logger.error(f"等待任务取消时发生异常: {e}")
return False
# 确保环境文件存在
ensure_env_file()
@staticmethod
async def stop_async_tasks():
"""停止所有异步任务"""
try:
from src.manager.async_task_manager import async_task_manager
await async_task_manager.stop_and_wait_all_tasks()
return True
except ImportError:
logger.warning("异步任务管理器不可用,跳过任务停止")
return False
except Exception as e:
logger.error(f"停止异步任务失败: {e}")
return False
# 加载环境变量
load_dotenv()
class ShutdownManager:
"""关闭管理器"""
@staticmethod
async def graceful_shutdown(loop=None):
"""优雅关闭程序"""
try:
logger.info("正在优雅关闭麦麦...")
start_time = time.time()
# 停止异步任务
tasks_stopped = await TaskManager.stop_async_tasks()
# 取消待处理任务
tasks_cancelled = True
if loop and not loop.is_closed():
tasks_cancelled = await TaskManager.cancel_pending_tasks(loop)
shutdown_time = time.time() - start_time
success = tasks_stopped and tasks_cancelled
if success:
logger.info(f"麦麦优雅关闭完成,耗时: {shutdown_time:.2f}")
else:
logger.warning(f"麦麦关闭完成,但部分操作未成功,耗时: {shutdown_time:.2f}")
return success
confirm_logger = get_logger("confirm")
# 获取没有加载env时的环境变量
except Exception as e:
logger.error(f"麦麦关闭失败: {e}", exc_info=True)
return False
uvicorn_server = None
driver = None
app = None
loop = None
main_system = None
async def request_shutdown() -> bool:
"""请求关闭程序"""
@asynccontextmanager
async def create_event_loop_context():
"""创建事件循环的上下文管理器"""
loop = None
try:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
yield loop
except Exception as e:
logger.error(f"创建事件循环失败: {e}")
raise
finally:
if loop and not loop.is_closed():
try:
loop.run_until_complete(graceful_shutdown(maibot.main_system))
except Exception as ge: # 捕捉优雅关闭时可能发生的错误
logger.error(f"优雅关闭时发生错误: {ge}")
return False
return True
except Exception as e:
logger.error(f"请求关闭程序时发生错误: {e}")
await ShutdownManager.graceful_shutdown(loop)
except Exception as e:
logger.error(f"关闭事件循环时出错: {e}")
finally:
try:
loop.close()
logger.info("事件循环已关闭")
except Exception as e:
logger.error(f"关闭事件循环失败: {e}")
class DatabaseManager:
"""数据库连接管理器"""
def __init__(self):
self._connection = None
async def __aenter__(self):
"""异步上下文管理器入口"""
try:
from src.common.database.database import initialize_sql_database
from src.config.config import global_config
logger.info("正在初始化数据库连接...")
start_time = time.time()
# 使用线程执行器运行潜在的阻塞操作
await asyncio.to_thread(initialize_sql_database, global_config.database)
elapsed_time = time.time() - start_time
logger.info(f"数据库连接初始化成功,使用 {global_config.database.database_type} 数据库,耗时: {elapsed_time:.2f}")
return self
except Exception as e:
logger.error(f"数据库连接初始化失败: {e}")
raise
async def __aexit__(self, exc_type, exc_val, exc_tb):
"""异步上下文管理器出口"""
if exc_type:
logger.error(f"数据库操作发生异常: {exc_val}")
return False
def easter_egg():
# 彩蛋
init()
text = "多年以后面对AI行刑队张三将会回想起他2023年在会议上讨论人工智能的那个下午"
rainbow_colors = [Fore.RED, Fore.YELLOW, Fore.GREEN, Fore.CYAN, Fore.BLUE, Fore.MAGENTA]
rainbow_text = ""
for i, char in enumerate(text):
rainbow_text += rainbow_colors[i % len(rainbow_colors)] + char
logger.info(rainbow_text)
async def graceful_shutdown(main_system_instance):
"""优雅地关闭所有系统组件"""
try:
logger.info("正在优雅关闭麦麦...")
# 停止MainSystem中的组件它会处理服务器等
if main_system_instance and hasattr(main_system_instance, "shutdown"):
logger.info("正在关闭MainSystem...")
await main_system_instance.shutdown()
# 停止聊天管理器
class ConfigurationValidator:
"""配置验证器"""
@staticmethod
def validate_configuration():
"""验证关键配置"""
try:
from src.chat.message_receive.chat_stream import get_chat_manager
chat_manager = get_chat_manager()
if hasattr(chat_manager, "_stop_auto_save"):
logger.info("正在停止聊天管理器...")
chat_manager._stop_auto_save()
from src.config.config import global_config
# 检查必要的配置节
required_sections = ['database', 'bot']
for section in required_sections:
if not hasattr(global_config, section):
logger.error(f"配置中缺少{section}配置节")
return False
# 验证数据库配置
db_config = global_config.database
if not hasattr(db_config, 'database_type') or not db_config.database_type:
logger.error("数据库配置缺少database_type字段")
return False
if db_config.database_type not in SUPPORTED_DATABASES:
logger.error(f"不支持的数据库类型: {db_config.database_type}")
logger.info(f"支持的数据库类型: {', '.join(SUPPORTED_DATABASES)}")
return False
logger.info("配置验证通过")
return True
except ImportError:
logger.error("无法导入全局配置模块")
return False
except Exception as e:
logger.warning(f"停止聊天管理器时出错: {e}")
logger.error(f"配置验证失败: {e}")
return False
# 停止情绪管理器
try:
from src.mood.mood_manager import mood_manager
if hasattr(mood_manager, "stop"):
logger.info("正在停止情绪管理器...")
await mood_manager.stop()
except Exception as e:
logger.warning(f"停止情绪管理器时出错: {e}")
class EasterEgg:
"""彩蛋功能"""
_initialized = False
@classmethod
def show(cls):
"""显示彩色文本"""
if not cls._initialized:
init()
cls._initialized = True
text = "多年以后面对AI行刑队张三将会回想起他2023年在会议上讨论人工智能的那个下午"
rainbow_colors = [Fore.RED, Fore.YELLOW, Fore.GREEN, Fore.CYAN, Fore.BLUE, Fore.MAGENTA]
rainbow_text = ""
for i, char in enumerate(text):
rainbow_text += rainbow_colors[i % len(rainbow_colors)] + char
logger.info(rainbow_text)
# 停止记忆系统
try:
from src.chat.memory_system.memory_manager import memory_manager
if hasattr(memory_manager, "shutdown"):
logger.info("正在停止记忆系统...")
await memory_manager.shutdown()
except Exception as e:
logger.warning(f"停止记忆系统时出错: {e}")
# 停止所有异步任务
try:
await async_task_manager.stop_and_wait_all_tasks()
except Exception as e:
logger.warning(f"停止异步任务管理器时出错: {e}")
# 获取所有剩余任务,排除当前任务
remaining_tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
if remaining_tasks:
logger.info(f"正在取消 {len(remaining_tasks)} 个剩余任务...")
# 取消所有剩余任务
for task in remaining_tasks:
if not task.done():
task.cancel()
# 等待所有任务完成,设置超时
try:
await asyncio.wait_for(asyncio.gather(*remaining_tasks, return_exceptions=True), timeout=15.0)
logger.info("所有剩余任务已成功取消")
except asyncio.TimeoutError:
logger.warning("等待任务取消超时,强制继续关闭")
except Exception as e:
logger.error(f"等待任务取消时发生异常: {e}")
logger.info("麦麦优雅关闭完成")
# 关闭日志系统,释放文件句柄
shutdown_logging()
# 尝试停止事件循环
try:
loop = asyncio.get_running_loop()
if loop.is_running():
loop.stop()
logger.info("事件循环已请求停止")
except RuntimeError:
pass # 没有正在运行的事件循环
except Exception as e:
logger.error(f"麦麦关闭失败: {e}", exc_info=True)
def check_eula():
"""检查EULA和隐私条款确认状态 - 环境变量版类似Minecraft"""
# 检查环境变量中的EULA确认
eula_confirmed = os.getenv("EULA_CONFIRMED", "").lower()
if eula_confirmed == "true":
logger.info("EULA已通过环境变量确认")
return
# 如果没有确认,提示用户
confirm_logger.critical("您需要同意EULA和隐私条款才能使用MoFox_Bot")
confirm_logger.critical("请阅读以下文件:")
confirm_logger.critical(" - EULA.md (用户许可协议)")
confirm_logger.critical(" - PRIVACY.md (隐私条款)")
confirm_logger.critical("然后编辑 .env 文件,将 'EULA_CONFIRMED=false' 改为 'EULA_CONFIRMED=true'")
# 等待用户确认
while True:
try:
load_dotenv(override=True) # 重新加载.env文件
eula_confirmed = os.getenv("EULA_CONFIRMED", "").lower()
if eula_confirmed == "true":
confirm_logger.info("EULA确认成功感谢您的同意")
return
confirm_logger.critical("请修改 .env 文件中的 EULA_CONFIRMED=true 后重新启动程序")
input("按Enter键检查.env文件状态...")
except KeyboardInterrupt:
confirm_logger.info("用户取消,程序退出")
sys.exit(0)
except Exception as e:
confirm_logger.error(f"检查EULA状态失败: {e}")
sys.exit(1)
class MaiBotMain(BaseMain):
class MaiBotMain:
"""麦麦机器人主程序类"""
def __init__(self):
super().__init__()
self.main_system = None
def setup_timezone(self):
"""设置时区"""
if platform.system().lower() != "windows":
time.tzset() # type: ignore
def check_and_confirm_eula(self):
"""检查并确认EULA和隐私条款"""
check_eula()
logger.info("检查EULA和隐私条款完成")
try:
if platform.system().lower() != "windows":
time.tzset()
logger.info("时区设置完成")
else:
logger.info("Windows系统跳过时区设置")
except Exception as e:
logger.warning(f"时区设置失败: {e}")
async def initialize_database(self):
"""初始化数据库"""
logger.info("正在初始化数据库连接...")
try:
await initialize_sql_database(global_config.database)
logger.info(f"数据库连接初始化成功,使用 {global_config.database.database_type} 数据库")
except Exception as e:
logger.error(f"数据库连接初始化失败: {e}")
raise e
"""初始化数据库连接"""
async with DatabaseManager():
pass
async def initialize_database_async(self):
"""异步初始化数据库表结构"""
logger.info("正在初始化数据库表结构...")
try:
start_time = time.time()
from src.common.database.sqlalchemy_models import initialize_database as init_db
await init_db()
logger.info("数据库表结构初始化完成")
elapsed_time = time.time() - start_time
logger.info(f"数据库表结构初始化完成,耗时: {elapsed_time:.2f}")
except Exception as e:
logger.error(f"数据库表结构初始化失败: {e}")
raise e
raise
def create_main_system(self):
"""创建MainSystem实例"""
from src.main import MainSystem
self.main_system = MainSystem()
return self.main_system
async def run(self):
"""运行主程序"""
async def run_sync_init(self):
"""执行同步初始化步骤"""
self.setup_timezone()
self.check_and_confirm_eula()
await self.initialize_database()
await EULAManager.check_eula()
if not ConfigurationValidator.validate_configuration():
raise RuntimeError("配置验证失败,请检查配置文件")
return self.create_main_system()
async def run_async_init(self, main_system):
"""执行异步初始化步骤"""
# 初始化数据库连接
await self.initialize_database()
# 初始化数据库表结构
await self.initialize_database_async()
# 初始化主系统
await main_system.initialize()
# 初始化知识库
from src.chat.knowledge.knowledge_lib import initialize_lpmm_knowledge
initialize_lpmm_knowledge()
# 显示彩蛋
EasterEgg.show()
if __name__ == "__main__":
exit_code = 0 # 用于记录程序最终的退出状态
async def wait_for_user_input():
"""等待用户输入(异步方式)"""
try:
# 创建MaiBotMain实例并获取MainSystem
maibot = MaiBotMain()
# 创建事件循环
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
# 在非生产环境下,使用异步方式等待输入
if os.getenv('ENVIRONMENT') != 'production':
logger.info("程序执行完成,按 Ctrl+C 退出...")
# 简单的异步等待,避免阻塞事件循环
while True:
await asyncio.sleep(1)
except KeyboardInterrupt:
logger.info("用户中断程序")
return True
except Exception as e:
logger.error(f"等待用户输入时发生错误: {e}")
return False
async def main_async():
"""主异步函数"""
exit_code = 0
main_task = None
async with create_event_loop_context() as loop:
try:
# 异步初始化数据库和表结构
main_system = loop.run_until_complete(maibot.run())
loop.run_until_complete(maibot.initialize_database_async())
# 执行初始化和任务调度
loop.run_until_complete(main_system.initialize())
initialize_lpmm_knowledge()
# Schedule tasks returns a future that runs forever.
# We can run console_input_loop concurrently.
main_tasks = loop.create_task(main_system.schedule_tasks())
loop.run_until_complete(main_tasks)
# 确保环境文件存在
ConfigManager.ensure_env_file()
# 创建主程序实例并执行初始化
maibot = MaiBotMain()
main_system = await maibot.run_sync_init()
await maibot.run_async_init(main_system)
# 运行主任务
main_task = asyncio.create_task(main_system.schedule_tasks())
logger.info("麦麦机器人启动完成,开始运行主任务...")
# 同时运行主任务和用户输入等待
user_input_done = asyncio.create_task(wait_for_user_input())
# 使用wait等待任意一个任务完成
done, pending = await asyncio.wait(
[main_task, user_input_done],
return_when=asyncio.FIRST_COMPLETED
)
# 如果用户输入任务完成用户按了Ctrl+C取消主任务
if user_input_done in done and main_task not in done:
logger.info("用户请求退出,正在取消主任务...")
main_task.cancel()
try:
await main_task
except asyncio.CancelledError:
logger.info("主任务已取消")
except Exception as e:
logger.error(f"主任务取消时发生错误: {e}")
except KeyboardInterrupt:
logger.warning("收到中断信号,正在优雅关闭...")
# The actual shutdown logic is now in the finally block.
if main_task and not main_task.done():
main_task.cancel()
except Exception as e:
logger.error(f"主程序发生异常: {e}")
logger.debug(f"异常详情: {traceback.format_exc()}")
exit_code = 1
return exit_code
if __name__ == "__main__":
exit_code = 0
try:
exit_code = asyncio.run(main_async())
except KeyboardInterrupt:
logger.info("程序被用户中断")
exit_code = 130
except Exception as e:
logger.error(f"程序发生异常: {e!s} {traceback.format_exc()!s}")
exit_code = 1 # 标记发生错误
logger.error(f"程序启动失败: {e}")
exit_code = 1
finally:
# 确保 loop 在任何情况下都尝试关闭(如果存在且未关闭)
if "loop" in locals() and loop and not loop.is_closed():
logger.info("开始执行最终关闭流程...")
try:
# 传递main_system实例
loop.run_until_complete(graceful_shutdown(maibot.main_system))
except Exception as ge:
logger.error(f"优雅关闭时发生错误: {ge}")
loop.close()
logger.info("事件循环已关闭")
# 在程序退出前暂停,让你有机会看到输出
# input("按 Enter 键退出...") # <--- 添加这行
sys.exit(exit_code) # <--- 使用记录的退出码
# 确保日志系统正确关闭
try:
shutdown_logging()
except Exception as e:
print(f"关闭日志系统时出错: {e}")
sys.exit(exit_code)