workbuddy
405e836ee2
Single genuinely-missing module after the deletion. Rest of the 44 'missing'
paths were false positives (app/ prefix artifacts in the deletion commit).
Recovery now matches the 1346fbb baseline used by the restored admin
modules and models.py.
348 lines
13 KiB
Python
348 lines
13 KiB
Python
"""
|
|
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)
|