AgentCoreService/pqai_agent/rate_limiter.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8

import time
from typing import Optional, Union, Dict

from pqai_agent.logging_service import logger
from pqai_agent.mq_message import MessageType


class MessageSenderRateLimiter:
    MAX_CHAR_PER_SECOND = 5

    def __init__(self):
        self.last_send_time = {}

    def wait_for_sending(self, sender_id: str, next_message: Union[str, Dict]):
        current_time = time.time()
        last_send_time = self.last_send_time.get(sender_id, 0)
        elapsed_time = current_time - last_send_time
        if isinstance(next_message, str) or next_message["type"] in (MessageType.TEXT, MessageType.VOICE):
            required_time = len(next_message) / self.MAX_CHAR_PER_SECOND
        else:
            # FIXME: 非文字消息的判断方式
            required_time = 2
        if elapsed_time < required_time:
            logger.debug(f"Rate limit exceeded. Waiting for {required_time - elapsed_time:.2f} seconds.")
            time.sleep(required_time - elapsed_time)
        current_time = time.time()
        self.last_send_time[sender_id] = current_time
        return