From c9131b333b5a0223cfdb97ded7c2b68866fefc6b Mon Sep 17 00:00:00 2001
From: lanvent <lanvent@qq.com>
Date: Fri, 21 Apr 2023 13:41:21 +0800
Subject: [PATCH] feat: add clear_quota_v2 method to clear API quota when it's
 used up

---
 channel/wechatmp/wechatmp_channel.py |  2 --
 channel/wechatmp/wechatmp_client.py  | 26 ++++++++++++++++----------
 requirements-optional.txt            |  1 +
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/channel/wechatmp/wechatmp_channel.py b/channel/wechatmp/wechatmp_channel.py
index 9da2aa8..afa5ac5 100644
--- a/channel/wechatmp/wechatmp_channel.py
+++ b/channel/wechatmp/wechatmp_channel.py
@@ -92,7 +92,6 @@ class WechatMPChannel(ChatChannel):
             elif reply.type == ReplyType.IMAGE_URL:  # 从网络下载图片
                 img_url = reply.content
                 pic_res = requests.get(img_url, stream=True)
-                print(pic_res.headers)
                 image_storage = io.BytesIO()
                 for block in pic_res.iter_content(1024):
                     image_storage.write(block)
@@ -159,7 +158,6 @@ class WechatMPChannel(ChatChannel):
             elif reply.type == ReplyType.IMAGE_URL:  # 从网络下载图片
                 img_url = reply.content
                 pic_res = requests.get(img_url, stream=True)
-                print(pic_res.headers)
                 image_storage = io.BytesIO()
                 for block in pic_res.iter_content(1024):
                     image_storage.write(block)
diff --git a/channel/wechatmp/wechatmp_client.py b/channel/wechatmp/wechatmp_client.py
index 7938ead..ee0ec84 100644
--- a/channel/wechatmp/wechatmp_client.py
+++ b/channel/wechatmp/wechatmp_client.py
@@ -1,11 +1,9 @@
 import time
-import json
-import requests
 import threading
 from channel.wechatmp.common import *
 from wechatpy.client import WeChatClient
 from common.log import logger
-from config import conf
+from wechatpy.exceptions import APILimitedException
 
 
 class WechatMPClient(WeChatClient):
@@ -16,13 +14,13 @@ class WechatMPClient(WeChatClient):
         )
         self.fetch_access_token_lock = threading.Lock()
 
-    def fetch_access_token(self):
-        """
-        获取 access token
-        详情请参考 http://mp.weixin.qq.com/wiki/index.php?title=通用接口文档
+    def clear_quota(self):
+        return self.post("clear_quota", data={"appid": self.appid})
 
-        :return: 返回的 JSON 数据包
-        """
+    def clear_quota_v2(self):
+        return self.post("clear_quota/v2", params={"appid": self.appid, "appsecret": self.secret})
+
+    def fetch_access_token(self): # 重载父类方法，加锁避免多线程重复获取access_token
         with self.fetch_access_token_lock:
             access_token = self.session.get(self.access_token_key)
             if access_token:
@@ -32,4 +30,12 @@ class WechatMPClient(WeChatClient):
                 if self.expires_at - timestamp > 60:
                     return access_token
             return super().fetch_access_token()
-        
\ No newline at end of file
+
+    def _request(self, method, url_or_endpoint, **kwargs): # 重载父类方法，遇到API限流时，清除quota后重试
+        try:
+            return super()._request(method, url_or_endpoint, **kwargs)
+        except APILimitedException as e:
+            logger.error("[wechatmp] API quata has been used up. {}".format(e))
+            response = self.clear_quota_v2()
+            logger.debug("[wechatmp] API quata has been cleard, {}".format(response))
+            return super()._request(method, url_or_endpoint, **kwargs)
\ No newline at end of file
diff --git a/requirements-optional.txt b/requirements-optional.txt
index cfb52c9..ba453cb 100644
--- a/requirements-optional.txt
+++ b/requirements-optional.txt
@@ -18,6 +18,7 @@ pysilk_mod>=1.6.0 # needed by send voice
 
 # wechatmp
 web.py
+wechatpy
 
 # chatgpt-tool-hub plugin