Merge branch 'zhayujie:master' into master

README.md

@@ -134,6 +134,7 @@ pip3 install --upgrade openai
 + 对于图像生成，在满足个人或群组触发条件外，还需要额外的关键词前缀来触发，对应配置 `image_create_prefix `
 + 关于OpenAI对话及图片接口的参数配置（内容自由度、回复字数限制、图片大小等），可以参考 [对话接口](https://beta.openai.com/docs/api-reference/completions) 和 [图像接口](https://beta.openai.com/docs/api-reference/completions)  文档直接在 [代码](https://github.com/zhayujie/chatgpt-on-wechat/blob/master/bot/openai/open_ai_bot.py) `bot/openai/open_ai_bot.py` 中进行调整。
 + `conversation_max_tokens`：表示能够记忆的上下文最大字数（一问一答为一组对话，如果累积的对话字数超出限制，就会优先移除最早的一组对话）
++ `rate_limit_chatgpt`，`rate_limit_dalle`：每分钟最高问答速率、画图速率，超速后排队按序处理。
 + `clear_memory_commands`: 对话内指令，主动清空前文记忆，字符串数组可自定义指令别名。
 + `hot_reload`: 程序退出后，暂存微信扫码状态，默认关闭。
 + `character_desc` 配置中保存着你对机器人说的一段话，他会记住这段话并作为他的设定，你可以为他定制任何人格      (关于会话上下文的更多内容参考该 [issue](https://github.com/zhayujie/chatgpt-on-wechat/issues/43))

bot/chatgpt/chat_gpt_bot.py

@@ -3,6 +3,7 @@
 from bot.bot import Bot
 from config import conf, load_config
 from common.log import logger
+from common.token_bucket import TokenBucket
 from common.expired_dict import ExpiredDict
 import openai
 import time
@@ -21,6 +22,10 @@ class ChatGPTBot(Bot):
         proxy = conf().get('proxy')
         if proxy:
             openai.proxy = proxy
+        if conf().get('rate_limit_chatgpt'):
+            self.tb4chatgpt = TokenBucket(conf().get('rate_limit_chatgpt', 20))
+        if conf().get('rate_limit_dalle'):
+            self.tb4dalle = TokenBucket(conf().get('rate_limit_dalle', 50))
 
     def reply(self, query, context=None):
         # acquire reply content
@@ -63,6 +68,8 @@ class ChatGPTBot(Bot):
         :return: {}
         '''
         try:
+            if conf().get('rate_limit_chatgpt') and not self.tb4chatgpt.get_token():
+                return {"completion_tokens": 0, "content": "提问太快啦，请休息一下再问我吧"}
             response = openai.ChatCompletion.create(
                 model= conf().get("model") or "gpt-3.5-turbo",  # 对话模型的名称
                 messages=session,
@@ -102,6 +109,8 @@ class ChatGPTBot(Bot):
 
     def create_img(self, query, retry_count=0):
         try:
+            if conf().get('rate_limit_dalle') and not self.tb4dalle.get_token():
+                return "请求太快了，请休息一下再问我吧"
             logger.info("[OPEN_AI] image_query={}".format(query))
             response = openai.Image.create(
                 prompt=query,    #图片描述
@@ -118,7 +127,7 @@ class ChatGPTBot(Bot):
                 logger.warn("[OPEN_AI] ImgCreate RateLimit exceed, 第{}次重试".format(retry_count+1))
                 return self.create_img(query, retry_count+1)
             else:
-                return "提问太快啦，请休息一下再问我吧"
+                return "请求太快啦，请休息一下再问我吧"
         except Exception as e:
             logger.exception(e)
             return None

channel/wechat/wechat_channel.py

@@ -14,6 +14,7 @@ from common.tmp_dir import TmpDir
 from config import conf
 import requests
 import io
+import time
 
 thread_pool = ThreadPoolExecutor(max_workers=8)
 
@@ -74,7 +75,11 @@ class WechatChannel(Channel):
         from_user_id = msg['FromUserName']
         to_user_id = msg['ToUserName']              # 接收人id
         other_user_id = msg['User']['UserName']     # 对手方id
+        create_time = msg['CreateTime']             # 消息时间
         match_prefix = self.check_prefix(content, conf().get('single_chat_prefix'))
+        if conf().get('hot_reload') == True and int(create_time) < int(time.time()) - 60:    #跳过1分钟前的历史消息
+            logger.debug("[WX]history message skipped")
+            return
         if "」\n- - - - - - - - - - - - - - -" in content:
             logger.debug("[WX]reference query skipped")
             return
@@ -108,6 +113,10 @@ class WechatChannel(Channel):
         logger.debug("[WX]receive group msg: " + json.dumps(msg, ensure_ascii=False))
         group_name = msg['User'].get('NickName', None)
         group_id = msg['User'].get('UserName', None)
+        create_time = msg['CreateTime']             # 消息时间
+        if conf().get('hot_reload') == True and int(create_time) < int(time.time()) - 60:    #跳过1分钟前的历史消息
+            logger.debug("[WX]history group message skipped")
+            return
         if not group_name:
             return ""
         origin_content = msg['Content']

channel/wechat/wechaty_channel.py

@@ -147,6 +147,13 @@ class WechatyChannel(Channel):
             match_prefix = (is_at and not config.get("group_at_off", False)) \
                            or self.check_prefix(content, config.get('group_chat_prefix')) \
                            or self.check_contain(content, config.get('group_chat_keyword'))
+            # Wechaty判断is_at为True，返回的内容是过滤掉@之后的内容；而is_at为False，则会返回完整的内容
+            # 故判断如果匹配到自定义前缀，则返回过滤掉前缀+空格后的内容，用于实现类似自定义+前缀触发生成AI图片的功能
+            prefixes = config.get('group_chat_prefix')
+            for prefix in prefixes:
+                if content.startswith(prefix):
+                    content = content.replace(prefix, '', 1).strip()
+                    break
             if ('ALL_GROUP' in config.get('group_name_white_list') or room_name in config.get(
                     'group_name_white_list') or self.check_contain(room_name, config.get(
                 'group_name_keyword_white_list'))) and match_prefix:

common/token_bucket.py

@@ -0,0 +1,45 @@
+import threading
+import time
+
+
+class TokenBucket:
+    def __init__(self, tpm, timeout=None):
+        self.capacity = int(tpm)  # 令牌桶容量
+        self.tokens = 0  # 初始令牌数为0
+        self.rate = int(tpm) / 60  # 令牌每秒生成速率
+        self.timeout = timeout  # 等待令牌超时时间
+        self.cond = threading.Condition()  # 条件变量
+        self.is_running = True
+        # 开启令牌生成线程
+        threading.Thread(target=self._generate_tokens).start()
+
+    def _generate_tokens(self):
+        """生成令牌"""
+        while self.is_running:
+            with self.cond:
+                if self.tokens < self.capacity:
+                    self.tokens += 1
+                self.cond.notify()  # 通知获取令牌的线程
+            time.sleep(1 / self.rate)
+
+    def get_token(self):
+        """获取令牌"""
+        with self.cond:
+            while self.tokens <= 0:
+                flag = self.cond.wait(self.timeout)
+                if not flag:  # 超时
+                    return False
+            self.tokens -= 1
+        return True
+
+    def close(self):
+        self.is_running = False
+
+
+if __name__ == "__main__":
+    token_bucket = TokenBucket(20, None)  # 创建一个每分钟生产20个tokens的令牌桶
+    # token_bucket = TokenBucket(20, 0.1)
+    for i in range(3):
+        if token_bucket.get_token():
+            print(f"第{i+1}次请求成功")
+    token_bucket.close()

config.py

@@ -9,7 +9,7 @@ config = {}
 
 def load_config():
     global config
-    config_path = "config.json"
+    config_path = "./config.json"
     if not os.path.exists(config_path):
         raise Exception('配置文件不存在，请根据config-template.json模板创建config.json文件')