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