diff --git a/backend/python/app/pipeline.py b/backend/python/app/pipeline.py new file mode 100644 index 0000000..7d7f9a9 --- /dev/null +++ b/backend/python/app/pipeline.py @@ -0,0 +1,347 @@ +""" +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)