import json import threading import time import uuid from typing import Any from agno.db.base import SessionType from agno.db.mysql import AsyncMySQLDb from agno.db.postgres import AsyncPostgresDb from agno.db.sqlite import AsyncSqliteDb from agno.session.team import TeamSession from app.common.enums import RET from app.config.setting import settings from app.core.base_schema import AuthSchema from app.core.exceptions import CustomException from app.core.logger import logger from .schema import ChatSessionCreateSchema, ChatSessionUpdateSchema # Team 标识:agno 会话按 user_id + team_id 隔离 TEAM_ID = "default" # agno 异步 DB 实例为进程级单例:内部维护异步引擎与连接池,避免每请求重建 _db_instance: Any | None = None _db_lock = threading.Lock() def _get_agno_db() -> Any: """获取 agno 异步数据库实例(进程级单例)""" global _db_instance if _db_instance is None: with _db_lock: if _db_instance is None: db_type = settings.DATABASE_TYPE # agno 异步 DB 内部走 SQLAlchemy asyncio 扩展(create_async_engine), # 必须用异步驱动 URI;传同步的 DB_URI(pymysql/psycopg)会报 InvalidRequestError db_uri = settings.ASYNC_DB_URI if db_type == "mysql": _db_instance = AsyncMySQLDb(db_url=db_uri, db_schema=settings.DATABASE_NAME, create_schema=False) elif db_type == "postgres": _db_instance = AsyncPostgresDb(db_url=db_uri, db_schema="public", create_schema=False) elif db_type == "sqlite": _db_instance = AsyncSqliteDb(db_file=db_uri.replace("sqlite+aiosqlite:///", "")) else: # settings.DATABASE_TYPE 是服务端配置:取值不在支持范围属部署配置错误,按 500 返回 logger.error(f"不支持的数据库类型: {db_type}") raise RuntimeError(f"不支持的数据库类型: {db_type}") return _db_instance async def init_agno_tables() -> None: """应用启动时确保 agno 表结构存在(幂等,表已存在则跳过)。 agno 的 create_schema 仅控制是否自动建库,建表需调用方显式触发; 不初始化则首次读写会话时报 "Table has an invalid schema"。 """ db = _get_agno_db() try: await db._create_all_tables() except Exception as e: # 多 worker 并发首启时可能同时建表,"已存在"类冲突可安全忽略 logger.warning(f"agno 表初始化告警(多为并发建表冲突,可忽略): {e}") def _row_session_name(row: dict[str, Any]) -> str: """从 deserialize=False 的原始会话行中提取会话标题(session_data 可能是 JSON 字符串或字典)""" session_data = row.get("session_data") if isinstance(session_data, str): try: session_data = json.loads(session_data) except (json.JSONDecodeError, TypeError): return "" if isinstance(session_data, dict): return str(session_data.get("session_name") or "") return "" class ChatSessionCRUD: """聊天会话数据层 - 使用 agno 异步数据库存储""" SESSION_TYPE = SessionType.TEAM def __init__(self, auth: AuthSchema) -> None: """初始化CRUD数据层 参数: - auth (AuthSchema): 认证信息模型 """ self.auth = auth self.user_id = auth.user.username or "user" self.team_id = TEAM_ID self.db = _get_agno_db() async def get_by_id(self, session_id: str) -> TeamSession | None: """获取会话详情;不存在或失败时返回 None""" try: return await self.db.get_session( session_id=session_id, session_type=self.SESSION_TYPE, user_id=self.user_id, ) except Exception as e: logger.error(f"获取会话详情失败: {e}") return None async def list_page( self, page_no: int, page_size: int, title: str | None = None, order_by: list[dict[str, str]] | None = None, ) -> tuple[list[dict[str, Any]], int]: """分页查询会话(分页/排序/计数均在 DB 层完成)。 使用 deserialize=False 返回原始字典,避免全量反序列化 runs 消息; 返回 (会话字典列表, 总数)。 """ # 前端排序字段(created_time/updated_time)映射为 agno 会话表列名 sort_by, sort_order = "updated_at", "desc" if order_by: column, direction = next(iter(order_by[0])) sort_by = { "created_time": "created_at", "updated_time": "updated_at", "created_at": "created_at", "updated_at": "updated_at", }.get(column, "updated_at") sort_order = direction if direction in ("asc", "desc") else "desc" # 标题搜索为低频操作:agno 未实现该过滤,需全量拉取当前用户会话后内存过滤+分页 fetch_page = page_no if not title else None fetch_limit = page_size if not title else None try: result = await self.db.get_sessions( session_type=self.SESSION_TYPE, user_id=self.user_id, limit=fetch_limit, page=fetch_page, sort_by=sort_by, sort_order=sort_order, deserialize=False, ) except Exception as e: logger.error(f"获取会话列表失败: {e}") raise rows: list[dict[str, Any]] = [] total = 0 if isinstance(result, tuple) and len(result) == 2: rows, total = list(result[0] or []), int(result[1] or 0) else: rows, total = list(result or []), len(result or []) if title: keyword = title.strip().lower() filtered = [r for r in rows if keyword in _row_session_name(r).lower()] start = (page_no - 1) * page_size return filtered[start : start + page_size], len(filtered) return rows, total async def create(self, data: ChatSessionCreateSchema) -> TeamSession: """创建会话(Team 在运行时自动创建并管理 session)""" now = int(time.time()) session_data = {"session_name": data.title} if data.title else {} session = TeamSession( session_id=str(uuid.uuid4()), user_id=self.user_id, team_id=self.team_id, session_data=session_data, created_at=now, updated_at=now, ) try: result = await self.db.upsert_session(session=session) except Exception as e: logger.exception(f"创建会话失败: {e}") raise if result is None: logger.error("创建会话失败") raise RuntimeError("创建会话失败") return result async def rename(self, session_id: str, data: ChatSessionUpdateSchema) -> None: """重命名会话""" try: result = await self.db.rename_session( session_id=session_id, session_type=self.SESSION_TYPE, session_name=data.title, user_id=self.user_id, ) except Exception as e: logger.error(f"更新会话失败: {e}") raise if result is None: raise CustomException(msg="会话不存在", code=RET.NOT_FOUND.code) async def delete(self, session_ids: list[str]) -> None: """批量删除会话""" for session_id in session_ids: try: await self.db.delete_session(session_id=session_id, user_id=self.user_id) except Exception as e: logger.error(f"删除会话失败: {session_id} - {e}") # agno/SQLite 侧删除失败属内部故障,按 500 返回(session_ids 为客户端参数,但失败原因在服务端) raise