diff --git a/config-template.py b/config-template.py index c62dc0e5..5fc5e283 100644 --- a/config-template.py +++ b/config-template.py @@ -52,10 +52,10 @@ response_rules = { "regexp": [] # "为什么.*", "怎么?样.*", "怎么.*", "如何.*", "[Hh]ow to.*", "[Ww]hy not.*", "[Ww]hat is.*", ".*怎么办", ".*咋办" } -# 单个api-key的使用量警告阈值 -# 当使用此api-key进行请求的文字量达到此阈值时，会在控制台输出警告并通知管理员 +# 单个api-key的费用警告阈值 +# 当使用此api-key进行请求所消耗的费用估算达到此阈值时，会在控制台输出警告并通知管理员 # 若之后还有未使用超过此值的api-key，则会切换到新的api-key进行请求 -api_key_usage_threshold = 900000 +api_key_fee_threshold = 18.0 # 敏感词过滤开关，以同样数量的*代替敏感词回复 # 请在sensitive.json中添加敏感词 @@ -80,7 +80,7 @@ completion_api_params = { # OpenAI的Image API的参数 # 具体请查看OpenAI的文档: https://beta.openai.com/docs/api-reference/images/create image_api_params = { - "size": "256x256", + "size": "256x256", # 图片尺寸，支持256x256, 512x512, 1024x1024 } # 消息处理的超时时间，单位为秒 diff --git a/main.py b/main.py index 7352942c..8c19f5df 100644 --- a/main.py +++ b/main.py @@ -78,7 +78,7 @@ def main(): time.sleep(86400) except KeyboardInterrupt: try: - pkg.openai.manager.get_inst().key_mgr.dump_usage() + pkg.openai.manager.get_inst().key_mgr.dump_fee() for session in pkg.openai.session.sessions: logging.info('持久化session: %s', session) pkg.openai.session.sessions[session].persistence() diff --git a/pkg/database/manager.py b/pkg/database/manager.py index ecd38002..3ede5f57 100644 --- a/pkg/database/manager.py +++ b/pkg/database/manager.py @@ -58,6 +58,15 @@ class DatabaseManager: `usage` bigint not null ) """) + + self.execute(""" + create table if not exists `account_fee`( + `id` INTEGER PRIMARY KEY AUTOINCREMENT, + `key_md5` varchar(255) not null, + `timestamp` bigint not null, + `fee` DECIMAL(9,3) not null + ) + """) print('Database initialized.') # session持久化 @@ -264,6 +273,45 @@ class DatabaseManager: usage[key_md5] = usage_count return usage + def dump_api_key_fee(self, api_keys: dict, fee: dict): + logging.debug("dumping api key fee...") + logging.debug(api_keys) + logging.debug(fee) + for api_key in api_keys: + # 计算key的md5值 + key_md5 = hashlib.md5(api_keys[api_key].encode('utf-8')).hexdigest() + # 获取使用量 + fee_count = 0 + if key_md5 in fee: + fee_count = fee[key_md5] + # 将使用量存进数据库 + # 先检查是否已存在 + self.execute(""" + select count(*) from `account_fee` where `key_md5` = '{}'""".format(key_md5)) + result = self.cursor.fetchone() + if result[0] == 0: + # 不存在则插入 + self.execute(""" + insert into `account_fee` (`key_md5`, `fee`,`timestamp`) values ('{}', {}, {}) + """.format(key_md5, fee_count, int(time.time()))) + else: + # 存在则更新，timestamp设置为当前 + self.execute(""" + update `account_fee` set `fee` = {}, `timestamp` = {} where `key_md5` = '{}' + """.format(fee_count, int(time.time()), key_md5)) + + def load_api_key_fee(self): + self.execute(""" + select `key_md5`, `fee` from `account_fee` + """) + results = self.cursor.fetchall() + fee = {} + for result in results: + key_md5 = result[0] + fee_count = result[1] + fee[key_md5] = fee_count + return fee + def get_inst() -> DatabaseManager: global inst return inst diff --git a/pkg/openai/keymgr.py b/pkg/openai/keymgr.py index 75d7b196..728cea3e 100644 --- a/pkg/openai/keymgr.py +++ b/pkg/openai/keymgr.py @@ -14,7 +14,11 @@ class KeysManager: # 其中键为api-key的md5值，值为使用量 usage = {} - api_key_usage_threshold = 900000 + fee = {} + + api_key_usage_threshold = 900000 # 已弃用 + + api_key_fee_threshold = 18.0 using_key = "" @@ -24,9 +28,11 @@ class KeysManager: return self.using_key def __init__(self, api_key): - if hasattr(config, 'api_key_usage_threshold'): - self.api_key_usage_threshold = config.api_key_usage_threshold - self.load_usage() + # if hasattr(config, 'api_key_usage_threshold'): + # self.api_key_usage_threshold = config.api_key_usage_threshold + if hasattr(config, 'api_key_fee_threshold'): + self.api_key_fee_threshold = config.api_key_fee_threshold + self.load_fee() if type(api_key) is dict: self.api_key = api_key @@ -45,9 +51,9 @@ class KeysManager: # 根据使用量自动切换到可用的api-key # 返回是否切换成功, 切换后的api-key的别名 def auto_switch(self) -> (bool, str): - self.dump_usage() + self.dump_fee() for key_name in self.api_key: - if self.get_usage(self.api_key[key_name]) < self.api_key_usage_threshold: + if self.get_fee(self.api_key[key_name]) < self.api_key_fee_threshold: self.using_key = self.api_key[key_name] logging.info("使用api-key:" + key_name) return True, key_name @@ -57,30 +63,76 @@ class KeysManager: return False, "" - def get_usage(self, api_key): - md5 = hashlib.md5(api_key.encode('utf-8')).hexdigest() - if md5 not in self.usage: - self.usage[md5] = 0 - return self.usage[md5] - def add(self, key_name, key): self.api_key[key_name] = key + # def get_usage(self, api_key): + # md5 = hashlib.md5(api_key.encode('utf-8')).hexdigest() + # if md5 not in self.usage: + # self.usage[md5] = 0 + # return self.usage[md5] + # 报告使用 # 返回是否需要将openai的api-key切换 - def report_usage(self, new_content: str) -> bool: + # def report_usage(self, new_content: str) -> bool: + # md5 = hashlib.md5(self.using_key.encode('utf-8')).hexdigest() + # if md5 not in self.usage: + # self.usage[md5] = 0 + # + # # 经测算得出的理论与实际的偏差比例 + # salt_rate = 0.91 + # + # self.usage[md5] += ( (len(new_content.encode('utf-8')) - len(new_content)) / 2 + len(new_content) )*salt_rate + # + # self.usage[md5] = int(self.usage[md5]) + # + # if self.usage[md5] >= self.api_key_usage_threshold: + # switch_result, key_name = self.auto_switch() + # + # # 检查是否切换到新的 + # if switch_result: + # if key_name not in self.alerted: + # # 通知管理员 + # pkg.qqbot.manager.get_inst().notify_admin("api-key已切换到:" + key_name) + # self.alerted.append(key_name) + # return True + # else: + # if key_name not in self.alerted: + # # 通知管理员 + # pkg.qqbot.manager.get_inst().notify_admin("api-key已用完，无未使用的api-key可供切换") + # self.alerted.append(key_name) + # return False + + # 设置当前使用的api-key使用量超限 + # 这是在尝试调用api时发生超限异常时调用的 + def set_current_exceeded(self): md5 = hashlib.md5(self.using_key.encode('utf-8')).hexdigest() - if md5 not in self.usage: - self.usage[md5] = 0 + # self.usage[md5] = self.api_key_usage_threshold + self.fee[md5] = self.api_key_fee_threshold + self.dump_fee() - # 经测算得出的理论与实际的偏差比例 - salt_rate = 0.91 + # def dump_usage(self): + # pkg.database.manager.get_inst().dump_api_key_usage(api_keys=self.api_key, usage=self.usage) - self.usage[md5] += ( (len(new_content.encode('utf-8')) - len(new_content)) / 2 + len(new_content) )*salt_rate + # def load_usage(self): + # self.usage = pkg.database.manager.get_inst().load_api_key_usage() + # logging.debug("load usage:" + str(self.usage)) + # print("load usage:" + str(self.usage)) - self.usage[md5] = int(self.usage[md5]) + def get_fee(self, api_key): + md5 = hashlib.md5(api_key.encode('utf-8')).hexdigest() + if md5 not in self.fee: + self.fee[md5] = 0 + return self.fee[md5] - if self.usage[md5] >= self.api_key_usage_threshold: + def report_fee(self, fee: float) -> bool: + md5 = hashlib.md5(self.using_key.encode('utf-8')).hexdigest() + if md5 not in self.fee: + self.fee[md5] = 0 + + self.fee[md5] += fee + + if self.fee[md5] >= self.api_key_fee_threshold: switch_result, key_name = self.auto_switch() # 检查是否切换到新的 @@ -97,17 +149,9 @@ class KeysManager: self.alerted.append(key_name) return False - # 设置当前使用的api-key使用量超限 - # 这是在尝试调用api时发生超限异常时调用的 - def set_current_exceeded(self): - md5 = hashlib.md5(self.using_key.encode('utf-8')).hexdigest() - self.usage[md5] = self.api_key_usage_threshold - self.dump_usage() + def dump_fee(self): + pkg.database.manager.get_inst().dump_api_key_fee(api_keys=self.api_key, fee=self.fee) - def dump_usage(self): - pkg.database.manager.get_inst().dump_api_key_usage(api_keys=self.api_key, usage=self.usage) - - def load_usage(self): - self.usage = pkg.database.manager.get_inst().load_api_key_usage() - logging.debug("load usage:" + str(self.usage)) - print("load usage:" + str(self.usage)) + def load_fee(self): + self.fee = pkg.database.manager.get_inst().load_api_key_fee() + logging.info("load fee:" + str(self.fee)) \ No newline at end of file diff --git a/pkg/openai/manager.py b/pkg/openai/manager.py index e9279cfd..d65f0e54 100644 --- a/pkg/openai/manager.py +++ b/pkg/openai/manager.py @@ -5,6 +5,7 @@ import openai import config import pkg.openai.keymgr +import pkg.openai.pricing as pricing inst = None @@ -37,19 +38,30 @@ class OpenAIInteract: timeout=config.process_message_timeout, **config.completion_api_params ) - switched = self.key_mgr.report_usage(prompt + response['choices'][0]['text']) + + switched = self.key_mgr.report_fee(pricing.language_base_price(config.completion_api_params['model'], + prompt + response['choices'][0]['text'])) + if switched: openai.api_key = self.key_mgr.get_using_key() return response def request_image(self, prompt): + + params = config.image_api_params if hasattr(config, "image_api_params") else self.default_image_api_params + response = openai.Image.create( prompt=prompt, n=1, - **config.image_api_params if hasattr(config, "image_api_params") else self.default_image_api_params + **params ) + switched = self.key_mgr.report_fee(pricing.image_price(params['size'])) + + if switched: + openai.api_key = self.key_mgr.get_using_key() + return response diff --git a/pkg/openai/pricing.py b/pkg/openai/pricing.py new file mode 100644 index 00000000..1040a04c --- /dev/null +++ b/pkg/openai/pricing.py @@ -0,0 +1,21 @@ +pricing = { + "base": { # 文字模型单位是1000字符 + "text-davinci-003": 0.02, + }, + "image": { + "256x256": 0.016, + "512x512": 0.018, + "1024x1024": 0.02, + } +} + + +def language_base_price(model, text): + salt_rate = 0.93 + length = ((len(text.encode('utf-8')) - len(text)) / 2 + len(text)) * salt_rate + + return pricing["base"][model] * length / 1000 + + +def image_price(size): + return pricing["image"][size] diff --git a/pkg/qqbot/process.py b/pkg/qqbot/process.py index 5e9a5821..eadde547 100644 --- a/pkg/qqbot/process.py +++ b/pkg/qqbot/process.py @@ -113,17 +113,17 @@ def process_message(launcher_type: str, launcher_id: int, text_message: str) -> elif cmd == 'usage': api_keys = pkg.openai.manager.get_inst().key_mgr.api_key reply_str = "[bot]api-key使用情况:(阈值:{})\n\n".format( - pkg.openai.manager.get_inst().key_mgr.api_key_usage_threshold) + pkg.openai.manager.get_inst().key_mgr.api_key_fee_threshold) using_key_name = "" for api_key in api_keys: - reply_str += "{}:\n - {}字 {}%\n".format(api_key, - pkg.openai.manager.get_inst().key_mgr.get_usage( + reply_str += "{}:\n - {}元 {}%\n".format(api_key, + pkg.openai.manager.get_inst().key_mgr.get_fee( api_keys[api_key]), round( - pkg.openai.manager.get_inst().key_mgr.get_usage( + pkg.openai.manager.get_inst().key_mgr.get_fee( api_keys[ - api_key]) / pkg.openai.manager.get_inst().key_mgr.api_key_usage_threshold * 100, + api_key]) / pkg.openai.manager.get_inst().key_mgr.api_key_fee_threshold * 100, 3)) if api_keys[api_key] == pkg.openai.manager.get_inst().key_mgr.using_key: using_key_name = api_key @@ -158,7 +158,7 @@ def process_message(launcher_type: str, launcher_id: int, text_message: str) -> reply = ["[bot]err:调用API失败，请重试或联系作者，或等待修复"] except openai.error.RateLimitError as e: # 尝试切换api-key - current_tokens_amt = pkg.openai.manager.get_inst().key_mgr.get_usage( + current_tokens_amt = pkg.openai.manager.get_inst().key_mgr.get_fee( pkg.openai.manager.get_inst().key_mgr.get_using_key()) pkg.openai.manager.get_inst().key_mgr.set_current_exceeded() switched, name = pkg.openai.manager.get_inst().key_mgr.auto_switch()