From 3b2ac891b2a30f81ad7ff3c3b7e2c2510b71bfa2 Mon Sep 17 00:00:00 2001 From: Soulter <905617992@qq.com> Date: Mon, 3 Feb 2025 18:38:02 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E9=99=90=E6=B5=81?= =?UTF-8?q?=E5=99=A8=E4=B8=8D=E5=8F=AF=E7=94=A8=E7=9A=84=E9=97=AE=E9=A2=98?= =?UTF-8?q?=20#263?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- astrbot/core/pipeline/__init__.py | 4 +++- astrbot/core/pipeline/rate_limit_check/stage.py | 7 ++++--- 2 files changed, 7 insertions(+), 4 deletions(-) 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))