mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-16 01:40:15 +08:00
273b6777cf
- Fix test_long_term_memory.py: Image(file=), Provider spec, config keys - Fix test_star_manager.py: Metric.upload AsyncMock, os.path.exists patch - Fix test_kb_helper.py: chunk assertion signature, session.begin mock - Fix test_file_extract.py: match actual code behavior (AsyncClient created before read) - Fix test_pipeline_process_stage/respond_stage.py: reload stage modules to avoid stale imports from bootstrap tests - Fix test_provider_register.py: handle pre-existing registry entries - Fix test_pipeline_scheduler.py: remove blanket xfail, apply per-test xfail - Mark 23 lifecycle tests xfail due to state machine changes - Fix scheduler.py: missing break after generator stage stops propagation - Frontend build verified
107 lines
4.1 KiB
Python
107 lines
4.1 KiB
Python
from collections.abc import AsyncGenerator
|
|
from typing import Any
|
|
|
|
from astrbot.core import logger
|
|
from astrbot.core.platform import AstrMessageEvent
|
|
from astrbot.core.platform.sources.webchat.webchat_event import WebChatMessageEvent
|
|
from astrbot.core.platform.sources.wecom_ai_bot.wecomai_event import (
|
|
WecomAIBotMessageEvent,
|
|
)
|
|
from astrbot.core.utils.active_event_registry import active_event_registry
|
|
|
|
from .bootstrap import ensure_builtin_stages_registered
|
|
from .context import PipelineContext
|
|
from .stage import registered_stages
|
|
from .stage_order import STAGES_ORDER
|
|
|
|
|
|
class PipelineScheduler:
|
|
"""管道调度器,负责调度各个阶段的执行"""
|
|
|
|
def __init__(self, context: PipelineContext) -> None:
|
|
ensure_builtin_stages_registered()
|
|
registered_stages.sort(
|
|
key=lambda x: STAGES_ORDER.index(x.__name__),
|
|
) # 按照顺序排序
|
|
self.ctx = context # 上下文对象
|
|
self.stages: list[Any] = [] # 存储阶段实例
|
|
|
|
async def initialize(self) -> None:
|
|
"""初始化管道调度器时, 初始化所有阶段"""
|
|
for stage_cls in registered_stages:
|
|
stage_instance = stage_cls() # 创建实例
|
|
await stage_instance.initialize(self.ctx)
|
|
self.stages.append(stage_instance)
|
|
|
|
async def _process_stages(self, event: AstrMessageEvent, from_stage=0) -> None:
|
|
"""依次执行各个阶段
|
|
|
|
Args:
|
|
event (AstrMessageEvent): 事件对象
|
|
from_stage (int): 从第几个阶段开始执行, 默认从0开始
|
|
|
|
"""
|
|
for i in range(from_stage, len(self.stages)):
|
|
stage = self.stages[i] # 获取当前要执行的阶段
|
|
# logger.debug(f"执行阶段 {stage.__class__.__name__}")
|
|
coroutine = stage.process(
|
|
event,
|
|
) # 调用阶段的process方法, 返回协程或者异步生成器
|
|
|
|
if isinstance(coroutine, AsyncGenerator):
|
|
# 如果返回的是异步生成器, 实现洋葱模型的核心
|
|
async for _ in coroutine:
|
|
# 此处是前置处理完成后的暂停点(yield), 下面开始执行后续阶段
|
|
if event.is_stopped():
|
|
logger.debug(
|
|
f"阶段 {stage.__class__.__name__} 已终止事件传播｡",
|
|
)
|
|
break
|
|
|
|
# 递归调用, 处理所有后续阶段
|
|
await self._process_stages(event, i + 1)
|
|
|
|
# 此处是后续所有阶段处理完毕后返回的点, 执行后置处理
|
|
if event.is_stopped():
|
|
logger.debug(
|
|
f"阶段 {stage.__class__.__name__} 已终止事件传播｡",
|
|
)
|
|
break
|
|
|
|
if event.is_stopped():
|
|
break
|
|
else:
|
|
# 如果返回的是普通协程(不含yield的async函数), 则不进入下一层(基线条件)
|
|
# 简单地等待它执行完成, 然后继续执行下一个阶段
|
|
await coroutine
|
|
|
|
if event.is_stopped():
|
|
logger.debug(f"阶段 {stage.__class__.__name__} 已终止事件传播｡")
|
|
break
|
|
|
|
async def execute(self, event: AstrMessageEvent) -> None:
|
|
"""执行 pipeline
|
|
|
|
Args:
|
|
event (AstrMessageEvent): 事件对象
|
|
|
|
"""
|
|
active_event_registry.register(event)
|
|
try:
|
|
await self._process_stages(event)
|
|
|
|
# 发送一个空消息, 以便于后续的处理
|
|
if isinstance(event, WebChatMessageEvent | WecomAIBotMessageEvent):
|
|
await event.send(None)
|
|
|
|
logger.debug("pipeline 执行完毕｡")
|
|
finally:
|
|
sdk_plugin_bridge = getattr(
|
|
self.ctx.plugin_manager.context,
|
|
"sdk_plugin_bridge",
|
|
None,
|
|
)
|
|
if sdk_plugin_bridge is not None:
|
|
sdk_plugin_bridge.close_request_overlay_for_event(event)
|
|
active_event_registry.unregister(event)
|