""" GLM聊天服务，包含角色配置管理 """ import logging import asyncio import os from typing import Optional from pymongo import AsyncMongoClient from beanie import init_beanie from app.models.glm_role import GLMRole from app.models import models logger = logging.getLogger(__name__) class GLMService: """GLM聊天服务，包含角色配置管理""" def __init__(self): self._current_role: Optional[GLMRole] = None self._max_retries = 3 self._retry_delay = 1 # 秒 self._db_client: Optional[AsyncMongoClient] = None self._db_initialized = False async def get_current_role(self) -> GLMRole: """获取当前激活的角色配置""" if self._current_role is None: logger.info("当前无缓存角色，开始加载角色配置") self._current_role = await self._load_or_create_role_with_retry() logger.info(f"角色加载完成: name='{self._current_role.name}'") else: logger.debug(f"使用缓存角色: name='{self._current_role.name}'") return self._current_role async def update_role(self, role_data: dict) -> GLMRole: """更新角色配置 Args: role_data: 角色配置数据，包含name, description, skills, restrictions字段 Returns: 更新后的GLMRole实例 """ logger.info(f"开始更新角色配置，更新字段: {list(role_data.keys())}") role = await self.get_current_role() # 记录更新前的值 changes = [] if "name" in role_data and role_data["name"] != role.name: changes.append(f"name: '{role.name}' -> '{role_data['name']}'") role.name = role_data["name"] if "description" in role_data and role_data["description"] != role.description: changes.append( f"description: {len(role.description)}字符 -> {len(role_data['description'])}字符" ) role.description = role_data["description"] if "skills" in role_data and role_data["skills"] != role.skills: old_skills_count = len(role.skills.split("\n")) if role.skills else 0 new_skills_count = ( len(role_data["skills"].split("\n")) if role_data["skills"] else 0 ) changes.append(f"skills: {old_skills_count}条 -> {new_skills_count}条") role.skills = role_data["skills"] if ( "restrictions" in role_data and role_data["restrictions"] != role.restrictions ): old_restrictions_count = ( len(role.restrictions.split("\n")) if role.restrictions else 0 ) new_restrictions_count = ( len(role_data["restrictions"].split("\n")) if role_data["restrictions"] else 0 ) changes.append( f"restrictions: {old_restrictions_count}条 -> {new_restrictions_count}条" ) role.restrictions = role_data["restrictions"] role.update_timestamp() if changes: logger.info(f"角色配置变更: {', '.join(changes)}") else: logger.info("角色配置无实际变更") # 使用重试机制保存 await self._save_with_retry(role) # 更新缓存 self._current_role = role logger.info(f"角色配置更新完成: name='{role.name}'") return role async def _load_or_create_role_with_retry(self) -> GLMRole: """使用重试机制加载或创建默认角色""" for attempt in range(self._max_retries): try: return await self._load_or_create_role() except Exception as e: if attempt < self._max_retries - 1: logger.warning(f"加载角色配置失败，第{attempt + 1}次重试: {e}") await asyncio.sleep(self._retry_delay * (attempt + 1)) else: logger.error(f"加载角色配置失败，已达到最大重试次数: {e}") raise async def _ensure_db_connection(self) -> None: """确保数据库连接正常，如果连接已关闭则重新连接""" if not self._db_initialized: await self._init_db_connection() async def _init_db_connection(self) -> None: """初始化数据库连接""" try: # 从环境变量获取数据库配置 mongodb_url = os.getenv("MONGODB_URL", "mongodb://localhost:27017") database_name = os.getenv("DATABASE_NAME", "douba_chat_dev") # 创建新的数据库连接 if self._db_client: try: await self._db_client.close() except: pass self._db_client = AsyncMongoClient(mongodb_url) await init_beanie( database=self._db_client[database_name], document_models=models ) self._db_initialized = True logger.info(f"数据库连接初始化成功: {database_name}") except Exception as e: logger.error(f"数据库连接初始化失败: {e}") self._db_initialized = False raise async def _load_or_create_role(self) -> GLMRole: """加载或创建默认角色""" try: # 确保数据库连接正常 await self._ensure_db_connection() # 1. 尝试获取激活的角色 active_role = await GLMRole.find_one(GLMRole.is_active == True) if active_role: logger.info(f"加载激活角色: {active_role.name}") return active_role # 2. 尝试获取默认角色 default_role = await GLMRole.find_one(GLMRole.is_default == True) if default_role: default_role.is_active = True await default_role.save() logger.info(f"加载默认角色并激活: {default_role.name}") return default_role # 3. 创建默认角色 logger.info("创建默认角色配置") return await self._create_default_role() except Exception as e: logger.error(f"加载或创建角色失败: {e}") # 检查是否是数据库连接错误 if ( "Cannot use AsyncMongoClient after close" in str(e) or "not connected" in str(e).lower() ): logger.error("数据库连接已关闭，尝试重新连接...") self._db_initialized = False # 尝试重新连接 try: await self._init_db_connection() logger.info("数据库重新连接成功，重试操作...") # 重试操作 return await self._retry_after_reconnect() except Exception as reconnect_error: logger.error(f"数据库重新连接失败: {reconnect_error}") raise async def _retry_after_reconnect(self) -> GLMRole: """重新连接后重试操作""" # 1. 尝试获取激活的角色 active_role = await GLMRole.find_one(GLMRole.is_active == True) if active_role: logger.info(f"重新连接后加载激活角色: {active_role.name}") return active_role # 2. 尝试获取默认角色 default_role = await GLMRole.find_one(GLMRole.is_default == True) if default_role: default_role.is_active = True await default_role.save() logger.info(f"重新连接后加载默认角色并激活: {default_role.name}") return default_role # 3. 创建默认角色 logger.info("重新连接后创建默认角色配置") return await self._create_default_role() async def _create_default_role(self) -> GLMRole: """创建默认角色配置""" # 导入DEFAULT_ROLE_CONFIG from app.routes.glm_routes import DEFAULT_ROLE_CONFIG role = GLMRole( name=DEFAULT_ROLE_CONFIG.name, description=DEFAULT_ROLE_CONFIG.description, skills=DEFAULT_ROLE_CONFIG.skills, restrictions=DEFAULT_ROLE_CONFIG.restrictions, is_active=True, is_default=True, ) await role.save() logger.info(f"默认角色创建成功: {role.name}") return role async def _save_with_retry(self, role: GLMRole) -> None: """使用重试机制保存角色""" for attempt in range(self._max_retries): try: # 确保数据库连接正常 await self._ensure_db_connection() await role.save() return except Exception as e: if attempt < self._max_retries - 1: logger.warning(f"保存角色配置失败，第{attempt + 1}次重试: {e}") # 检查是否是数据库连接错误 if ( "Cannot use AsyncMongoClient after close" in str(e) or "not connected" in str(e).lower() ): logger.warning("检测到数据库连接问题，尝试重新连接...") self._db_initialized = False try: await self._init_db_connection() logger.info("数据库重新连接成功，继续重试...") except Exception as reconnect_error: logger.error(f"数据库重新连接失败: {reconnect_error}") await asyncio.sleep(self._retry_delay * (attempt + 1)) else: logger.error(f"保存角色配置失败，已达到最大重试次数: {e}") raise async def ensure_default_role(self): """确保存在默认角色配置（应用启动时调用）""" try: # 确保数据库连接正常 await self._ensure_db_connection() role = await self._load_or_create_role_with_retry() logger.info(f"默认角色已确保存在: {role.name}") return role except Exception as e: logger.error(f"确保默认角色存在时出错: {e}") raise async def close(self): """关闭服务，清理资源""" self._current_role = None if self._db_client: try: await self._db_client.close() logger.info("数据库连接已关闭") except Exception as e: logger.error(f"关闭数据库连接时出错: {e}") self._db_initialized = False logger.info("GLM服务已关闭") # 全局服务实例（单例模式） _glm_service: Optional[GLMService] = None def get_glm_service() -> GLMService: """获取GLM服务实例（单例）""" global _glm_service if _glm_service is None: _glm_service = GLMService() return _glm_service def close_glm_service(): """关闭GLM服务""" global _glm_service if _glm_service: _glm_service.close() _glm_service = None