chat-ai-api/app/db/redis_client.py

import redis
import json
import requests
import time
import random
from typing import Optional, Tuple
from contextlib import contextmanager
from dotenv import load_dotenv
import os

load_dotenv()

REDIS_HOST = os.getenv("REDIS_HOST")
REDIS_PORT = os.getenv("REDIS_PORT")
REDIS_TTL = 1800  # 30分钟

THIRD_PARTY_API = "http://eco.zhongsou.com/eco/user/user.redis.info.groovy"

try:
    redis_client = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, decode_responses=True)
except Exception:
    redis_client = None


def get_app_user(token: str) -> Optional[dict]:
    """
    通过第三方 token 获取用户信息
    先查 Redis 缓存，未命中则调第三方接口并写入缓存

    Args:
        token: App 端传入的第三方 token

    Returns:
        用户信息字典，token 无效或接口异常时返回 None
    """
    cache_key = f"app_user:{token}"

    # 1. 查 Redis 缓存
    try:
        if redis_client:
            cached = redis_client.get(cache_key)
            if cached:
                return json.loads(cached)
    except Exception as e:
        print(f"Redis 读取失败，降级调第三方接口: {e}")

    # 2. 调第三方接口
    try:
        resp = requests.get(THIRD_PARTY_API, params={"token": token}, timeout=5)
        data = resp.json()

        if data.get("head", {}).get("status") != 200:
            return None

        body = data.get("body")
        if not body or not body.get("userId"):
            return None

        user_info = {
            "userId": body["userId"],
            "userName": body["userName"],
            "nickName": body["nickName"],
            "mobile": body["mobile"],
            "imageUrl": body.get("imageUrl", ""),
            "appName": body["appName"],
        }

        # 3. 写入 Redis 缓存
        try:
            if redis_client:
                redis_client.setex(cache_key, REDIS_TTL, json.dumps(user_info))
        except Exception as e:
            print(f"Redis 写入失败: {e}")

        return user_info

    except Exception as e:
        print(f"第三方接口调用失败: {e}")
        return None


# ==================== 群聊AI频率控制 ====================

class RedisReplyController:
    """
    基于Redis的AI回复频率控制器
    - 支持多实例部署
    - 冷却时间控制
    - 每小时配额限制
    - 群消息密度检测
    - 分布式锁
    """

    def __init__(self, redis_client):
        self.redis = redis_client

    def can_reply(self, ai_role: str, priority: str, group_id: str) -> Tuple[bool, str]:
        """
        判断AI是否可以回复

        Args:
            ai_role: AI角色（entrepreneur | femininity）
            priority: 优先级（mentioned | force | optional）
            group_id: 群ID

        Returns:
            (can_reply: bool, reason: str)
        """
        if not self.redis:
            return True, "Redis不可用，直接通过"

        try:
            now = int(time.time())

            # 1. 冷却时间检查
            last_reply_key = f"im:ai:{group_id}:{ai_role}:last_reply"
            last_reply = self.redis.get(last_reply_key)

            if last_reply:
                elapsed = now - int(last_reply)

                # 根据优先级设置冷却时间
                if priority == 'mentioned':
                    # 明确@：随机30-60秒（防止恶意探测规律）
                    cooldown = random.randint(30, 60)
                elif priority == 'force':
                    cooldown = 60
                else:
                    cooldown = 300

                if elapsed < cooldown:
                    return False, f"冷却中，还需 {cooldown - elapsed} 秒"

            # 2. 每小时配额检查（明确@和强制触发不受限）
            if priority not in ['mentioned', 'force']:
                hour_key = f"im:ai:{group_id}:{ai_role}:count:{now // 3600}"
                count = self.redis.get(hour_key)
                count = int(count) if count else 0

                if count >= 5:  # v2.0: 从12次降为5次/小时
                    return False, f"本小时配额已用完 ({count}/5)"

            # 3. 群消息密度检查（明确@和强制触发不受限）
            if priority not in ['mentioned', 'force']:
                density_key = f"im:group:{group_id}:msg_density"
                recent_count = self.redis.llen(density_key)
                if recent_count > 20:
                    return False, f"群内讨论热烈 ({recent_count}条/5分钟)"

            return True, "通过"

        except Exception as e:
            print(f"Redis频率检查失败: {e}")
            return True, "Redis异常，直接通过"

    def record_reply(self, ai_role: str, group_id: str):
        """
        记录AI回复

        Args:
            ai_role: AI角色
            group_id: 群ID
        """
        if not self.redis:
            return

        try:
            now = int(time.time())

            # 记录最后回复时间（1小时有效期）
            last_reply_key = f"im:ai:{group_id}:{ai_role}:last_reply"
            self.redis.set(last_reply_key, now, ex=3600)

            # 增加小时计数
            hour_key = f"im:ai:{group_id}:{ai_role}:count:{now // 3600}"
            self.redis.incr(hour_key)
            self.redis.expire(hour_key, 3600)

        except Exception as e:
            print(f"Redis记录回复失败: {e}")

    def record_group_message(self, group_id: str, msg_seq: int):
        """
        记录群消息（用于密度检测）

        Args:
            group_id: 群ID
            msg_seq: 消息序列号
        """
        if not self.redis:
            return

        try:
            now = int(time.time())
            density_key = f"im:group:{group_id}:msg_density"

            # 添加到列表头部
            self.redis.lpush(density_key, f"{now}:{msg_seq}")
            # 只保留最近50条
            self.redis.ltrim(density_key, 0, 49)
            # 5分钟过期
            self.redis.expire(density_key, 300)

            # 同时更新群最后消息时间
            last_msg_key = f"im:group:{group_id}:last_msg_time"
            self.redis.set(last_msg_key, now, ex=3600)

        except Exception as e:
            print(f"Redis记录群消息失败: {e}")

    def get_last_group_message_time(self, group_id: str) -> int:
        """
        获取群最后消息时间

        Args:
            group_id: 群ID

        Returns:
            时间戳（秒），如果不存在返回0
        """
        if not self.redis:
            return 0

        try:
            key = f"im:group:{group_id}:last_msg_time"
            timestamp = self.redis.get(key)
            return int(timestamp) if timestamp else 0
        except Exception as e:
            print(f"Redis获取群最后消息时间失败: {e}")
            return 0


# ==================== 消息幂等去重 ====================

def is_message_processed(group_id: str, msg_seq: int) -> bool:
    """
    检查消息是否已处理

    Args:
        group_id: 群ID
        msg_seq: 消息序列号

    Returns:
        True: 已处理，False: 未处理
    """
    if not redis_client:
        return False

    try:
        key = f"im:processed:{group_id}:{msg_seq}"
        return redis_client.exists(key) > 0
    except Exception as e:
        print(f"Redis检查消息处理状态失败: {e}")
        return False


def mark_message_processed(group_id: str, msg_seq: int):
    """
    标记消息已处理（10分钟有效期）

    Args:
        group_id: 群ID
        msg_seq: 消息序列号
    """
    if not redis_client:
        return

    try:
        key = f"im:processed:{group_id}:{msg_seq}"
        redis_client.setex(key, 600, "1")  # 10分钟TTL
    except Exception as e:
        print(f"Redis标记消息已处理失败: {e}")


# ==================== 分布式锁 ====================

@contextmanager
def acquire_reply_lock(group_id: str, timeout: int = 5):
    """
    获取回复锁（分布式锁）

    同一时刻只允许一个AI回复，防止并发抢答

    Args:
        group_id: 群ID
        timeout: 锁超时时间（秒）

    Yields:
        bool: True表示获取到锁，False表示未获取到

    Example:
        with acquire_reply_lock(group_id) as locked:
            if not locked:
                return
            # 执行AI回复逻辑
    """
    if not redis_client:
        yield True
        return

    lock_key = f"im:reply_lock:{group_id}"
    lock_value = str(time.time())
    acquired = False

    try:
        # 尝试获取锁（NX: 不存在才设置，EX: 过期时间）
        acquired = redis_client.set(lock_key, lock_value, nx=True, ex=timeout)
        yield bool(acquired)
    finally:
        # 释放锁
        if acquired:
            try:
                redis_client.delete(lock_key)
            except Exception:
                pass  # 锁会自动过期，删除失败不影响


# ==================== 群聊上下文缓存 ====================

def get_cached_group_context(group_id: str) -> Optional[list]:
    """
    从Redis获取缓存的群聊上下文

    Args:
        group_id: 群ID

    Returns:
        群聊上下文列表，未命中返回None
    """
    if not redis_client:
        return None

    cache_key = f"im:group_context:{group_id}"

    try:
        cached = redis_client.get(cache_key)
        if cached:
            return json.loads(cached)
    except Exception as e:
        print(f"Redis获取群上下文缓存失败: {e}")

    return None


def cache_group_context(group_id: str, context: list, ttl: int = 120):
    """
    将群聊上下文缓存到Redis

    Args:
        group_id: 群ID
        context: 群聊上下文列表
        ttl: 缓存时间（秒），默认2分钟
    """
    if not redis_client:
        return

    cache_key = f"im:group_context:{group_id}"

    try:
        redis_client.setex(cache_key, ttl, json.dumps(context, ensure_ascii=False))
    except Exception as e:
        print(f"Redis缓存群上下文失败: {e}")


# ==================== 全局实例 ====================

# 频率控制器实例
reply_controller = RedisReplyController(redis_client)