diff --git a/astrbot/core/pipeline/__init__.py b/astrbot/core/pipeline/__init__.py index bdba27699..c2dd62054 100644 --- a/astrbot/core/pipeline/__init__.py +++ b/astrbot/core/pipeline/__init__.py @@ -2,6 +2,7 @@ from astrbot.core.message.message_event_result import MessageEventResult, EventR from .waking_check.stage import WakingCheckStage from .whitelist_check.stage import WhitelistCheckStage +from .rate_limit_check.stage import RateLimitStage from .content_safety_check.stage import ContentSafetyCheckStage from .preprocess_stage.stage import PreProcessStage from .process_stage.stage import ProcessStage @@ -11,7 +12,7 @@ from .respond.stage import RespondStage STAGES_ORDER = [ "WakingCheckStage", # 检查是否需要唤醒 "WhitelistCheckStage", # 检查是否在群聊/私聊白名单 - "RateLimitCheckStage", # 检查会话是否超过频率限制 + "RateLimitStage", # 检查会话是否超过频率限制 "ContentSafetyCheckStage", # 检查内容安全 "PreProcessStage", # 预处理 "ProcessStage", # 交由 Stars 处理(a.k.a 插件),或者 LLM 调用 @@ -22,6 +23,7 @@ STAGES_ORDER = [ __all__ = [ "WakingCheckStage", "WhitelistCheckStage", + "RateLimitStage", "ContentSafetyCheckStage", "PreProcessStage", "ProcessStage", diff --git a/astrbot/core/pipeline/rate_limit_check/stage.py b/astrbot/core/pipeline/rate_limit_check/stage.py index 8a2bba6ee..4b4ac5b3e 100644 --- a/astrbot/core/pipeline/rate_limit_check/stage.py +++ b/astrbot/core/pipeline/rate_limit_check/stage.py @@ -61,11 +61,12 @@ class RateLimitStage(Stage): stall_duration = (next_window_time - now).total_seconds() match self.rl_strategy: - case RateLimitStrategy.STALL: + case RateLimitStrategy.STALL.value: logger.info(f"会话 {session_id} 被限流。根据限流策略,此会话处理将被暂停 {stall_duration:.2f} 秒。") await asyncio.sleep(stall_duration) - case RateLimitStrategy.DISCARD: - event.set_result(MessageEventResult().message(f"会话 {session_id} 被限流。根据限流策略,此请求已被丢弃,直到您的限额于 {stall_duration:.2f} 秒后重置。")) + case RateLimitStrategy.DISCARD.value: + # event.set_result(MessageEventResult().message(f"会话 {session_id} 被限流。根据限流策略,此请求已被丢弃,直到您的限额于 {stall_duration:.2f} 秒后重置。")) + logger.info(f"会话 {session_id} 被限流。根据限流策略,此请求已被丢弃,直到限额于 {stall_duration:.2f} 秒后重置。") return event.stop_event() self._remove_expired_timestamps(timestamps, now + timedelta(seconds=stall_duration))