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)