""" Pipeline 消息处理管道 - 参照AstrBot Stage链式设计,同步版兼容FastAPI 消息处理阶段可插拔,链式调用 """ from abc import ABC, abstractmethod from typing import Any import logging import time logger = logging.getLogger(__name__) class StageResult: """阶段处理结果""" def __init__(self, success: bool = True, data: dict = None, error: str = None, stop: bool = False): self.success = success self.data = data or {} self.error = error self.stop = stop # 是否停止后续阶段 class PipelineStage(ABC): """管道阶段基类""" def __init__(self, name: str): self.name = name @abstractmethod def process(self, context: dict) -> StageResult: """处理消息上下文,返回StageResult""" ... # === 具体阶段实现 === class WakingCheckStage(PipelineStage): """唤醒检查:检查是否需要AI响应""" def __init__(self): super().__init__("waking_check") def process(self, context: dict) -> StageResult: message = context.get("message", "") platform = context.get("platform", "") config = context.get("channel_config", {}) bot_prefix = config.get("bot_prefix", "@bot") is_private = context.get("is_private", True) is_at_bot = bot_prefix.lower() in message.lower() if bot_prefix else False if is_private or is_at_bot or message.strip().startswith("/"): return StageResult(success=True, data={"should_respond": True}) # 群聊需要@或唤醒词 if not is_private and not is_at_bot: return StageResult(success=True, data={"should_respond": False}, stop=True) return StageResult(success=True, data={"should_respond": True}) class RateLimitCheckStage(PipelineStage): """速率限制检查""" def __init__(self, max_per_minute: int = 30): super().__init__("rate_limit_check") self.max_per_minute = max_per_minute self._counter = {} def process(self, context: dict) -> StageResult: user_id = context.get("sender_id", "unknown") now = time.time() key = f"{user_id}:{int(now // 60)}" count = self._counter.get(key, 0) if count >= self.max_per_minute: return StageResult(success=False, error="请求过于频繁,请稍后再试", stop=True) self._counter[key] = count + 1 # 清理旧计数器 keys_to_del = [k for k in self._counter if int(k.split(":")[-1]) < int(now // 60) - 1] for k in keys_to_del: del self._counter[k] return StageResult(success=True) class ContentSafetyCheckStage(PipelineStage): """内容安全检查""" def __init__(self): super().__init__("content_safety_check") self._blocked_words = [] # 可从配置加载 def process(self, context: dict) -> StageResult: message = context.get("message", "") for word in self._blocked_words: if word in message: return StageResult(success=False, error="消息包含违规内容", stop=True) return StageResult(success=True) class UnsupportedStage(PipelineStage): """业务逻辑处理阶段基类 - 子类实现具体处理""" def __init__(self): super().__init__("process") def process(self, context: dict) -> StageResult: # 此阶段由channels.py中的实际AI处理逻辑填充 return StageResult(success=True, data={"processed": True}) class ResponseBuildStage(PipelineStage): """响应构建阶段""" def __init__(self): super().__init__("respond") def process(self, context: dict) -> StageResult: ai_reply = context.get("ai_reply", "") if not ai_reply: return StageResult(success=True, data={"reply": ""}) return StageResult(success=True, data={"reply": ai_reply}) # === Pipeline 管道调度器 === class PipelineScheduler: """管道调度器:按顺序执行各个Stage""" def __init__(self, stages: list[PipelineStage] = None): self.stages = stages if stages is not None else self._default_stages() @staticmethod def _default_stages() -> list[PipelineStage]: return [ WakingCheckStage(), RateLimitCheckStage(max_per_minute=30), ContentSafetyCheckStage(), UnsupportedStage(), ResponseBuildStage(), ] def add_stage(self, stage: PipelineStage, position: int = None): """插入新阶段""" if position is not None: self.stages.insert(position, stage) else: self.stages.append(stage) def remove_stage(self, stage_name: str): """移除阶段""" self.stages = [s for s in self.stages if s.name != stage_name] def execute(self, context: dict) -> dict: """执行完整的管道处理""" context["_pipeline_log"] = [] for stage in self.stages: try: result = stage.process(context) context["_pipeline_log"].append({ "stage": stage.name, "success": result.success, "error": result.error, "stop": result.stop, }) if not result.success and result.stop: return { "success": False, "error": result.error, "stage": stage.name, "pipeline_log": context["_pipeline_log"], } if result.stop and result.success: # 正常停止(如waking_check判断不需要响应) return { "success": True, "data": result.data, "stage": stage.name, "pipeline_log": context["_pipeline_log"], } context.update(result.data) except Exception as e: logger.error(f"[Pipeline] Stage {stage.name} 异常: {e}") return { "success": False, "error": f"管道处理异常({stage.name}): {str(e)}", "pipeline_log": context["_pipeline_log"], } return { "success": True, "data": {"reply": context.get("reply", context.get("ai_reply", ""))}, "pipeline_log": context["_pipeline_log"], } # ============================================================ # 单元测试 # ============================================================ if __name__ == "__main__": print("=" * 60) print("Pipeline 管道模块 单元测试") print("=" * 60) # 测试1: Pipeline 创建和执行 print("\n[测试1] Pipeline 创建和执行") pipeline = PipelineScheduler() context = { "message": "你好,帮我查询天气", "platform": "feishu", "sender_id": "user_001", "is_private": True, "channel_config": {"bot_prefix": "@bot"}, "ai_reply": "你好!今天天气晴好,气温22-30℃。", } result = pipeline.execute(context) assert result["success"] is True, f"Pipeline执行应该成功: {result}" assert "pipeline_log" in result, "应该有pipeline_log" # 验证所有5个阶段都执行了 stage_names = [log["stage"] for log in result["pipeline_log"]] assert len(stage_names) == 5, f"应该执行5个阶段,实际执行了{len(stage_names)}个" assert stage_names == ["waking_check", "rate_limit_check", "content_safety_check", "process", "respond"], \ f"阶段顺序不对: {stage_names}" assert result["data"]["reply"] == "你好!今天天气晴好,气温22-30℃。" print(f" PASS: 所有5个阶段按顺序执行,回复={result['data']['reply'][:30]}...") # 测试2: 唤醒检查阶段 - 私聊 print("\n[测试2] 唤醒检查阶段 - 私聊消息") waking = WakingCheckStage() ctx = { "message": "你好", "platform": "telegram", "is_private": True, "channel_config": {"bot_prefix": "@bot"}, } result = waking.process(ctx) assert result.success is True, "私聊应该唤醒" assert result.data["should_respond"] is True, "私聊应该响应" print(f" PASS: 私聊消息正确唤醒, should_respond={result.data['should_respond']}") # 测试3: 唤醒检查阶段 - 群聊 @bot print("\n[测试3] 唤醒检查阶段 - 群聊@bot") ctx = { "message": "@bot 帮我查一下数据", "platform": "wecom", "is_private": False, "channel_config": {"bot_prefix": "@bot"}, } result = waking.process(ctx) assert result.success is True, "群聊@bot应该唤醒" assert result.data["should_respond"] is True, "群聊@bot应该响应" print(f" PASS: 群聊@bot正确唤醒, should_respond={result.data['should_respond']}") # 测试4: 唤醒检查阶段 - 群聊未@bot (应停止) print("\n[测试4] 唤醒检查阶段 - 群聊未@bot(应停止管道)") ctx = { "message": "大家今天天气不错", "platform": "wecom", "is_private": False, "channel_config": {"bot_prefix": "@bot"}, } result = waking.process(ctx) assert result.success is True, "处理应该成功" assert result.data["should_respond"] is False, "群聊未@bot不应该响应" assert result.stop is True, "应该停止管道" print(f" PASS: 群聊未@bot正确停止管道, should_respond={result.data['should_respond']}, stop={result.stop}") # 测试5: 速率限制 print("\n[测试5] 速率限制检查") rate_limiter = RateLimitCheckStage(max_per_minute=3) # 前3次应该成功 for i in range(3): ctx = {"sender_id": "user_test"} result = rate_limiter.process(ctx) assert result.success is True, f"第{i+1}次请求应该通过速率限制" print(" PASS: 前3次请求通过速率限制") # 第4次应该失败 ctx = {"sender_id": "user_test"} result = rate_limiter.process(ctx) assert result.success is False, "第4次请求应该被限流" assert "频繁" in result.error, f"错误信息应包含'频繁': {result.error}" assert result.stop is True, "限流应该停止管道" print(f" PASS: 第4次请求被限流, error={result.error}") # 测试6: 完整管道链 - 群聊未唤醒场景 print("\n[测试6] 完整管道链 - 群聊未唤醒(应被waking_check停止)") pipeline = PipelineScheduler() context = { "message": "今天天气真好啊", "platform": "wecom", "sender_id": "user_002", "is_private": False, "channel_config": {"bot_prefix": "@bot"}, } result = pipeline.execute(context) assert result["success"] is True, "管道应该正常完成" assert result["stage"] == "waking_check", f"应在waking_check阶段停止,实际: {result['stage']}" assert len(result["pipeline_log"]) == 1, f"应该只执行了1个阶段,实际: {len(result['pipeline_log'])}" print(f" PASS: 群聊未唤醒场景正确在waking_check阶段停止") # 测试7: 内容安全检查 - 无违规词 print("\n[测试7] 内容安全检查 - 无违规词") safety = ContentSafetyCheckStage() ctx = {"message": "正常的聊天内容"} result = safety.process(ctx) assert result.success is True, "正常内容应该通过检查" print(f" PASS: 正常内容通过安全检查") # 测试8: 异常处理 - Stage抛出异常 print("\n[测试8] Pipeline异常处理") class BrokenStage(PipelineStage): def __init__(self): super().__init__("broken") def process(self, context): raise RuntimeError("模拟阶段异常") pipeline2 = PipelineScheduler(stages=[BrokenStage()]) result = pipeline2.execute({"message": "test"}) assert result["success"] is False, "异常应该导致管道失败" assert "broken" in result["error"], f"错误信息应包含阶段名: {result['error']}" print(f" PASS: 阶段异常被正确捕获, error={result['error']}") # 测试9: Pipeline 添加和移除阶段 print("\n[测试9] Pipeline 阶段增删") pipeline3 = PipelineScheduler(stages=[]) assert len(pipeline3.stages) == 0, "初始应无阶段" test_stage = WakingCheckStage() pipeline3.add_stage(test_stage) assert len(pipeline3.stages) == 1, "添加后应有1个阶段" pipeline3.add_stage(RateLimitCheckStage(), position=0) assert pipeline3.stages[0].name == "rate_limit_check", "插入到位置0应为rate_limit_check" pipeline3.remove_stage("waking_check") assert len(pipeline3.stages) == 1, "移除后应有1个阶段" assert pipeline3.stages[0].name == "rate_limit_check", "剩余应为rate_limit_check" print(" PASS: 阶段增删操作正确") print("\n" + "=" * 60) print("所有测试通过!") print("=" * 60)