Tolan/memory_summaries/summarize_memory.py

"""对话记忆总结服务（定时扫描）

从数据库读取需要总结的用户对话，使用 LLM 生成总结并更新记忆。

Prerequisites:
- Python >= 3.7
- 配置文件：summarize.conf（与脚本同目录）
  - 必需配置：[database] 下的数据库配置
  - 必需配置：[ai] 下的 DashScope API Key

运行: python summarize_memory.py
"""

import asyncio
import configparser
import os
import sys
import logging
from datetime import datetime, timedelta
from pathlib import Path
from typing import List, Dict, Optional

from openai import OpenAI
from sqlalchemy import select, func

# 加载配置文件
config_file = Path(__file__).parent / "summarize.conf"
config = configparser.ConfigParser()
if config_file.exists():
    config.read(config_file, encoding='utf-8')
else:
    raise FileNotFoundError(f"配置文件不存在: {config_file}")

# 导入数据库相关模块（需要从 backend 目录导入）
backend_dir = Path(__file__).parent.parent / "backend"
sys.path.insert(0, str(backend_dir))
from database import DatabaseManager
from models_db import User, Message, ConversationSummary

# 导入记忆服务
sys.path.insert(0, str(Path(__file__).parent))
from memory_service import (
    load_memory,
    save_memory,
    extract_and_update_memory_session_level,
    build_session_snippet,
)

logger = logging.getLogger(__name__)

# ------- 从配置文件读取参数 -------
# 数据库配置
MYSQL_HOST = config.get("database", "mysql_host", fallback="").strip()
MYSQL_PORT = config.get("database", "mysql_port", fallback="").strip()
MYSQL_USER = config.get("database", "mysql_user", fallback="").strip()
MYSQL_PASSWORD = config.get("database", "mysql_password", fallback="").strip()
MYSQL_DATABASE = config.get("database", "mysql_database", fallback="").strip()

# AI配置
DASHSCOPE_API_KEY = config.get("ai", "dashscope_api_key", fallback="").strip()
AI_MODEL = config.get("ai", "model", fallback="qwen-flash-2025-07-28").strip()

# 服务配置
POLL_INTERVAL = float(config.get("service", "poll_interval", fallback="60.0"))
SUMMARY_DELAY_HOURS = float(config.get("service", "summary_delay_hours", fallback="1.0"))

# 初始化 OpenAI 客户端（DashScope兼容模式）
client = None
if DASHSCOPE_API_KEY:
    client = OpenAI(
        api_key=DASHSCOPE_API_KEY,
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )


async def find_users_need_summary(db_manager: DatabaseManager) -> List[Dict]:
    """
    查找需要总结的用户
    
    条件：
    1. last_active_at 距离现在 >= SUMMARY_DELAY_HOURS 小时
    2. total_rounds > 上一次总结的 end_round（有未总结的对话）
    
    Returns:
        需要总结的用户列表，每个元素包含 user_id, last_summary_end_round, total_rounds
    """
    try:
        cutoff_time = datetime.now() - timedelta(hours=SUMMARY_DELAY_HOURS)
        
        async with db_manager.get_session() as session:
            # 查询符合条件的用户（last_active_at < cutoff_time，确保至少1小时过去）
            users_query = await session.execute(
                select(User)
                .where(
                    User.last_active_at.isnot(None),
                    User.last_active_at < cutoff_time,  # ✅ 使用 < 确保至少1小时过去
                    User.total_rounds > 0
                )
            )
            users = users_query.scalars().all()
            
            if not users:
                return []
            
            result = []
            
            for user in users:
                user_id = user.user_id
                total_rounds = user.total_rounds or 0
                
                # 查询该用户的最新总结
                summary_query = await session.execute(
                    select(ConversationSummary)
                    .where(ConversationSummary.user_id == user_id)
                    .order_by(ConversationSummary.end_round.desc())
                    .limit(1)
                )
                latest_summary = summary_query.scalar_one_or_none()
                
                last_summary_end_round = 0
                if latest_summary:
                    last_summary_end_round = latest_summary.end_round
                
                # 如果有未总结的对话
                if total_rounds > last_summary_end_round:
                    result.append({
                        "user_id": user_id,
                        "last_summary_end_round": last_summary_end_round,
                        "total_rounds": total_rounds,
                        "last_active_at": user.last_active_at,
                    })
            
            return result
        
    except Exception as e:
        logger.error(f"❌ 查找需要总结的用户失败: {e}", exc_info=True)
        return []


async def get_messages_for_summary(
    db_manager: DatabaseManager,
    user_id: str,
    start_round: int,
    end_round: int
) -> List[Dict]:
    """
    获取需要总结的消息
    
    Args:
        db_manager: 数据库管理器
        user_id: 用户ID
        start_round: 起始轮次（不包含，即 start_round + 1）
        end_round: 结束轮次（包含）
    
    Returns:
        消息列表，格式: [{"role": "user|assistant", "content": "...", "round_number": 1}, ...]
    """
    try:
        async with db_manager.get_session() as session:
            messages_query = await session.execute(
                select(Message)
                .where(
                    Message.user_id == user_id,
                    Message.round_number > start_round,
                    Message.round_number <= end_round
                )
                .order_by(Message.round_number.asc(), Message.id.asc())
            )
            messages_list = messages_query.scalars().all()
            
            messages = [
                {
                    "role": msg.role,
                    "content": msg.content,
                    "round_number": msg.round_number,
                }
                for msg in messages_list
            ]
            
            return messages
            
    except Exception as e:
        logger.error(f"❌ 获取消息失败: user_id={user_id}, error={e}", exc_info=True)
        return []


async def process_user_summary(
    user_id: str,
    last_summary_end_round: int,
    total_rounds: int,
    db_manager: DatabaseManager
) -> bool:
    """
    处理单个用户的总结
    
    Args:
        user_id: 用户ID
        last_summary_end_round: 上一次总结的结束轮次
        total_rounds: 用户当前总轮次
        db_manager: 数据库管理器
    
    Returns:
        是否成功
    """
    try:
        start_round = last_summary_end_round
        end_round = total_rounds
        
        logger.info(f"🔄 开始处理用户总结: user_id={user_id}, rounds={start_round + 1}-{end_round}")
        
        # 1. 获取需要总结的消息
        messages = await get_messages_for_summary(
            db_manager, user_id, start_round, end_round
        )
        
        if not messages:
            logger.warning(f"⚠️ 没有消息需要总结: user_id={user_id}, rounds={start_round + 1}-{end_round}")
            return False
        
        logger.info(f"📊 获取到 {len(messages)} 条消息需要总结")
        
        # 2. 加载现有记忆
        current_memory = load_memory(user_id)
        
        # 3. 构建消息列表（用于记忆提取）
        # 需要转换为 extract_and_update_memory_session_level 需要的格式
        messages_for_memory = [
            {"role": msg["role"], "content": msg["content"]}
            for msg in messages
        ]
        
        # 4. 调用记忆提取和更新
        if not client:
            logger.error("❌ OpenAI 客户端未初始化，请检查配置文件中的 dashscope_api_key")
            return False
        
        new_memory = extract_and_update_memory_session_level(
            client=client,
            model=AI_MODEL,
            user_id=user_id,
            messages=messages_for_memory,
            current_memory=current_memory,
            max_turns=80,
        )
        
        # 5. 获取 meta.memory_summary（增量摘要）
        memory_summary = new_memory.get("meta.memory_summary", "")
        
        # 6. 保存到 conversation_summaries 表（完整总结）
        # 使用 memory_summary 作为总结文本（这是增量摘要）
        summary_id = f"summary_{user_id}_{datetime.now().strftime('%Y%m%d%H%M%S')}"
        
        async with db_manager.get_session() as session:
            # 检查是否已存在相同范围的总结
            existing_query = await session.execute(
                select(ConversationSummary)
                .where(
                    ConversationSummary.user_id == user_id,
                    ConversationSummary.start_round == start_round + 1,
                    ConversationSummary.end_round == end_round
                )
            )
            existing = existing_query.scalar_one_or_none()
            
            if existing:
                logger.info(f"⏭️ 总结已存在，跳过: user_id={user_id}, rounds={start_round + 1}-{end_round}")
                return True
            
            # 创建新总结记录
            summary = ConversationSummary(
                summary_id=summary_id,
                user_id=user_id,
                start_round=start_round + 1,
                end_round=end_round,
                summary_text=memory_summary or f"对话总结（第{start_round + 1}-{end_round}轮）",
                summary_type="auto",
                message_count=len(messages),
                token_count=len(memory_summary) // 2 if memory_summary else 0
            )
            session.add(summary)
            await session.commit()
        
        logger.info(f"✅ 用户总结完成: user_id={user_id}, rounds={start_round + 1}-{end_round}, messages={len(messages)}")
        return True
        
    except Exception as e:
        logger.error(f"❌ 处理用户总结失败: user_id={user_id}, error={e}", exc_info=True)
        return False


async def process_summaries(db_manager: DatabaseManager):
    """扫描并处理需要总结的用户"""
    if not DASHSCOPE_API_KEY:
        logger.error("❌ DashScope API Key 配置缺失，请在 summarize.conf 文件中设置 dashscope_api_key")
        return 0, 0
    
    try:
        # 查找需要总结的用户
        users_to_summarize = await find_users_need_summary(db_manager)
        
        if not users_to_summarize:
            return 0, 0
        
        logger.info(f"📋 找到 {len(users_to_summarize)} 个用户需要总结")
        
        success_count = 0
        fail_count = 0
        
        for user_info in users_to_summarize:
            user_id = user_info["user_id"]
            last_summary_end_round = user_info["last_summary_end_round"]
            total_rounds = user_info["total_rounds"]
            
            success = await process_user_summary(
                user_id=user_id,
                last_summary_end_round=last_summary_end_round,
                total_rounds=total_rounds,
                db_manager=db_manager
            )
            
            if success:
                success_count += 1
            else:
                fail_count += 1
        
        if success_count > 0 or fail_count > 0:
            logger.info(f"📊 处理完成: 成功 {success_count} 个, 失败 {fail_count} 个")
        
        return success_count, fail_count
        
    except Exception as e:
        logger.error(f"❌ 处理总结时出错: {e}", exc_info=True)
        return 0, 0


async def main_loop():
    """主循环：定时扫描并处理总结"""
    # 设置数据库配置到环境变量（DatabaseManager 从环境变量读取）
    if MYSQL_HOST:
        os.environ["MYSQL_HOST"] = MYSQL_HOST
    if MYSQL_PORT:
        os.environ["MYSQL_PORT"] = MYSQL_PORT
    if MYSQL_USER:
        os.environ["MYSQL_USER"] = MYSQL_USER
    if MYSQL_PASSWORD:
        os.environ["MYSQL_PASSWORD"] = MYSQL_PASSWORD
    if MYSQL_DATABASE:
        os.environ["MYSQL_DATABASE"] = MYSQL_DATABASE
    
    db_manager = DatabaseManager()
    
    try:
        # 初始化数据库连接
        await db_manager.connect()
        logger.info("✅ 数据库连接成功")
        
        logger.info(f"🚀 开始记忆总结服务（轮询间隔: {POLL_INTERVAL}秒，延迟: {SUMMARY_DELAY_HOURS}小时）...")
        
        # 主循环：持续轮询
        while True:
            try:
                await process_summaries(db_manager)
            except Exception as e:
                logger.error(f"❌ 处理总结时出错: {e}", exc_info=True)
            
            # 等待下次轮询
            await asyncio.sleep(POLL_INTERVAL)
            
    except KeyboardInterrupt:
        logger.info("🛑 收到中断信号，正在关闭...")
    except Exception as e:
        logger.error(f"❌ 服务异常: {e}", exc_info=True)
        sys.exit(1)
    finally:
        await db_manager.disconnect()
        logger.info("✅ 服务已关闭")


def setup_logging():
    """配置日志：正常日志输出到 summarize_memory.log，错误日志输出到 summarize_memory_error.log"""
    log_dir = Path(__file__).parent
    log_file = log_dir / "logs" / "summarize_memory.log"
    error_log_file = log_dir / "logs" / "summarize_memory_error.log"
    
    # 确保日志目录存在
    log_dir.mkdir(parents=True, exist_ok=True)
    (log_dir / "logs").mkdir(parents=True, exist_ok=True)
    
    # 创建格式化器
    formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    
    # 配置根日志记录器
    root_logger = logging.getLogger()
    root_logger.setLevel(logging.INFO)
    
    # 清除已有的处理器
    root_logger.handlers.clear()
    
    # 1. 控制台处理器（INFO及以上）
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(formatter)
    root_logger.addHandler(console_handler)
    
    # 2. 文件处理器（所有日志，INFO及以上）
    file_handler = logging.FileHandler(log_file, encoding='utf-8')
    file_handler.setLevel(logging.INFO)
    file_handler.setFormatter(formatter)
    root_logger.addHandler(file_handler)
    
    # 3. 错误日志文件处理器（ERROR及以上）
    error_file_handler = logging.FileHandler(error_log_file, encoding='utf-8')
    error_file_handler.setLevel(logging.ERROR)
    error_file_handler.setFormatter(formatter)
    root_logger.addHandler(error_file_handler)
    
    # 4. 设置 SQLAlchemy 连接池日志级别为 WARNING（避免 DEBUG 日志）
    logging.getLogger("sqlalchemy.pool").setLevel(logging.WARNING)
    logging.getLogger("sqlalchemy.engine").setLevel(logging.WARNING)
    
    logger.info(f"✅ 日志配置完成：正常日志 -> {log_file}, 错误日志 -> {error_log_file}")


if __name__ == "__main__":
    setup_logging()
    asyncio.run(main_loop())