LongjoeAgent/core/chat_result_manager.py

import os
import sqlite3
import uuid
import json
import datetime
import re
from typing import Dict, Any, Optional
from pathlib import Path
from utils.logger import chat_logger


class ChatResultManager:
    """聊天结果管理器 - 支持按用户+日期分割的SQLite数据库存储"""

    def __init__(self):
        self._db_connections = {}  # 存储不同用户+日期组合的连接
        self._task_db_cache = {}  # 缓存任务ID对应的数据库文件
        self._init_tables()

    def _get_day_number(self, date=None):
        """获取日期编号 (YYYYMMDD 格式)"""
        if date is None:
            date = datetime.datetime.now()
        return date.strftime("%Y%m%d")

    def _get_db_key(self, username, day_number=None):
        """生成数据库连接缓存key"""
        if day_number is None:
            day_number = self._get_day_number()
        return f"{username}_{day_number}"

    def _get_db_path(self, username, day_number=None):
        """获取数据库文件路径"""
        if day_number is None:
            day_number = self._get_day_number()

        # 数据库文件存放目录
        project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        base_dir = os.path.join(project_root, "data", "chat_history")
        os.makedirs(base_dir, exist_ok=True)

        # 数据库文件名格式: chat_history_{username}_{day_number}.db
        # 对用户名进行安全处理，避免特殊字符
        safe_username = re.sub(r'[\\/:*?"<>|\0]', "_", username)
        print(f"_get_db_path用户名: {safe_username}")
        db_filename = f"chat_history_{safe_username}_{day_number}.db"
        return os.path.join(base_dir, db_filename)

    def _get_db_connection(self, username, day_number=None):
        """获取数据库连接（按用户+日期分割）"""
        if day_number is None:
            day_number = self._get_day_number()

        db_key = self._get_db_key(username, day_number)

        # 如果连接不存在或需要重新连接，创建新连接
        if db_key not in self._db_connections:
            db_path = self._get_db_path(username, day_number)
            conn = sqlite3.connect(db_path, check_same_thread=False)
            conn.execute("PRAGMA journal_mode=WAL")
            self._db_connections[db_key] = conn

            # 初始化该数据库的表
            self._init_tables_for_connection(conn)

        return self._db_connections[db_key]

    def _init_tables_for_connection(self, conn):
        """为指定数据库连接初始化表结构"""
        cursor = conn.cursor()

        # 创建聊天任务表
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS chat_tasks (
                task_id TEXT PRIMARY KEY,
                status TEXT NOT NULL,  -- pending, processing, completed, failed
                request_data TEXT NOT NULL,
                response_data TEXT,
                error_message TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                thread_id TEXT,
                username TEXT
            )
        """
        )

        # 创建索引
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_status ON chat_tasks(status)")
        cursor.execute(
            "CREATE INDEX IF NOT EXISTS idx_thread_id ON chat_tasks(thread_id)"
        )
        cursor.execute(
            "CREATE INDEX IF NOT EXISTS idx_username ON chat_tasks(username)"
        )
        cursor.execute(
            "CREATE INDEX IF NOT EXISTS idx_created_at ON chat_tasks(created_at)"
        )

        conn.commit()

    def _init_tables(self):
        """初始化所有数据库表（兼容性方法）"""
        # 为默认用户初始化表，确保系统启动时至少有一个数据库可用
        default_username = "default"
        conn = self._get_db_connection(default_username)
        self._init_tables_for_connection(conn)

    def create_task(self, request_data: Dict[str, Any]) -> str:
        """创建新的聊天任务"""
        task_id = str(uuid.uuid4())
        username = request_data.get("username", "default")
        print(f"创建任务: {task_id} (用户: {username})")
        conn = self._get_db_connection(username)
        cursor = conn.cursor()

        cursor.execute(
            """
            INSERT INTO chat_tasks 
            (task_id, status, request_data, thread_id, username)
            VALUES (?, ?, ?, ?, ?)
        """,
            (
                task_id,
                "pending",
                json.dumps(request_data, ensure_ascii=False),
                request_data.get("thread_id", "default"),
                username,
            ),
        )

        conn.commit()

        # 缓存任务ID对应的数据库信息
        self._task_db_cache[task_id] = {
            "username": username,
            "day_number": self._get_day_number(),
            "db_file": os.path.basename(self._get_db_path(username)),
        }

        chat_logger.info(f"创建聊天任务: {task_id} (用户: {username})")
        return task_id

    def update_task_status(
        self,
        task_id: str,
        status: str,
        response_data: Dict[str, Any] = None,
        error_message: str = None,
    ):
        """更新任务状态"""
        # 首先检查缓存中是否有该任务对应的数据库信息
        if task_id in self._task_db_cache:
            cache_info = self._task_db_cache[task_id]
            username = cache_info["username"]
            day_number = cache_info["day_number"]

            # 尝试在缓存的数据库中更新
            db_path = self._get_db_path(username, day_number)
            if os.path.exists(db_path):
                try:
                    self._update_task_in_db(
                        username,
                        task_id,
                        status,
                        response_data,
                        error_message,
                        day_number,
                    )
                    return
                except Exception as e:
                    chat_logger.warning(
                        f"使用缓存更新任务状态失败: {task_id} - {str(e)}"
                    )
                    # 缓存可能已失效，继续使用备用策略
                    del self._task_db_cache[task_id]

        # 备用策略：如果缓存不存在或失效，使用原来的get_task方法查找
        task_info = self.get_task(task_id)
        if not task_info:
            chat_logger.error(f"更新任务状态失败: 任务不存在 - {task_id}")
            return

        username = task_info.get("username", "default")
        self._update_task_in_db(username, task_id, status, response_data, error_message)

    def _update_task_in_db(
        self,
        username,
        task_id,
        status,
        response_data=None,
        error_message=None,
        day_number=None,
    ):
        """在指定数据库中更新任务状态（内部方法）"""
        conn = self._get_db_connection(username, day_number)
        cursor = conn.cursor()

        cursor.execute(
            """
            UPDATE chat_tasks 
            SET status = ?, 
                response_data = ?, 
                error_message = ?,
                updated_at = CURRENT_TIMESTAMP
            WHERE task_id = ?
        """,
            (
                status,
                (
                    json.dumps(response_data, ensure_ascii=False)
                    if response_data
                    else None
                ),
                error_message,
                task_id,
            ),
        )

        conn.commit()
        chat_logger.info(f"更新任务状态: {task_id} -> {status} (用户: {username})")

    def _search_task_in_db(self, db_path, task_id):
        """在指定数据库中搜索任务"""
        try:
            conn = sqlite3.connect(db_path, check_same_thread=False)
            cursor = conn.cursor()

            cursor.execute(
                """
                SELECT task_id, status, request_data, response_data, 
                       error_message, created_at, updated_at, thread_id, username
                FROM chat_tasks 
                WHERE task_id = ?
            """,
                (task_id,),
            )

            row = cursor.fetchone()
            conn.close()

            if row:
                return {
                    "task_id": row[0],
                    "status": row[1],
                    "request_data": json.loads(row[2]) if row[2] else {},
                    "response_data": json.loads(row[3]) if row[3] else None,
                    "error_message": row[4],
                    "created_at": row[5],
                    "updated_at": row[6],
                    "thread_id": row[7],
                    "username": row[8],
                }
        except Exception as e:
            chat_logger.warning(
                f"搜索任务时数据库访问异常: {os.path.basename(db_path)} - {str(e)}"
            )

        return None

    def get_task(self, task_id: str) -> Optional[Dict[str, Any]]:
        """获取任务信息 - 智能查找策略"""
        # 1. 首先检查缓存中是否有该任务对应的数据库信息
        if task_id in self._task_db_cache:
            cache_info = self._task_db_cache[task_id]
            username = cache_info["username"]
            day_number = cache_info["day_number"]

            # 尝试在缓存的数据库中查找
            db_path = self._get_db_path(username, day_number)
            if os.path.exists(db_path):
                result = self._search_task_in_db(db_path, task_id)
                if result:
                    return result
            else:
                # 如果缓存的文件不存在，清除缓存
                del self._task_db_cache[task_id]

        # 2. 如果没有缓存，尝试从请求上下文推断用户名
        # 这里可以扩展为从请求头或其他上下文获取用户名
        # 目前先使用默认策略

        # 3. 按优先级搜索策略
        project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        base_dir = os.path.join(project_root, "data", "chat_history")

        if not os.path.exists(base_dir):
            return None

        # 策略1: 搜索当天所有用户的数据库（最可能的情况）
        today = self._get_day_number()
        for db_file in os.listdir(base_dir):
            if db_file.endswith(f"_{today}.db") and db_file.startswith("chat_history_"):
                db_path = os.path.join(base_dir, db_file)
                result = self._search_task_in_db(db_path, task_id)
                if result:
                    # 找到后缓存数据库信息
                    self._task_db_cache[task_id] = {
                        "username": result["username"],
                        "day_number": today,
                        "db_file": db_file,
                    }
                    return result

        # 策略2: 搜索最近7天内的数据库（次优情况）
        for days_back in range(1, 8):
            search_date = datetime.datetime.now() - datetime.timedelta(days=days_back)
            search_day = self._get_day_number(search_date)

            for db_file in os.listdir(base_dir):
                if db_file.endswith(f"_{search_day}.db") and db_file.startswith(
                    "chat_history_"
                ):
                    db_path = os.path.join(base_dir, db_file)
                    result = self._search_task_in_db(db_path, task_id)
                    if result:
                        # 找到后缓存数据库信息
                        self._task_db_cache[task_id] = {
                            "username": result["username"],
                            "day_number": search_day,
                            "db_file": db_file,
                        }
                        return result

        # 策略3: 最后搜索所有数据库（最坏情况）
        for db_file in os.listdir(base_dir):
            if db_file.endswith(".db") and db_file.startswith("chat_history_"):
                db_path = os.path.join(base_dir, db_file)
                result = self._search_task_in_db(db_path, task_id)
                if result:
                    # 从文件名中提取用户名和日期
                    parts = (
                        db_file.replace("chat_history_", "")
                        .replace(".db", "")
                        .split("_")
                    )
                    if len(parts) >= 2:
                        username_part = "_".join(parts[:-1])  # 用户名可能包含下划线
                        day_number = parts[-1]

                        self._task_db_cache[task_id] = {
                            "username": result["username"],
                            "day_number": day_number,
                            "db_file": db_file,
                        }
                    return result

        return None

    def cleanup_old_tasks(self, max_days=7):
        """清理超过指定天数的旧任务数据"""
        cutoff_date = datetime.datetime.now() - datetime.timedelta(days=max_days)
        cutoff_day = self._get_day_number(cutoff_date)

        project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        base_dir = os.path.join(project_root, "data", "chat_history")

        if not os.path.exists(base_dir):
            return

        cleaned_count = 0
        # 清理旧的数据库文件
        for db_file in os.listdir(base_dir):
            if db_file.endswith(".db") and db_file.startswith("chat_history_"):
                # 从文件名中提取日期部分
                # 格式: chat_history_{username}_{YYYYMMDD}.db
                parts = db_file.split("_")
                if len(parts) >= 3:
                    try:
                        file_day = parts[-1].replace(".db", "")
                        if file_day < cutoff_day:
                            db_path = os.path.join(base_dir, db_file)
                            os.remove(db_path)
                            cleaned_count += 1
                            chat_logger.info(f"清理旧数据库文件: {db_file}")

                            # 清理对应的缓存
                            task_ids_to_remove = []
                            for tid, cache_info in self._task_db_cache.items():
                                if cache_info["db_file"] == db_file:
                                    task_ids_to_remove.append(tid)

                            for tid in task_ids_to_remove:
                                del self._task_db_cache[tid]
                    except (ValueError, IndexError):
                        continue

        chat_logger.info(f"清理完成: 删除了 {cleaned_count} 个旧数据库文件")

    def close(self):
        """关闭所有数据库连接"""
        for db_key, conn in self._db_connections.items():
            try:
                conn.close()
            except Exception as e:
                chat_logger.warning(f"关闭数据库连接失败: {db_key} - {str(e)}")
        self._db_connections.clear()
        self._task_db_cache.clear()
        chat_logger.info("所有数据库连接已关闭")


# 全局实例
chat_result_manager = ChatResultManager()