edbddf855d
- 实现Agent管理，支持AI辅助生成系统提示词 - 支持多个AI提供商（OpenRouter、智谱、MiniMax等） - 实现聊天室和讨论引擎 - WebSocket实时消息推送 - 前端使用React + Ant Design - 后端使用FastAPI + MongoDB Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
95 lines
2.3 KiB
Python
95 lines
2.3 KiB
Python
"""
|
||
MongoDB数据库连接模块
|
||
使用Motor异步驱动
|
||
"""
|
||
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
|
||
from beanie import init_beanie
|
||
from loguru import logger
|
||
from typing import Optional
|
||
|
||
from config import settings
|
||
|
||
# 全局数据库客户端和数据库实例
|
||
_client: Optional[AsyncIOMotorClient] = None
|
||
_database: Optional[AsyncIOMotorDatabase] = None
|
||
|
||
|
||
async def connect_db() -> None:
|
||
"""
|
||
连接MongoDB数据库
|
||
初始化Beanie ODM
|
||
"""
|
||
global _client, _database
|
||
|
||
try:
|
||
_client = AsyncIOMotorClient(settings.MONGODB_URL)
|
||
_database = _client[settings.MONGODB_DB]
|
||
|
||
# 导入所有文档模型用于初始化Beanie
|
||
from models.ai_provider import AIProvider
|
||
from models.agent import Agent
|
||
from models.chatroom import ChatRoom
|
||
from models.message import Message
|
||
from models.discussion_result import DiscussionResult
|
||
from models.agent_memory import AgentMemory
|
||
|
||
# 初始化Beanie
|
||
await init_beanie(
|
||
database=_database,
|
||
document_models=[
|
||
AIProvider,
|
||
Agent,
|
||
ChatRoom,
|
||
Message,
|
||
DiscussionResult,
|
||
AgentMemory,
|
||
]
|
||
)
|
||
|
||
logger.info(f"已连接到MongoDB数据库: {settings.MONGODB_DB}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"数据库连接失败: {e}")
|
||
raise
|
||
|
||
|
||
async def close_db() -> None:
|
||
"""
|
||
关闭数据库连接
|
||
"""
|
||
global _client
|
||
|
||
if _client:
|
||
_client.close()
|
||
logger.info("数据库连接已关闭")
|
||
|
||
|
||
def get_database() -> AsyncIOMotorDatabase:
|
||
"""
|
||
获取数据库实例
|
||
|
||
Returns:
|
||
MongoDB数据库实例
|
||
|
||
Raises:
|
||
RuntimeError: 数据库未初始化
|
||
"""
|
||
if _database is None:
|
||
raise RuntimeError("数据库未初始化，请先调用connect_db()")
|
||
return _database
|
||
|
||
|
||
def get_client() -> AsyncIOMotorClient:
|
||
"""
|
||
获取数据库客户端
|
||
|
||
Returns:
|
||
MongoDB客户端实例
|
||
|
||
Raises:
|
||
RuntimeError: 客户端未初始化
|
||
"""
|
||
if _client is None:
|
||
raise RuntimeError("数据库客户端未初始化")
|
||
return _client
|