fix: restore backend/python/app/pipeline.py (lost in mass-deletion 5d7f32d)
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.
This commit is contained in:
workbuddy committed
1 parent
da0ab6496c
commit
405e836ee2
1 file changed
+347
@@ -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)
|
||||
Reference in new issue
Block a user