from typing import Any from agno.db.base import SessionType from agno.db.mysql import MySQLDb from agno.db.postgres import PostgresDb from agno.db.sqlite import SqliteDb from agno.session.team import TeamSession from app.api.v1.module_system.auth.schema import AuthSchema from app.config.setting import settings from app.core.logger import log from .schema import ChatSessionCreateSchema, ChatSessionUpdateSchema class ChatSessionCRUD: """聊天会话数据层 - 使用 agno 数据库存储""" # 会话类型配置 - 使用 TEAM 类型因为创建的是 Team SESSION_TYPE = SessionType.TEAM def __init__(self, auth: AuthSchema) -> None: """初始化CRUD数据层""" self.auth = auth self.user_id = auth.user.username if auth and auth.user else "user" self.team_id = str(auth.user.dept_id) if auth and auth.user and hasattr(auth.user, 'dept_id') and auth.user.dept_id else None self.db = self._get_db() def _get_db(self) -> Any: """获取数据库连接""" db_type = settings.DATABASE_TYPE db_uri = settings.DB_URI db_mapping = { "mysql": lambda: MySQLDb(db_url=db_uri), "postgres": lambda: PostgresDb(db_url=db_uri), "sqlite": lambda: SqliteDb(db_file=db_uri.replace("sqlite:///", "")), } if db_type not in db_mapping: raise ValueError(f"不支持的数据库类型: {db_type}") return db_mapping[db_type]() def __del__(self) -> None: """析构时关闭数据库连接""" self.db.close() async def get_by_id_crud(self, session_id: str) -> TeamSession | None: """获取会话详情""" try: return self.db.get_session( session_id=session_id, session_type=self.SESSION_TYPE, user_id=self.user_id ) except Exception as e: log.error(f"获取会话详情失败: {e}") return None async def list_crud( self, search: dict[str, Any] | None = None, order_by: list[dict[str, str]] | None = None, ) -> list[TeamSession]: """列表查询 - 获取所有会话""" try: result = self.db.get_sessions( session_type=self.SESSION_TYPE, user_id=self.user_id ) if isinstance(result, tuple) and len(result) == 2: return result[0] return result if isinstance(result, list) else [] except Exception as e: log.error(f"获取会话列表失败: {e}") return [] async def create_crud(self, data: ChatSessionCreateSchema) -> TeamSession | None: """创建会话 - Team 会在运行时自动创建和管理 session""" import time import uuid try: session_id = str(uuid.uuid4()) now = int(time.time()) # 创建 session_data,包含 session_name session_data = {} if data.title: session_data["session_name"] = data.title # 创建 TeamSession 对象 session = TeamSession( session_id=session_id, user_id=self.user_id, team_id=self.team_id, session_data=session_data, created_at=now, updated_at=now, ) # 保存会话 result = self.db.upsert_session(session=session) return result except Exception as e: log.exception(f"创建会话失败: {e}") return None async def update_crud(self, session_id: str, data: ChatSessionUpdateSchema) -> bool: """更新会话""" try: self.db.rename_session( session_id=session_id, session_type=self.SESSION_TYPE, session_name=data.title, user_id=self.user_id ) return True except Exception as e: log.error(f"更新会话失败: {e}") return False async def delete_crud(self, session_ids: list[str]) -> bool: """批量删除会话""" try: for session_id in session_ids: self.db.delete_session( session_id=session_id, user_id=self.user_id ) return True except Exception as e: log.error(f"删除会话失败: {e}") return False