Tolan/backend/services/context_service.py

"""
上下文管理服务（SQLAlchemy ORM 版本）
"""
import json
import asyncio
from typing import Awaitable, Callable, List, Dict, Optional
import logging
import redis.asyncio as aioredis
from sqlalchemy import select

logger = logging.getLogger(__name__)


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

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

    async def set_db_manager(self, db_manager):
        """设置数据库管理器（用于延迟初始化）"""
        self.db = db_manager
        self._redis_provider = db_manager.get_redis
        self.redis = None
    
    async def _ensure_redis(self):
        """确保Redis连接可用"""
        if not self._redis_provider:
            if not self.db:
                raise RuntimeError("Redis 未初始化，缺少数据库管理器")
            self._redis_provider = self.db.get_redis
        redis_client = await self._redis_provider()
        if not redis_client:
            raise RuntimeError("Redis 客户端不可用")
        self.redis = redis_client
        return redis_client
    
    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}")
            
            # 使用 ORM 查询消息
            from database import db_manager
            from models_db import Message
            
            async with db_manager.get_session() as session:
                result_query = await session.execute(
                    select(Message)
                    .where(Message.user_id == user_id)
                    .where(Message.round_number >= start_round)
                    .where(Message.round_number < end_round)
                    .order_by(Message.round_number.asc(), Message.id.asc())
                )
                messages = result_query.scalars().all()
                result = [
                    {
                        "role": msg.role,
                        "content": msg.content,
                        "round_number": msg.round_number,
                        "emotion": msg.emotion,
                        "created_at": msg.created_at.isoformat() if msg.created_at else None,
                    }
                    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:
            # 使用 ORM 查询最新总结
            from database import db_manager
            from models_db import ConversationSummary
            from sqlalchemy import func
            
            logger.debug(f"🔍 从MySQL查询最新总结: user={user_id}, up_to={up_to_round}")
            
            async with db_manager.get_session() as session:
            # 先查询所有总结，用于调试
                all_summaries_query = await session.execute(
                    select(ConversationSummary)
                    .where(ConversationSummary.user_id == user_id)
                    .order_by(ConversationSummary.end_round.desc())
                    .limit(10)
                )
                all_summaries = all_summaries_query.scalars().all()
            
            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}")
                
                # 查询最新的一个总结
                result_query = await session.execute(
                    select(ConversationSummary)
                    .where(ConversationSummary.user_id == user_id)
                    .where(ConversationSummary.end_round <= up_to_round)
                    .order_by(ConversationSummary.end_round.desc())
                    .limit(1)
                )
                summary = result_query.scalar_one_or_none()
                
                if summary:
                    logger.info(f"✅ 获取最新总结: user={user_id}, end_round={summary.end_round}")
                    return summary.summary_text
            
                logger.warning(f"⚠️ 未找到总结: user={user_id}, up_to={up_to_round}, 数据库中总结数={len(all_summaries)}")
            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()