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}")


# ==================== 情绪状态机 ====================

# 群情绪标签及其数值权重（正值=正面，负值=负面）
_MOOD_SCORES = {"excited": 2.0, "calm": 0.0, "low": -1.5, "tense": -2.0}

# 各情绪档位对应的语气指令
_MOOD_TONE = {
    "excited": "群里气氛热烈，回复可以带点感染力，用“确实”“厉害”接住情绪再展开。",
    "calm": "保持你一贯的沉稳企业家语气。",
    "low": "群里气氛偏低落，语气放缓，先共情再给建议，不要上来就分析。",
    "tense": "群里情绪紧张或焦虑，回复要稳，给确定性而非更多问题，避免追问施压。",
}

# 滑动窗口最近5条，权重从旧到新（越新权重越高），归一化后使用
_MOOD_WEIGHTS_BASE = [0.10, 0.10, 0.15, 0.25, 0.40]


def classify_mood(text: str) -> tuple:
    """
    规则法快速情绪分类，零延迟、零调用成本。
    返回 (mood: str, intensity: float)
    """
    if any(w in text for w in ["太难了", "撑不住", "不想干", "亏了", "失败", "崩了", "没希望"]):
        return "low", 0.7
    if any(w in text for w in ["焦虑", "慌", "急", "来不及", "压力"]) or text.count("？") >= 3:
        return "tense", 0.6
    if any(w in text for w in ["牛", "太好了", "成了", "恭喜", "涨了", "爆了"]) or "！" in text:
        return "excited", 0.6
    return "calm", 0.3


def update_group_mood(group_id: str, user_id: str, text: str) -> tuple:
    """
    用新消息更新群情绪状态，写入 Redis（TTL 30 分钟）。
    每条消息以 user_id+text[:20] 作简单去重，避免相同历史消息重复计入。
    返回 (group_mood: str, intensity: float)
    """
    if not redis_client:
        return "calm", 0.0

    key = f"mood:group:{group_id}"
    dedup_key = f"{user_id}:{text[:20]}"

    try:
        raw = redis_client.get(key)
        state = json.loads(raw) if raw else {"recent_signals": [], "seen_keys": []}

        # 去重：同一消息不重复写入
        seen_keys = state.get("seen_keys", [])
        if dedup_key in seen_keys:
            # 消息已处理，直接返回当前状态
            return state.get("group_mood", "calm"), state.get("intensity", 0.0)

        mood, _ = classify_mood(text)
        signals = state.get("recent_signals", [])
        signals.append({"mood": mood})
        signals = signals[-5:]  # 保留最近5条
        seen_keys.append(dedup_key)
        seen_keys = seen_keys[-50:]  # 去重集合最多50条

        # 取与信号数量对应的尾部权重，然后归一化到和为1
        weights_raw = _MOOD_WEIGHTS_BASE[-len(signals):]
        total_w = sum(weights_raw)
        weights = [w / total_w for w in weights_raw]

        scores = [_MOOD_SCORES[s["mood"]] for s in signals]
        avg = sum(w * s for w, s in zip(weights, scores))

        if avg >= 1.0:
            group_mood = "excited"
        elif avg >= -0.5:
            group_mood = "calm"
        elif avg >= -1.5:
            group_mood = "low"
        else:
            group_mood = "tense"

        intensity = round(min(abs(avg) / 2.0, 1.0), 2)

        new_state = {
            "group_mood": group_mood,
            "intensity": intensity,
            "recent_signals": signals,
            "seen_keys": seen_keys,
            "updated_at": time.time(),
        }
        redis_client.setex(key, 1800, json.dumps(new_state, ensure_ascii=False))
        return group_mood, intensity

    except Exception as e:
        print(f"Redis 更新群情绪失败: {e}")
        return "calm", 0.0


def get_group_mood(group_id: str) -> tuple:
    """读取群当前情绪状态，Redis 不可用时回退到 calm。"""
    if not redis_client:
        return "calm", 0.0
    try:
        raw = redis_client.get(f"mood:group:{group_id}")
        if not raw:
            return "calm", 0.0
        state = json.loads(raw)
        return state.get("group_mood", "calm"), state.get("intensity", 0.0)
    except Exception as e:
        print(f"Redis 读取群情绪失败: {e}")
        return "calm", 0.0


def build_mood_block(group_id: str) -> str:
    """构造注入 system prompt 的群氛围说明段。"""
    mood, intensity = get_group_mood(group_id)
    if mood == "calm" and intensity < 0.1:
        return ""  # 平静且强度极低时不注入，减少 prompt 噪声
    tone = _MOOD_TONE.get(mood, "")
    return f"\n\n## 当前群氛围\n{mood}（强度 {intensity}）。{tone}"


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

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