diff --git a/config-template.py b/config-template.py index 3820459a..62f5783d 100644 --- a/config-template.py +++ b/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 diff --git a/main.py b/main.py index 38e1d5e9..988626fc 100644 --- a/main.py +++ b/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__': diff --git a/pkg/database/manager.py b/pkg/database/manager.py index 30e508f5..8b04218c 100644 --- a/pkg/database/manager.py +++ b/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 diff --git a/pkg/openai/manager.py b/pkg/openai/manager.py index fed9fd2c..3e9daec0 100644 --- a/pkg/openai/manager.py +++ b/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