""" Redis事件总线 - 利用已有Redis实现消息队列 参照AstrBot EventBus设计，但使用Redis替代asyncio.Queue """ import json import logging import threading import time from typing import Callable import redis logger = logging.getLogger(__name__) # Redis配置（与项目已有配置一致） try: from .config import REDIS_HOST, REDIS_PORT, REDIS_PASSWORD except ImportError: REDIS_HOST = "localhost" REDIS_PORT = 6379 REDIS_PASSWORD = "pwd2026" def _get_redis(): """获取Redis连接""" return redis.Redis( host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True, ) # Redis键名前缀 QUEUE_KEY = "aiflow:event_queue" PROCESSING_KEY = "aiflow:event_processing" class MessageEvent: """统一消息事件""" def __init__(self, platform: str, channel_id: str, content: str, sender_id: str = "", sender_name: str = "", message_id: str = "", group_id: str = None, content_type: str = "text", raw: dict = None): self.platform = platform self.channel_id = channel_id self.content = content self.sender_id = sender_id self.sender_name = sender_name self.message_id = message_id self.group_id = group_id self.content_type = content_type self.raw = raw or {} self.timestamp = int(time.time()) def to_dict(self) -> dict: return { "platform": self.platform, "channel_id": self.channel_id, "content": self.content, "sender_id": self.sender_id, "sender_name": self.sender_name, "message_id": self.message_id, "group_id": self.group_id, "content_type": self.content_type, "raw": self.raw, "timestamp": self.timestamp, } @classmethod def from_dict(cls, data: dict) -> "MessageEvent": return cls( platform=data.get("platform", ""), channel_id=data.get("channel_id", ""), content=data.get("content", ""), sender_id=data.get("sender_id", ""), sender_name=data.get("sender_name", ""), message_id=data.get("message_id", ""), group_id=data.get("group_id"), content_type=data.get("content_type", "text"), raw=data.get("raw", {}), ) class RedisEventBus: """Redis事件总线 - 生产者-消费者模式""" def __init__(self): self._handlers: dict[str, list[Callable]] = {} self._consumers: list[threading.Thread] = [] self._running = False def publish(self, event: MessageEvent) -> bool: """发布事件到Redis队列""" try: r = _get_redis() event_data = json.dumps(event.to_dict(), ensure_ascii=False) r.lpush(QUEUE_KEY, event_data) logger.debug(f"[EventBus] 发布事件: {event.platform}:{event.channel_id}") return True except Exception as e: logger.error(f"[EventBus] 发布事件失败: {e}") return False def subscribe(self, platform: str, handler: Callable[[MessageEvent], None]): """订阅特定平台的事件""" if platform not in self._handlers: self._handlers[platform] = [] self._handlers[platform].append(handler) logger.info(f"[EventBus] 注册处理器: {platform}") def _consume_loop(self): """消费者循环：从Redis队列取出事件并分发给处理器""" r = _get_redis() while self._running: try: # BRPOP 阻塞等待新事件 result = r.brpop(QUEUE_KEY, timeout=2) if not result: continue _, event_data = result event_dict = json.loads(event_data) event = MessageEvent.from_dict(event_dict) # 分发给对应平台的处理器 platform = event.platform handlers = self._handlers.get(platform, []) + self._handlers.get("*", []) for handler in handlers: try: handler(event) except Exception as e: logger.error(f"[EventBus] 处理器异常(platform={platform}): {e}") except Exception as e: if self._running: logger.error(f"[EventBus] 消费循环异常: {e}") time.sleep(1) def start(self, num_consumers: int = 2): """启动消费者线程""" if self._running: return self._running = True for i in range(num_consumers): t = threading.Thread(target=self._consume_loop, name=f"EventBus-Consumer-{i}", daemon=True) t.start() self._consumers.append(t) logger.info(f"[EventBus] 已启动 {num_consumers} 个消费者") def stop(self): """停止所有消费者""" self._running = False for t in self._consumers: t.join(timeout=5) self._consumers.clear() logger.info("[EventBus] 已停止所有消费者") def queue_size(self) -> int: """获取队列长度""" try: r = _get_redis() return r.llen(QUEUE_KEY) except Exception: return -1 # 全局单例 _event_bus_instance = None def get_event_bus() -> RedisEventBus: """获取全局事件总线单例""" global _event_bus_instance if _event_bus_instance is None: _event_bus_instance = RedisEventBus() return _event_bus_instance # ============================================================ # 单元测试 # ============================================================ if __name__ == "__main__": print("=" * 60) print("Redis事件总线 单元测试") print("=" * 60) # 测试MessageEvent print("\n[测试1] MessageEvent 创建和序列化") event = MessageEvent( platform="feishu", channel_id="ch_001", content="你好，帮我查询天气", sender_id="user_001", sender_name="张三", message_id="msg_001", group_id="grp_001", content_type="text", raw={"custom": "data"}, ) event_dict = event.to_dict() assert event_dict["platform"] == "feishu" assert event_dict["channel_id"] == "ch_001" assert event_dict["content"] == "你好，帮我查询天气" assert event_dict["sender_id"] == "user_001" assert event_dict["sender_name"] == "张三" assert event_dict["message_id"] == "msg_001" assert event_dict["group_id"] == "grp_001" assert event_dict["content_type"] == "text" assert event_dict["raw"] == {"custom": "data"} assert "timestamp" in event_dict print(" PASS: MessageEvent序列化正确") # 测试MessageEvent反序列化 print("\n[测试2] MessageEvent 反序列化") event2 = MessageEvent.from_dict(event_dict) assert event2.platform == "feishu" assert event2.channel_id == "ch_001" assert event2.content == "你好，帮我查询天气" assert event2.sender_id == "user_001" assert event2.sender_name == "张三" assert event2.message_id == "msg_001" assert event2.group_id == "grp_001" assert event2.content_type == "text" assert event2.raw == {"custom": "data"} print(" PASS: MessageEvent反序列化正确") # 测试RedisEventBus发布事件 print("\n[测试3] RedisEventBus 发布事件") bus = RedisEventBus() result = bus.publish(event) assert result is True, "事件发布应该成功" print(" PASS: 事件发布成功") # 测试队列长度 print("\n[测试4] 队列长度查询") size = bus.queue_size() assert size >= 1, f"队列中应该至少有1条消息，实际: {size}" print(f" PASS: 队列长度={size}") # 测试事件消费 print("\n[测试5] 事件消费") received_events = [] def handler(evt: MessageEvent): received_events.append(evt) bus.subscribe("feishu", handler) bus.start(num_consumers=1) # 发布一个事件并等待消费 test_event = MessageEvent( platform="feishu", channel_id="ch_test", content="测试消息", sender_id="user_test", ) bus.publish(test_event) # 等待消费者处理 time.sleep(1) assert len(received_events) >= 1, f"应该接收到事件，实际: {len(received_events)}" if received_events: assert received_events[-1].platform == "feishu" assert received_events[-1].channel_id == "ch_test" assert received_events[-1].content == "测试消息" print(f" PASS: 事件被正确消费, platform={received_events[-1].platform}, content={received_events[-1].content}") # 清理 bus.stop() print(" PASS: EventBus正常停止") # 测试全局单例 print("\n[测试6] 全局单例") bus1 = get_event_bus() bus2 = get_event_bus() assert bus1 is bus2, "get_event_bus()应该返回同一个实例" print(" PASS: 全局单例正确") # 清理Redis测试队列 print("\n[测试7] 清理测试数据") r = _get_redis() r.delete(QUEUE_KEY) size = bus.queue_size() assert size == 0, f"清理后队列应该为空，实际: {size}" print(f" PASS: Redis测试队列已清理, 队列长度={size}") print("\n" + "=" * 60) print("所有测试通过!") print("=" * 60)