xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-06-02 03:55:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: 完全删除费用估算功能 (#81) (#96)

Browse Source
This commit is contained in:
Rock Chin
2023-01-10 14:54:50 +08:00
parent f884313d72
commit dfa8621a1a
8 changed files with 5 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
main.py
View File

@@ -204,7 +204,6 @@ def stop():
qqbot_inst = pkg.utils.context.get_qqbot_manager()
assert isinstance(qqbot_inst, pkg.qqbot.manager.QQBotManager)
pkg.utils.context.get_openai_manager().key_mgr.dump_fee()
for session in pkg.openai.session.sessions:
logging.info('持久化session: %s', session)
pkg.openai.session.sessions[session].persistence()

48
pkg/database/manager.py
View File

@@ -50,15 +50,6 @@ class DatabaseManager:
)
""")
# 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
# )
# """)
self.execute("""
create table if not exists `account_fee`(
`id` INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -281,45 +272,6 @@ 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 dump_usage_json(self, usage: dict):
json_str = json.dumps(usage)
self.execute("""

59
pkg/openai/keymgr.py
View File

@@ -12,10 +12,6 @@ class KeysManager:
# api-key的使用量
# 其中键为api-key的md5值，值为使用量
fee = {}
api_key_fee_threshold = 18.0
using_key = ""
alerted = []
@@ -34,11 +30,6 @@ class KeysManager:
# if hasattr(config, 'api_key_usage_threshold'):
# self.api_key_usage_threshold = config.api_key_usage_threshold
config = pkg.utils.context.get_config()
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
elif type(api_key) is str:
@@ -59,16 +50,11 @@ class KeysManager:
# 根据tested自动切换到可用的api-key
# 返回是否切换成功, 切换后的api-key的别名
def auto_switch(self) -> (bool, str):
self.dump_fee()
for key_name in self.api_key:
if self.api_key[key_name] not in self.exceeded:
self.using_key = self.api_key[key_name]
logging.info("使用api-key:" + key_name)
return True, key_name
# 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
self.using_key = list(self.api_key.values())[0]
logging.info("使用api-key:" + list(self.api_key.keys())[0])
@@ -85,53 +71,10 @@ class KeysManager:
# self.usage[md5] = self.api_key_usage_threshold
# self.fee[md5] = self.api_key_fee_threshold
self.exceeded.append(self.using_key)
self.dump_fee()
def get_key_name(self, api_key):
"""根据api-key获取其别名"""
for key_name in self.api_key:
if self.api_key[key_name] == api_key:
return key_name
return ""
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]
def report_fee(self, fee: float) -> bool:
logging.debug("report fee:" + str(fee))
md5 = hashlib.md5(self.using_key.encode('utf-8')).hexdigest()
if md5 not in self.fee:
self.fee[md5] = 0
self.fee[md5] += fee
config = pkg.utils.context.get_config()
if self.fee[md5] >= self.api_key_fee_threshold and \
hasattr(config, 'auto_switch_api_key') and config.auto_switch_api_key:
switch_result, key_name = self.auto_switch()
# 检查是否切换到新的
if switch_result:
if key_name not in self.alerted:
# 通知管理员
pkg.utils.context.get_qqbot_manager().notify_admin("api-key已切换到:" + key_name)
self.alerted.append(key_name)
return True
else:
if key_name not in self.alerted:
# 通知管理员
pkg.utils.context.get_qqbot_manager().notify_admin("api-key已用完，无未使用的api-key可供切换")
self.alerted.append(key_name)
return False
return False
def dump_fee(self):
pkg.utils.context.get_database_manager().dump_api_key_fee(api_keys=self.api_key, fee=self.fee)
def load_fee(self):
self.fee = pkg.utils.context.get_database_manager().load_api_key_fee()
logging.info("load fee:" + str(self.fee))
return ""

11
pkg/openai/manager.py
View File

@@ -3,7 +3,6 @@ import logging
import openai
import pkg.openai.keymgr
import pkg.openai.pricing as pricing
import pkg.utils.context
import pkg.audit.gatherer
@@ -43,11 +42,6 @@ class OpenAIInteract:
self.audit_mgr.report_text_model_usage(config.completion_api_params['model'],
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):
@@ -63,10 +57,5 @@ class OpenAIInteract:
self.audit_mgr.report_image_model_usage(params['size'])
switched = self.key_mgr.report_fee(pricing.image_price(params['size']))
if switched:
openai.api_key = self.key_mgr.get_using_key()
return response

3
pkg/openai/pricing.py → pkg/openai/pricing.bak.py
View File

@@ -1,3 +1,6 @@
# 计费模块
# 已弃用 https://github.com/RockChinQ/QChatGPT/issues/81
import logging
pricing = {

16
pkg/openai/session.py
View File

@@ -51,22 +51,6 @@ def dump_session(session_name: str):
del sessions[session_name]
# def blocked_func(lock: threading.Lock):
#
# def decorator(func):
# def wrapper(*args, **kwargs):
# print('lock acquire,{}'.format(lock))
# lock.acquire()
# try:
# return func(*args, **kwargs)
# finally:
# lock.release()
#
# return wrapper
#
# return decorator
# 通用的OpenAI API交互session
# session内部保留了对话的上下文，
# 收到用户消息后，将上下文提交给OpenAI API生成回复

1
pkg/qqbot/filter.py
View File

@@ -1,3 +1,4 @@
# 敏感词过滤模块
import re

21
pkg/qqbot/process.py
View File

@@ -212,27 +212,6 @@ def process_message(launcher_type: str, launcher_id: int, text_message: str, mes
to_send = session.undo()
text_message = to_send
is_message = True
elif cmd == 'fee':
api_keys = pkg.utils.context.get_openai_manager().key_mgr.api_key
reply_str = "[bot]api-key费用情况(估算):(阈值:{})\n\n".format(
pkg.utils.context.get_openai_manager().key_mgr.api_key_fee_threshold)
using_key_name = ""
for api_key in api_keys:
reply_str += "{}:\n - {}美元 {}%\n".format(api_key,
round(
pkg.utils.context.get_openai_manager().key_mgr.get_fee(
api_keys[api_key]), 6),
round(
pkg.utils.context.get_openai_manager().key_mgr.get_fee(
api_keys[
api_key]) / pkg.utils.context.get_openai_manager().key_mgr.api_key_fee_threshold * 100,
3))
if api_keys[api_key] == pkg.utils.context.get_openai_manager().key_mgr.using_key:
using_key_name = api_key
reply_str += "\n当前使用:{}".format(using_key_name)
reply = [reply_str]
elif cmd == 'usage':
reply_str = "[bot]各api-key使用情况:\n\n"