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

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