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:
        """
        获取指定轮次的上下文 - v0.3
        
        逻辑：
        - 如果 current_round <= 20: 返回前N轮所有消息
        - 如果 current_round > 20: 
            * 最近20轮消息 (round_number in [current_round-20, current_round-1])
            * 历史总结（从 conversation_summaries 表 + JSON memory）
        
        Args:
            user_id: 用户ID
            current_round: 当前轮次
        
        Returns:
            {
                "summary": str | None,  # 历史总结（数据库 + JSON memory）
                "recent_messages": List[Dict],  # 最近20轮消息
                "total_tokens": int  # 总token数
            }
        """
        
        try:
            # 🔥 确保Redis连接可用
            await self._ensure_redis()
            
            if current_round <= 20:
                # 前20轮：直接返回所有历史消息
                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:
                # 超过20轮：最新总结 + 最近20轮
                summary_end_round = current_round - 1
                
                logger.info(f"🔍 查询总结: user={user_id}, current_round={current_round}, summary_end_round={summary_end_round}")
                
                # 1. 获取最新的总结（数据库 + JSON memory）
                summary = await self._get_latest_summary(user_id, summary_end_round)
                
                # 2. 获取最近20轮消息
                recent_messages = await self._get_messages_from_cache_or_db(
                    user_id,
                    start_round=current_round - 20,
                    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]:
        """
        获取最新的总结（v0.3：数据库总结 + JSON memory摘要）
        
        Args:
            user_id: 用户ID
            up_to_round: 总结到的轮次（包含）
        
        Returns:
            合并后的总结文本，如果没有则返回None
        """
        
        try:
            # 1. 从数据库获取最新总结
            from database import db_manager
            from models_db import ConversationSummary
            
            logger.debug(f"🔍 从MySQL查询最新总结: user={user_id}, up_to={up_to_round}")
            
            db_summary = None
            async with db_manager.get_session() as session:
                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_record = result_query.scalar_one_or_none()
                
                if summary_record:
                    db_summary = summary_record.summary_text
                    logger.info(f"✅ 获取数据库总结: user={user_id}, end_round={summary_record.end_round}")
            
            # 2. 从 JSON 文件获取记忆摘要
            json_summary = None
            try:
                import sys
                from pathlib import Path
                project_root = Path(__file__).parent.parent.parent
                sys.path.insert(0, str(project_root))
                from memory_summaries.memory_service import load_memory
                
                memory = load_memory(user_id)
                json_summary = memory.get("meta.memory_summary", "")
                if isinstance(json_summary, str):
                    json_summary = json_summary.strip()
                    if json_summary and json_summary != "(none)":
                        logger.debug(f"✅ 获取JSON记忆摘要: user={user_id}")
                else:
                    json_summary = None
            except Exception as e:
                logger.debug(f"加载JSON记忆失败: user={user_id}, error={e}")
                json_summary = None
            
            # 3. 合并总结（优先使用数据库总结，JSON摘要作为补充）
            if db_summary and json_summary:
                # 合并两者
                combined = f"{db_summary}\n\n[JSON Memory Summary]\n{json_summary}"
                return combined
            elif db_summary:
                return db_summary
            elif json_summary:
                return json_summary
            else:
                logger.warning(f"⚠️ 未找到总结: user={user_id}, up_to={up_to_round}")
            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()