初始化

2025-08-11 19:34:18 +08:00
parent ef7a3aee23
commit 23ee3767ef
77 changed files with 10000 additions and 7525 deletions
--- a/src/common/database/sqlalchemy_database_api.py
+++ b/src/common/database/sqlalchemy_database_api.py
@@ -0,0 +1,420 @@
+"""SQLAlchemy数据库API模块
+
+提供基于SQLAlchemy的数据库操作，替换Peewee以解决MySQL连接问题
+支持自动重连、连接池管理和更好的错误处理
+"""
+
+import traceback
+import time
+from typing import Dict, List, Any, Union, Type, Optional
+from contextlib import contextmanager
+from sqlalchemy.orm import Session
+from sqlalchemy.exc import SQLAlchemyError, DisconnectionError, OperationalError
+from sqlalchemy import desc, asc, func, and_, or_
+from src.common.logger import get_logger
+from src.common.database.sqlalchemy_models import (
+    Base, get_db_session, Messages, ActionRecords, PersonInfo, ChatStreams,
+    LLMUsage, Emoji, Images, ImageDescriptions, OnlineTime, Memory,
+    Expression, ThinkingLog, GraphNodes, GraphEdges,get_session
+)
+
+logger = get_logger("sqlalchemy_database_api")
+
+# 模型映射表，用于通过名称获取模型类
+MODEL_MAPPING = {
+    'Messages': Messages,
+    'ActionRecords': ActionRecords,
+    'PersonInfo': PersonInfo,
+    'ChatStreams': ChatStreams,
+    'LLMUsage': LLMUsage,
+    'Emoji': Emoji,
+    'Images': Images,
+    'ImageDescriptions': ImageDescriptions,
+    'OnlineTime': OnlineTime,
+    'Memory': Memory,
+    'Expression': Expression,
+    'ThinkingLog': ThinkingLog,
+    'GraphNodes': GraphNodes,
+    'GraphEdges': GraphEdges,
+}
+
+
+@contextmanager
+def get_db_session():
+    """数据库会话上下文管理器，自动处理事务和连接错误"""
+    session = None
+    max_retries = 3
+    retry_delay = 1.0
+
+    for attempt in range(max_retries):
+        try:
+            session = get_session()
+            yield session
+            session.commit()
+            break
+        except (DisconnectionError, OperationalError) as e:
+            logger.warning(f"数据库连接错误 (尝试 {attempt + 1}/{max_retries}): {e}")
+            if session:
+                session.rollback()
+                session.close()
+            if attempt < max_retries - 1:
+                time.sleep(retry_delay * (attempt + 1))
+            else:
+                raise
+        except Exception as e:
+            if session:
+                session.rollback()
+            raise
+        finally:
+            if session:
+                session.close()
+
+
+def build_filters(session: Session, model_class: Type[Base], filters: Dict[str, Any]):
+    """构建查询过滤条件"""
+    conditions = []
+
+    for field_name, value in filters.items():
+        if not hasattr(model_class, field_name):
+            logger.warning(f"模型 {model_class.__name__} 中不存在字段 '{field_name}'")
+            continue
+
+        field = getattr(model_class, field_name)
+
+        if isinstance(value, dict):
+            # 处理 MongoDB 风格的操作符
+            for op, op_value in value.items():
+                if op == "$gt":
+                    conditions.append(field > op_value)
+                elif op == "$lt":
+                    conditions.append(field < op_value)
+                elif op == "$gte":
+                    conditions.append(field >= op_value)
+                elif op == "$lte":
+                    conditions.append(field <= op_value)
+                elif op == "$ne":
+                    conditions.append(field != op_value)
+                elif op == "$in":
+                    conditions.append(field.in_(op_value))
+                elif op == "$nin":
+                    conditions.append(~field.in_(op_value))
+                else:
+                    logger.warning(f"未知操作符 '{op}' (字段: '{field_name}')")
+        else:
+            # 直接相等比较
+            conditions.append(field == value)
+
+    return conditions
+
+
+async def db_query(
+    model_class: Type[Base],
+    data: Optional[Dict[str, Any]] = None,
+    query_type: Optional[str] = "get",
+    filters: Optional[Dict[str, Any]] = None,
+    limit: Optional[int] = None,
+    order_by: Optional[List[str]] = None,
+    single_result: Optional[bool] = False,
+) -> Union[List[Dict[str, Any]], Dict[str, Any], None]:
+    """执行数据库查询操作
+
+    Args:
+        model_class: SQLAlchemy模型类
+        data: 用于创建或更新的数据字典
+        query_type: 查询类型 ("get", "create", "update", "delete", "count")
+        filters: 过滤条件字典
+        limit: 限制结果数量
+        order_by: 排序字段，前缀'-'表示降序
+        single_result: 是否只返回单个结果
+
+    Returns:
+        根据查询类型返回相应结果
+    """
+    try:
+        if query_type not in ["get", "create", "update", "delete", "count"]:
+            raise ValueError("query_type must be 'get', 'create', 'update', 'delete' or 'count'")
+
+        with get_db_session() as session:
+            if query_type == "get":
+                query = session.query(model_class)
+
+                # 应用过滤条件
+                if filters:
+                    conditions = build_filters(session, model_class, filters)
+                    if conditions:
+                        query = query.filter(and_(*conditions))
+
+                # 应用排序
+                if order_by:
+                    for field_name in order_by:
+                        if field_name.startswith("-"):
+                            field_name = field_name[1:]
+                            if hasattr(model_class, field_name):
+                                query = query.order_by(desc(getattr(model_class, field_name)))
+                        else:
+                            if hasattr(model_class, field_name):
+                                query = query.order_by(asc(getattr(model_class, field_name)))
+
+                # 应用限制
+                if limit and limit > 0:
+                    query = query.limit(limit)
+
+                # 执行查询
+                results = query.all()
+
+                # 转换为字典格式
+                result_dicts = []
+                for result in results:
+                    result_dict = {}
+                    for column in result.__table__.columns:
+                        result_dict[column.name] = getattr(result, column.name)
+                    result_dicts.append(result_dict)
+
+                if single_result:
+                    return result_dicts[0] if result_dicts else None
+                return result_dicts
+
+            elif query_type == "create":
+                if not data:
+                    raise ValueError("创建记录需要提供data参数")
+
+                # 创建新记录
+                new_record = model_class(**data)
+                session.add(new_record)
+                session.flush()  # 获取自动生成的ID
+
+                # 转换为字典格式返回
+                result_dict = {}
+                for column in new_record.__table__.columns:
+                    result_dict[column.name] = getattr(new_record, column.name)
+                return result_dict
+
+            elif query_type == "update":
+                if not data:
+                    raise ValueError("更新记录需要提供data参数")
+
+                query = session.query(model_class)
+
+                # 应用过滤条件
+                if filters:
+                    conditions = build_filters(session, model_class, filters)
+                    if conditions:
+                        query = query.filter(and_(*conditions))
+
+                # 执行更新
+                affected_rows = query.update(data)
+                return affected_rows
+
+            elif query_type == "delete":
+                query = session.query(model_class)
+
+                # 应用过滤条件
+                if filters:
+                    conditions = build_filters(session, model_class, filters)
+                    if conditions:
+                        query = query.filter(and_(*conditions))
+
+                # 执行删除
+                affected_rows = query.delete()
+                return affected_rows
+
+            elif query_type == "count":
+                query = session.query(func.count(model_class.id))
+
+                # 应用过滤条件
+                if filters:
+                    base_query = session.query(model_class)
+                    conditions = build_filters(session, model_class, filters)
+                    if conditions:
+                        base_query = base_query.filter(and_(*conditions))
+                    query = session.query(func.count()).select_from(base_query.subquery())
+
+                return query.scalar()
+
+    except SQLAlchemyError as e:
+        logger.error(f"[SQLAlchemy] 数据库操作出错: {e}")
+        traceback.print_exc()
+
+        # 根据查询类型返回合适的默认值
+        if query_type == "get":
+            return None if single_result else []
+        elif query_type in ["create", "update", "delete", "count"]:
+            return None
+        return None
+
+    except Exception as e:
+        logger.error(f"[SQLAlchemy] 意外错误: {e}")
+        traceback.print_exc()
+
+        if query_type == "get":
+            return None if single_result else []
+        return None
+
+
+async def db_save(
+    model_class: Type[Base],
+    data: Dict[str, Any],
+    key_field: Optional[str] = None,
+    key_value: Optional[Any] = None
+) -> Optional[Dict[str, Any]]:
+    """保存数据到数据库（创建或更新）
+
+    Args:
+        model_class: SQLAlchemy模型类
+        data: 要保存的数据字典
+        key_field: 用于查找现有记录的字段名
+        key_value: 用于查找现有记录的字段值
+
+    Returns:
+        保存后的记录数据或None
+    """
+    try:
+        with get_db_session() as session:
+            # 如果提供了key_field和key_value，尝试更新现有记录
+            if key_field and key_value is not None:
+                if hasattr(model_class, key_field):
+                    existing_record = session.query(model_class).filter(
+                        getattr(model_class, key_field) == key_value
+                    ).first()
+
+                    if existing_record:
+                        # 更新现有记录
+                        for field, value in data.items():
+                            if hasattr(existing_record, field):
+                                setattr(existing_record, field, value)
+
+                        session.flush()
+
+                        # 转换为字典格式返回
+                        result_dict = {}
+                        for column in existing_record.__table__.columns:
+                            result_dict[column.name] = getattr(existing_record, column.name)
+                        return result_dict
+
+            # 创建新记录
+            new_record = model_class(**data)
+            session.add(new_record)
+            session.flush()
+
+            # 转换为字典格式返回
+            result_dict = {}
+            for column in new_record.__table__.columns:
+                result_dict[column.name] = getattr(new_record, column.name)
+            return result_dict
+
+    except SQLAlchemyError as e:
+        logger.error(f"[SQLAlchemy] 保存数据库记录出错: {e}")
+        traceback.print_exc()
+        return None
+    except Exception as e:
+        logger.error(f"[SQLAlchemy] 保存时意外错误: {e}")
+        traceback.print_exc()
+        return None
+
+
+async def db_get(
+    model_class: Type[Base],
+    filters: Optional[Dict[str, Any]] = None,
+    limit: Optional[int] = None,
+    order_by: Optional[str] = None,
+    single_result: Optional[bool] = False,
+) -> Union[List[Dict[str, Any]], Dict[str, Any], None]:
+    """从数据库获取记录
+
+    Args:
+        model_class: SQLAlchemy模型类
+        filters: 过滤条件
+        limit: 结果数量限制
+        order_by: 排序字段，前缀'-'表示降序
+        single_result: 是否只返回单个结果
+
+    Returns:
+        记录数据或None
+    """
+    order_by_list = [order_by] if order_by else None
+    return await db_query(
+        model_class=model_class,
+        query_type="get",
+        filters=filters,
+        limit=limit,
+        order_by=order_by_list,
+        single_result=single_result
+    )
+
+
+async def store_action_info(
+    chat_stream=None,
+    action_build_into_prompt: bool = False,
+    action_prompt_display: str = "",
+    action_done: bool = True,
+    thinking_id: str = "",
+    action_data: Optional[dict] = None,
+    action_name: str = "",
+) -> Optional[Dict[str, Any]]:
+    """存储动作信息到数据库
+
+    Args:
+        chat_stream: 聊天流对象
+        action_build_into_prompt: 是否将此动作构建到提示中
+        action_prompt_display: 动作的提示显示文本
+        action_done: 动作是否完成
+        thinking_id: 关联的思考ID
+        action_data: 动作数据字典
+        action_name: 动作名称
+
+    Returns:
+        保存的记录数据或None
+    """
+    try:
+        import json
+
+        # 构建动作记录数据
+        record_data = {
+            "action_id": thinking_id or str(int(time.time() * 1000000)),
+            "time": time.time(),
+            "action_name": action_name,
+            "action_data": json.dumps(action_data or {}, ensure_ascii=False),
+            "action_done": action_done,
+            "action_build_into_prompt": action_build_into_prompt,
+            "action_prompt_display": action_prompt_display,
+        }
+
+        # 从chat_stream获取聊天信息
+        if chat_stream:
+            record_data.update({
+                "chat_id": getattr(chat_stream, "stream_id", ""),
+                "chat_info_stream_id": getattr(chat_stream, "stream_id", ""),
+                "chat_info_platform": getattr(chat_stream, "platform", ""),
+            })
+        else:
+            record_data.update({
+                "chat_id": "",
+                "chat_info_stream_id": "",
+                "chat_info_platform": "",
+            })
+
+        # 保存记录
+        saved_record = await db_save(
+            ActionRecords,
+            data=record_data,
+            key_field="action_id",
+            key_value=record_data["action_id"]
+        )
+
+        if saved_record:
+            logger.debug(f"[SQLAlchemy] 成功存储动作信息: {action_name} (ID: {record_data['action_id']})")
+        else:
+            logger.error(f"[SQLAlchemy] 存储动作信息失败: {action_name}")
+
+        return saved_record
+
+    except Exception as e:
+        logger.error(f"[SQLAlchemy] 存储动作信息时发生错误: {e}")
+        traceback.print_exc()
+        return None
+
+
+# 兼容性函数，方便从Peewee迁移
+def get_model_class(model_name: str) -> Optional[Type[Base]]:
+    """根据模型名称获取模型类"""
+    return MODEL_MAPPING.get(model_name)