QChatGPT/pkg/qqbot/ratelimit.py

# 限速相关模块
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()