life-echo/api/services/redis_service.py

"""
Redis 服务模块：用于会话状态存储和缓存
"""
import os
import json
import logging
from typing import Optional, List, Dict, Any

import redis.asyncio as aioredis

logger = logging.getLogger(__name__)


class RedisService:
    """Redis 服务，用于存储对话历史和状态"""
    
    def __init__(self):
        """初始化 Redis 连接"""
        self.redis_url = os.getenv("REDIS_URL", "redis://localhost:6379/0")
        self._client: Optional[aioredis.Redis] = None
        # 会话过期时间（默认 24 小时）
        self.session_ttl = int(os.getenv("REDIS_SESSION_TTL", "86400"))
    
    async def get_client(self) -> aioredis.Redis:
        """获取 Redis 客户端（延迟初始化）"""
        if self._client is None:
            try:
                self._client = await aioredis.from_url(
                    self.redis_url,
                    encoding="utf-8",
                    decode_responses=True
                )
                # 测试连接
                await self._client.ping()
                logger.info(f"Redis 连接成功: {self.redis_url}")
            except Exception as e:
                logger.error(f"Redis 连接失败: {e}")
                raise
        return self._client
    
    async def close(self):
        """关闭 Redis 连接"""
        if self._client:
            await self._client.close()
            self._client = None
    
    # ==================== 对话历史管理 ====================
    
    def _conversation_key(self, conversation_id: str) -> str:
        """生成对话历史的 Redis key"""
        return f"conversation:history:{conversation_id}"
    
    async def get_conversation_history(self, conversation_id: str) -> List[Dict[str, Any]]:
        """
        获取对话历史
        
        Args:
            conversation_id: 对话 ID
        
        Returns:
            消息列表 [{"role": "human/ai", "content": "..."}]
        """
        try:
            client = await self.get_client()
            key = self._conversation_key(conversation_id)
            data = await client.get(key)
            if data:
                return json.loads(data)
            return []
        except Exception as e:
            logger.error(f"获取对话历史失败: {e}")
            return []
    
    async def add_message(
        self, 
        conversation_id: str, 
        role: str, 
        content: str
    ) -> bool:
        """
        添加消息到对话历史
        
        Args:
            conversation_id: 对话 ID
            role: 角色 ("human" 或 "ai")
            content: 消息内容
        
        Returns:
            是否成功
        """
        try:
            client = await self.get_client()
            key = self._conversation_key(conversation_id)
            
            # 获取现有历史
            history = await self.get_conversation_history(conversation_id)
            
            # 添加新消息
            history.append({"role": role, "content": content})
            
            # 保存回 Redis（带过期时间）
            await client.setex(key, self.session_ttl, json.dumps(history, ensure_ascii=False))
            return True
        except Exception as e:
            logger.error(f"添加消息失败: {e}")
            return False
    
    async def clear_conversation_history(self, conversation_id: str) -> bool:
        """
        清除对话历史
        
        Args:
            conversation_id: 对话 ID
        
        Returns:
            是否成功
        """
        try:
            client = await self.get_client()
            key = self._conversation_key(conversation_id)
            await client.delete(key)
            return True
        except Exception as e:
            logger.error(f"清除对话历史失败: {e}")
            return False
    
    async def extend_session_ttl(self, conversation_id: str) -> bool:
        """
        延长会话过期时间
        
        Args:
            conversation_id: 对话 ID
        
        Returns:
            是否成功
        """
        try:
            client = await self.get_client()
            key = self._conversation_key(conversation_id)
            await client.expire(key, self.session_ttl)
            return True
        except Exception as e:
            logger.error(f"延长会话TTL失败: {e}")
            return False
    
    # ==================== 通用缓存方法 ====================
    
    async def set_cache(self, key: str, value: Any, ttl: Optional[int] = None) -> bool:
        """设置缓存"""
        try:
            client = await self.get_client()
            data = json.dumps(value, ensure_ascii=False) if not isinstance(value, str) else value
            if ttl:
                await client.setex(key, ttl, data)
            else:
                await client.set(key, data)
            return True
        except Exception as e:
            logger.error(f"设置缓存失败: {e}")
            return False
    
    async def get_cache(self, key: str) -> Optional[Any]:
        """获取缓存"""
        try:
            client = await self.get_client()
            data = await client.get(key)
            if data:
                try:
                    return json.loads(data)
                except json.JSONDecodeError:
                    return data
            return None
        except Exception as e:
            logger.error(f"获取缓存失败: {e}")
            return None
    
    async def delete_cache(self, key: str) -> bool:
        """删除缓存"""
        try:
            client = await self.get_client()
            await client.delete(key)
            return True
        except Exception as e:
            logger.error(f"删除缓存失败: {e}")
            return False
    
    def is_available(self) -> bool:
        """检查 Redis 是否可用"""
        return self._client is not None


# 创建全局实例
redis_service = RedisService()