qihuaiflow/backend/python/app/pipeline.py

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