Browse Source

Update rate_limiter: support multimodal message

StrayWarrior 2 weeks ago
parent
commit
6a4807ed7d
1 changed files with 10 additions and 2 deletions
  1. 10 2
      pqai_agent/rate_limiter.py

+ 10 - 2
pqai_agent/rate_limiter.py

@@ -3,7 +3,11 @@
 # vim:fenc=utf-8
 # vim:fenc=utf-8
 
 
 import time
 import time
+from typing import Optional, Union, Dict
+
 from pqai_agent.logging_service import logger
 from pqai_agent.logging_service import logger
+from pqai_agent.mq_message import MessageType
+
 
 
 class MessageSenderRateLimiter:
 class MessageSenderRateLimiter:
     MAX_CHAR_PER_SECOND = 5
     MAX_CHAR_PER_SECOND = 5
@@ -11,11 +15,15 @@ class MessageSenderRateLimiter:
     def __init__(self):
     def __init__(self):
         self.last_send_time = {}
         self.last_send_time = {}
 
 
-    def wait_for_sending(self, sender_id: str, next_message: str):
+    def wait_for_sending(self, sender_id: str, next_message: Union[str, Dict]):
         current_time = time.time()
         current_time = time.time()
         last_send_time = self.last_send_time.get(sender_id, 0)
         last_send_time = self.last_send_time.get(sender_id, 0)
         elapsed_time = current_time - last_send_time
         elapsed_time = current_time - last_send_time
-        required_time = len(next_message) / self.MAX_CHAR_PER_SECOND
+        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:
         if elapsed_time < required_time:
             logger.debug(f"Rate limit exceeded. Waiting for {required_time - elapsed_time:.2f} seconds.")
             logger.debug(f"Rate limit exceeded. Waiting for {required_time - elapsed_time:.2f} seconds.")
             time.sleep(required_time - elapsed_time)
             time.sleep(required_time - elapsed_time)