Tolan/backend/services/context_service.py

"""
上下文管理服务
"""
import json
from typing import List, Dict, Optional
import logging

logger = logging.getLogger(__name__)


class ContextService:
    """上下文管理服务"""
    
    def __init__(self, db_manager=None):
        """
        初始化上下文服务

        Args:
            db_manager: 数据库管理器实例（可选，用于延迟初始化）
        """
        self.db = db_manager
        self.redis = None
        if db_manager:
            # 🔥 延迟获取Redis（在首次使用时）
            pass

    async def set_db_manager(self, db_manager):
        """设置数据库管理器（用于延迟初始化）"""
        self.db = db_manager
        self.redis = await db_manager.get_redis()
    
    async def _ensure_redis(self):
        """确保Redis连接可用"""
        if not self.redis and self.db:
            self.redis = await self.db.get_redis()
    
    async def get_context_for_round(
        self, 
        user_id: str, 
        current_round: int
    ) -> Dict:
        """
        获取指定轮次的上下文
        
        逻辑：
        - 如果 current_round <= 15: 返回前N轮所有消息
        - 如果 current_round > 15: 
            * 最近15轮消息 (round_number in [current_round-15, current_round-1])
            * 所有历史总结 (1 到 current_round-15)
        
        Args:
            user_id: 用户ID
            current_round: 当前轮次
        
        Returns:
            {
                "summary": str | None,  # 最新的一个总结（已包含所有历史信息）
                "recent_messages": List[Dict],  # 最近消息
                "total_tokens": int  # 总token数
            }
        """
        
        try:
            # 🔥 确保Redis连接可用
            await self._ensure_redis()
            
            if current_round <= 15:
                # 前15轮：直接返回所有历史消息
                messages = await self._get_messages_from_cache_or_db(
                    user_id, 
                    start_round=1, 
                    end_round=current_round
                )
                
                logger.info(f"[Context] Round {current_round}: {len(messages)} 条历史消息")
                
                return {
                    "summary": None,
                    "recent_messages": messages,
                    "total_tokens": self._count_tokens(messages=messages)
                }
            
            else:
                # 超过15轮：最新总结 + 最近15轮
                # 🔥 修复：查找 end_round < current_round 的最新总结
                # 因为总结是累积的，最新的总结包含了所有历史信息
                summary_end_round = current_round - 1
                
                logger.info(f"🔍 查询总结: user={user_id}, current_round={current_round}, summary_end_round={summary_end_round}")
                
                # 1. 获取最新的一个总结（已包含所有历史信息）
                summary = await self._get_latest_summary(user_id, summary_end_round)
                
                # 2. 获取最近15轮消息
                recent_messages = await self._get_messages_from_cache_or_db(
                    user_id,
                    start_round=current_round - 15,
                    end_round=current_round
                )
                
                logger.info(
                    f"[Context] Round {current_round}: "
                    f"总结={'有' if summary else '无'}, 消息={len(recent_messages)}条"
                )
                
                return {
                    "summary": summary,
                    "recent_messages": recent_messages,
                    "total_tokens": self._count_tokens(summary=summary, messages=recent_messages)
                }
                
        except Exception as e:
            logger.error(f"❌ 获取上下文失败: user={user_id}, round={current_round}, error={e}")
            # 降级方案：返回空上下文
            return {
                "summary": None,
                "recent_messages": [],
                "total_tokens": 0
            }
    
    async def _get_messages_from_cache_or_db(
        self, 
        user_id: str, 
        start_round: int, 
        end_round: int
    ) -> List[Dict]:
        """
        优先从Redis获取，未命中则从MySQL
        
        Args:
            user_id: 用户ID
            start_round: 起始轮次
            end_round: 结束轮次（不包含）
        
        Returns:
            消息列表
        """
        
        try:
            # 1. 如果请求的是最近15轮以内，尝试从Redis获取
            messages_needed = end_round - start_round
            
            if messages_needed <= 15:
                cached_messages = await self._get_from_redis(user_id, start_round, end_round)
                if cached_messages:
                    logger.debug(f"✅ 从Redis命中: user={user_id}, rounds={start_round}-{end_round}")
                    return cached_messages
            
            # 2. 从MySQL查询
            logger.debug(f"🔍 从MySQL查询: user={user_id}, rounds={start_round}-{end_round}")
            
            query = """
                SELECT role, content, round_number, emotion, created_at
                FROM messages
                WHERE user_id = %(user_id)s 
                  AND round_number >= %(start_round)s 
                  AND round_number < %(end_round)s
                ORDER BY round_number ASC, id ASC
            """
            
            messages = await self.db.execute_query(query, {
                "user_id": user_id,
                "start_round": start_round,
                "end_round": end_round
            })
            
            result = [dict(msg) for msg in messages]
            
            # 3. 如果是最近的消息，回填Redis
            # （这里简化处理，实际应判断end_round是否接近当前轮次）
            if messages_needed <= 15 and len(result) > 0:
                asyncio.create_task(self._cache_recent_messages(user_id, result))
            
            return result
            
        except Exception as e:
            logger.error(f"❌ 获取消息失败: user={user_id}, error={e}")
            return []
    
    async def _get_from_redis(
        self, 
        user_id: str, 
        start_round: int, 
        end_round: int
    ) -> Optional[List[Dict]]:
        """
        从Redis获取消息
        
        Args:
            user_id: 用户ID
            start_round: 起始轮次
            end_round: 结束轮次（不包含）
        
        Returns:
            消息列表或None
        """
        
        try:
            cache_key = f"user:{user_id}:recent_messages"
            
            # 获取所有缓存消息（最多30条）
            cached_data = await self.redis.lrange(cache_key, 0, -1)
            
            if not cached_data:
                return None
            
            # 解析并过滤
            messages = []
            for data in reversed(cached_data):  # Redis是LPUSH，需要反转
                try:
                    msg = json.loads(data)
                    round_num = msg.get('round_number', 0)
                    
                    if start_round <= round_num < end_round:
                        messages.append(msg)
                except:
                    continue
            
            # 按轮次排序
            messages.sort(key=lambda x: x.get('round_number', 0))
            
            # 检查是否完整覆盖请求范围
            if len(messages) > 0:
                rounds_in_cache = set(msg['round_number'] for msg in messages)
                rounds_needed = set(range(start_round, end_round))
                
                # 如果缓存不完整，返回None触发MySQL查询
                if not rounds_needed.issubset(rounds_in_cache):
                    return None
            
            return messages if messages else None
            
        except Exception as e:
            logger.warning(f"⚠️ Redis获取失败: user={user_id}, error={e}")
            return None
    
    async def _cache_recent_messages(self, user_id: str, messages: List[Dict]):
        """缓存最近消息到Redis（异步）"""
        try:
            # 这里简化处理，实际应该更智能地合并现有缓存
            pass
        except Exception as e:
            logger.warning(f"⚠️ 缓存消息失败: {e}")
    
    async def _get_latest_summary(
        self, 
        user_id: str, 
        up_to_round: int
    ) -> Optional[str]:
        """
        获取最新的一个总结（最新总结已包含所有历史信息）
        
        Args:
            user_id: 用户ID
            up_to_round: 总结到的轮次（包含）
        
        Returns:
            最新的总结文本，如果没有则返回None
        """
        
        try:
            # 从MySQL查询最新的一个总结
            logger.debug(f"🔍 从MySQL查询最新总结: user={user_id}, up_to={up_to_round}")
            
            # 先查询所有总结，用于调试
            all_summaries = await self.db.execute_query(
                """
                SELECT end_round, start_round, created_at
                FROM conversation_summaries
                WHERE user_id = %(user_id)s
                ORDER BY end_round DESC
                LIMIT 10
                """,
                {"user_id": user_id}
            )
            
            if all_summaries:
                logger.info(f"📊 数据库中的总结: user={user_id}, 总数={len(all_summaries)}, 列表={[(s['start_round'], s['end_round']) for s in all_summaries]}")
            else:
                logger.warning(f"⚠️ 数据库中没有任何总结: user={user_id}")
                
                # 调试：查询数据库中所有不同的 user_id
                try:
                    all_user_ids = await self.db.execute_query(
                        """
                        SELECT DISTINCT user_id, COUNT(*) as summary_count
                        FROM conversation_summaries
                        GROUP BY user_id
                        LIMIT 10
                        """
                    )
                    if all_user_ids:
                        logger.info(f"🔍 数据库中存在的其他 user_id: {[(u['user_id'], u['summary_count']) for u in all_user_ids]}")
                except Exception as e:
                    logger.debug(f"查询所有 user_id 失败: {e}")
            
            query = """
                SELECT summary_text, end_round
                FROM conversation_summaries
                WHERE user_id = %(user_id)s AND end_round <= %(up_to_round)s
                ORDER BY end_round DESC
                LIMIT 1
            """
            
            result = await self.db.execute_query(query, {
                "user_id": user_id,
                "up_to_round": up_to_round
            })
            
            if result:
                summary_text = result[0]['summary_text']
                end_round = result[0]['end_round']
                logger.info(f"✅ 获取最新总结: user={user_id}, end_round={end_round}")
                return summary_text
            
            logger.warning(f"⚠️ 未找到总结: user={user_id}, up_to={up_to_round}, 数据库中总结数={len(all_summaries) if all_summaries else 0}")
            return None
            
        except Exception as e:
            logger.error(f"❌ 获取总结失败: user={user_id}, error={e}")
            return None
    
    def _count_tokens(
        self, 
        summary: Optional[str] = None, 
        messages: Optional[List[Dict]] = None
    ) -> int:
        """
        估算token数量（简单实现）
        
        Args:
            summary: 总结文本（单个）
            messages: 消息列表
        
        Returns:
            估算的token数
        """
        
        total = 0
        
        if summary:
            # 中文：1字约等于1token，英文：1词约等于1token
            total += len(summary) // 2
        
        if messages:
            for msg in messages:
                content = msg.get('content', '')
                total += len(content) // 2
        
        return total


# 为了避免循环导入，延迟导入
import asyncio


# 全局上下文服务实例
context_service = ContextService()