feat: 限制机器人回复的频率 (#72)

pkg/qqbot/ratelimit.py

@@ -0,0 +1,86 @@
+# 限速相关模块
+import time
+import logging
+import threading
+
+__crt_minute_usage__ = {}
+"""当前分钟每个会话的对话次数"""
+
+
+__timer_thr__: threading.Thread = None
+
+
+def add_usage(session_name: str):
+    """增加会话的对话次数"""
+    global __crt_minute_usage__
+    if session_name in __crt_minute_usage__:
+        __crt_minute_usage__[session_name] += 1
+    else:
+        __crt_minute_usage__[session_name] = 1
+
+
+def start_timer():
+    """启动定时器"""
+    global __timer_thr__
+    __timer_thr__ = threading.Thread(target=run_timer, daemon=True)
+    __timer_thr__.start()
+
+
+def run_timer():
+    """启动定时器，每分钟清空一次对话次数"""
+    global __crt_minute_usage__
+    global __timer_thr__
+    
+    # 等待直到整分钟
+    time.sleep(60 - time.time() % 60)
+
+    while True:
+        if __timer_thr__ != threading.current_thread():
+            break
+        
+        logging.debug("清空当前分钟的对话次数")
+        __crt_minute_usage__ = {}
+        time.sleep(60)
+
+
+def get_usage(session_name: str) -> int:
+    """获取会话的对话次数"""
+    global __crt_minute_usage__
+    if session_name in __crt_minute_usage__:
+        return __crt_minute_usage__[session_name]
+    else:
+        return 0
+
+
+def get_rest_wait_time(session_name: str, spent: float) -> float:
+    """获取会话此回合的剩余等待时间"""
+    global __crt_minute_usage__
+
+    import config
+
+    if not hasattr(config, 'rate_limitation'):
+        return 0
+
+    min_seconds_per_round = 60.0 / config.rate_limitation
+
+    if session_name in __crt_minute_usage__:
+        return max(0, min_seconds_per_round - spent)
+    else:
+        return 0
+
+
+def is_reach_limit(session_name: str) -> bool:
+    """判断会话是否超过限制"""
+    global __crt_minute_usage__
+
+    import config
+
+    if not hasattr(config, 'rate_limitation'):
+        return False
+
+    if session_name in __crt_minute_usage__:
+        return __crt_minute_usage__[session_name] >= config.rate_limitation
+    else:
+        return False
+
+start_timer()