296 lines
11 KiB
Python
296 lines
11 KiB
Python
"""
|
||
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
|