Merge pull request #32 from RockChinQ/multi-keys

feat: 支持多个api-key并支持设置文字量阈值对其进行自动切换 #8

config-template.py

@@ -16,10 +16,23 @@ mirai_http_api_config = {
 
 # [必需] OpenAI的配置
 # api_key: OpenAI的API Key
+# 若只有一个api-key，请直接修改以下内容中的"openai_api_key"为你的api-key
+# 如准备了多个api-key，可以以字典的形式填写，程序会自动选择可用的api-key
+# 例如{
+#       "api0": "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
+#       "api1": "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
+#     }
 openai_config = {
-    "api_key": "openai_api_key",
+    "api_key": {
+        "default": "openai_api_key"
+    },
 }
 
+# 单个api-key的使用量警告阈值
+# 当使用此api-key进行请求的文字量达到此阈值时，会在控制台输出警告并通知管理员
+# 若之后还有未使用超过此值的api-key，则会切换到新的api-key进行请求
+api_key_usage_threshold = 895000
+
 # 管理员QQ号，用于接收报错等通知，为0时不发送通知
 admin_qq = 0
 

main.py

@@ -8,8 +8,8 @@ import logging
 import colorlog
 
 import sys
-sys.path.append(".")
 
+sys.path.append(".")
 
 log_colors_config = {
     'DEBUG': 'green',  # cyan white
@@ -55,12 +55,12 @@ def main():
     import pkg.qqbot.manager
 
     # 主启动流程
-    openai_interact = pkg.openai.manager.OpenAIInteract(config.openai_config['api_key'], config.completion_api_params)
-
     database = pkg.database.manager.DatabaseManager()
 
     database.initialize_database()
 
+    openai_interact = pkg.openai.manager.OpenAIInteract(config.openai_config['api_key'], config.completion_api_params)
+
     # 加载所有未超时的session
     pkg.openai.session.load_sessions()
 
@@ -78,6 +78,7 @@ def main():
             time.sleep(86400)
         except KeyboardInterrupt:
             try:
+                pkg.openai.manager.get_inst().key_mgr.dump_usage()
                 for session in pkg.openai.session.sessions:
                     logging.info('持久化session: %s', session)
                     pkg.openai.session.sessions[session].persistence()
@@ -85,7 +86,7 @@ def main():
                 if not isinstance(e, KeyboardInterrupt):
                     raise e
             print("程序退出")
-            break
+            sys.exit(0)
 
 
 if __name__ == '__main__':

pkg/database/manager.py

@@ -1,3 +1,4 @@
+import hashlib
 import logging
 import time
 from sqlite3 import Cursor
@@ -48,6 +49,15 @@ class DatabaseManager:
             `prompt` text not null
         )
         """)
+
+        self.execute("""
+        create table if not exists `api_key_usage`(
+            `id` INTEGER PRIMARY KEY AUTOINCREMENT,
+            `key_md5` varchar(255) not null,
+            `timestamp` bigint not null,
+            `usage` bigint not null
+        )
+        """)
         print('Database initialized.')
 
     # session持久化
@@ -214,6 +224,45 @@ class DatabaseManager:
 
         return sessions
 
+    # 将apikey的使用量存进数据库
+    def dump_api_key_usage(self, api_keys: dict, usage: dict):
+        logging.debug('dumping api key usage...')
+        logging.debug(api_keys)
+        logging.debug(usage)
+        for api_key in api_keys:
+            # 计算key的md5值
+            key_md5 = hashlib.md5(api_keys[api_key].encode('utf-8')).hexdigest()
+            # 获取使用量
+            usage_count = 0
+            if key_md5 in usage:
+                usage_count = usage[key_md5]
+            # 将使用量存进数据库
+            # 先检查是否已存在
+            self.execute("""
+            select count(*) from `api_key_usage` where `key_md5` = '{}'""".format(key_md5))
+            result = self.cursor.fetchone()
+            if result[0] == 0:
+                # 不存在则插入
+                self.execute("""
+                insert into `api_key_usage` (`key_md5`, `usage`,`timestamp`) values ('{}', {}, {})
+                """.format(key_md5, usage_count, int(time.time())))
+            else:
+                # 存在则更新，timestamp设置为当前
+                self.execute("""
+                update `api_key_usage` set `usage` = {}, `timestamp` = {} where `key_md5` = '{}'
+                """.format(usage_count, int(time.time()), key_md5))
+
+    def load_api_key_usage(self):
+        self.execute("""
+        select `key_md5`, `usage` from `api_key_usage`
+        """)
+        results = self.cursor.fetchall()
+        usage = {}
+        for result in results:
+            key_md5 = result[0]
+            usage_count = result[1]
+            usage[key_md5] = usage_count
+        return usage
 
 def get_inst() -> DatabaseManager:
     global inst

pkg/openai/manager.py

@@ -2,19 +2,24 @@ import openai
 
 import config
 
+import pkg.openai.keymgr
+
 inst = None
 
 
 # 为其他模块提供与OpenAI交互的接口
 class OpenAIInteract:
-    api_key = ''
     api_params = {}
 
+    key_mgr = None
+
     def __init__(self, api_key: str, api_params: dict):
-        self.api_key = api_key
+        # self.api_key = api_key
         self.api_params = api_params
 
-        openai.api_key = self.api_key
+        self.key_mgr = pkg.openai.keymgr.KeysManager(api_key)
+
+        openai.api_key = self.key_mgr.get_using_key()
 
         global inst
         inst = self
@@ -27,6 +32,10 @@ class OpenAIInteract:
             timeout=config.process_message_timeout,
             **self.api_params
         )
+        switched = self.key_mgr.report_usage(prompt + response['choices'][0]['text'])
+        if switched:
+            openai.api_key = self.key_mgr.get_using_key()
+
         return response