Merge pull request #414 from RockChinQ/detailed-rate-limit

[Feat] 速度限制支持细化到单个人或群

config-template.py

@@ -247,12 +247,35 @@ session_expire_time = 1200
 # 单会话内每分钟可进行的对话次数
 # 若不需要限速，可以设置为一个很大的值
 # 默认值60次，基本上不会触发限速
-rate_limitation = 60
+#
+# 若要设置针对某特定群的限速，请使用如下格式：
+# {
+#   "group_<群号>": 60,
+#   "default": 60,
+# }
+# 若要设置针对某特定用户私聊的限速，请使用如下格式：
+# {
+#   "person_<用户QQ>": 60,
+#   "default": 60,
+# }
+# 同时设置多个群和私聊的限速，示例：
+# {
+#   "group_12345678": 60,
+#   "group_87654321": 60,
+#   "person_234567890": 60,
+#   "person_345678901": 60,
+#   "default": 60,
+# }
+#
+# 注意: 未指定的都使用default的限速值，default不可删除
+rate_limitation = {
+    "default": 60,
+}
 
 # 会话限速策略
 # - "wait": 每次对话获取到回复时，等待一定时间再发送回复，保证其不会超过限速均值
 # - "drop": 此分钟内，若对话次数超过限速次数，则丢弃之后的对话，每自然分钟重置
-rate_limit_strategy = "wait"
+rate_limit_strategy = "drop"
 
 # 是否在启动时进行依赖库更新
 upgrade_dependencies = True

override-all.json

@@ -67,8 +67,10 @@
     "admin_pool_num": 2,
     "user_pool_num": 6,
     "session_expire_time": 1200,
-    "rate_limitation": 60,
+    "rate_limitation": {
-    "rate_limit_strategy": "wait",
+        "default": 60
+    },
+    "rate_limit_strategy": "drop",
     "upgrade_dependencies": true,
     "report_usage": true,
     "logging_level": 20

pkg/qqbot/ratelimit.py

@@ -10,6 +10,20 @@ __crt_minute_usage__ = {}
 __timer_thr__: threading.Thread = None
 
 
+def get_limitation(session_name: str) -> int:
+    """获取会话的限制次数"""
+    import config
+
+    if type(config.rate_limitation) == dict:
+        # 如果被指定了
+        if session_name in config.rate_limitation:
+            return config.rate_limitation[session_name]
+        else:
+            return config.rate_limitation["default"]
+    elif type(config.rate_limitation) == int:
+        return config.rate_limitation
+
+
 def add_usage(session_name: str):
     """增加会话的对话次数"""
     global __crt_minute_usage__
@@ -56,9 +70,7 @@ def get_rest_wait_time(session_name: str, spent: float) -> float:
     """获取会话此回合的剩余等待时间"""
     global __crt_minute_usage__
 
-    import config
+    min_seconds_per_round = 60.0 / get_limitation(session_name)
-
-    min_seconds_per_round = 60.0 / config.rate_limitation
 
     if session_name in __crt_minute_usage__:
         return max(0, min_seconds_per_round - spent)
@@ -70,10 +82,8 @@ def is_reach_limit(session_name: str) -> bool:
     """判断会话是否超过限制"""
     global __crt_minute_usage__
 
-    import config
-
     if session_name in __crt_minute_usage__:
-        return __crt_minute_usage__[session_name] >= config.rate_limitation
+        return __crt_minute_usage__[session_name] >= get_limitation(session_name)
     else:
         return False