Files
AstrBot/cores/qqbot/core.py
T
2022-12-15 15:14:40 +08:00

342 lines
12 KiB
Python

import botpy
from botpy.message import Message
import yaml
import asyncio
import cores.openai.core
import re
from util.errors.errors import PromptExceededError
from botpy.message import DirectMessage
import json
chatgpt = ""
# db = ""
session_dict = {}
max_tokens = 2000
version = "1.3"
gpt_config = {
'engine': '',
'temperature': '',
'top_p': '',
'frequency_penalty': '',
'presence_penalty': '',
'max_tokens': '',
}
count = {
}
stat_file = ''
uniqueSession = False
class botClient(botpy.Client):
async def on_at_message_create(self, message: Message):
global stat_file
try:
if str(message.guild_id) not in count:
count[str(message.guild_id)] = {
'count': 1,
'direct_count': 0,
}
else:
count[str(message.guild_id)]['count'] += 1
stat_file = open("./configs/stat", 'w', encoding='utf-8')
stat_file.write(json.dumps(count))
stat_file.flush()
stat_file.close()
except BaseException:
pass
await oper_msg(message=message, at=True)
async def on_direct_message_create(self, message: DirectMessage):
global stat_file
try:
if str(message.guild_id) not in count:
count[str(message.guild_id)] = {
'count': 1,
'direct_count': 1,
}
else:
count[str(message.guild_id)]['count'] += 1
count[str(message.guild_id)]['direct_count'] += 1
stat_file = open("./configs/stat", 'w', encoding='utf-8')
stat_file.write(json.dumps(count))
stat_file.flush()
stat_file.close()
except BaseException:
pass
await oper_msg(message=message, at=False)
def initBot(chatgpt_inst):
global chatgpt
chatgpt = chatgpt_inst
# global db
# db = db_inst
global max_tokens
max_tokens = int(chatgpt_inst.getConfigs()['total_tokens_limit'])
global gpt_config
gpt_config = chatgpt_inst.getConfigs()
gpt_config['key'] = "***"
# 读统计信息
global stat_file
stat_file = open("./configs/stat", 'r', encoding='utf-8')
global count
res = stat_file.read()
if res == '':
count = {}
else:
try:
count = json.loads(res)
except BaseException:
pass
global uniqueSession
with open("./configs/config.yaml", 'r', encoding='utf-8') as ymlfile:
cfg = yaml.safe_load(ymlfile)
if 'uniqueSessionMode' in cfg['qqbot'] and cfg['qqbot']['uniqueSessionMode'] == 'true':
uniqueSession = True
else:
uniqueSession = False
if cfg['qqbot']['appid'] != '' or cfg['qqbot']['token'] != '':
print("读取QQBot appid token 成功")
intents = botpy.Intents(public_guild_messages=True, direct_message=True)
client = botClient(intents=intents)
client.run(appid=cfg['qqbot']['appid'], token=cfg['qqbot']['token'])
else:
raise BaseException("请在config中完善你的appid和token")
async def get_chatGPT_response(prompts_str):
res = ''
usage = ''
res, usage = await chatgpt.chat(prompts_str)
# 处理结果文本
chatgpt_res = res.strip()
return res, usage
def get_prompts_by_cache_list(cache_data_list, divide=False, paging=False, size=5, page=1):
prompts = ""
if paging:
page_begin = (page-1)*size
page_end = page*size
if page_begin < 0:
page_begin = 0
if page_end > len(cache_data_list):
page_end = len(cache_data_list)
cache_data_list = cache_data_list[page_begin:page_end]
for item in cache_data_list:
prompts += str(item['prompt'])
if divide:
prompts += "----------\n"
return prompts
def get_user_usage_tokens(cache_list):
usage_tokens = 0
for item in cache_list:
usage_tokens += int(item['single_tokens'])
return usage_tokens
async def oper_msg(message, at=False):
print("[QQBOT] 接收到消息:"+ message.content)
qq_msg = ""
if at:
# 过滤用户id
pattern = r"<@!\d+>\s+(.+)"
# 多行匹配
pattern = re.compile(pattern, flags=re.MULTILINE)
result = re.search(pattern, message.content)
if result:
qq_msg = result.group(1).strip()
else:
qq_msg = message.content
user_id = message.author.id
if not at:
session_id = message.author.id
else:
if uniqueSession:
session_id = message.author.id
else:
session_id = message.guild_id
if session_id:
name = ''
if uniqueSession:
name = message.member.nick
else:
name = "频道"
if qq_msg == "/reset":
session_dict[session_id] = []
if at:
await message.reply(content=f"{name}(id: {session_id}) 的历史记录重置成功")
else:
await message.reply(content=f"你的历史记录重置成功")
return
if qq_msg[:4] == "/his":
#分页,每页5条
size_per_page = 3
page = 1
if qq_msg[5:]:
page = int(qq_msg[5:])
# 检查是否有过历史记录
if session_id not in session_dict:
await message.reply(content=f"{name} 的历史记录为空")
return
l = session_dict[session_id]
max_page = len(l)//size_per_page + 1 if len(l)%size_per_page != 0 else len(l)//size_per_page
p = get_prompts_by_cache_list(session_dict[session_id], divide=True, paging=True, size=size_per_page, page=page)
if at:
await message.reply(content=f"{name} 的历史记录如下:\n{p}\n{page}页 | 共{max_page}\n*输入/his 2跳转到第2页")
else:
await message.reply(content=f"历史记录如下:\n{p}\n{page}页 | 共{max_page}\n*输入/his 2跳转到第2页")
return
if qq_msg == "/token":
if at:
await message.reply(content=f"{name} 会话的token数: {get_user_usage_tokens(session_dict[session_id])}\n系统最大缓存token数: {max_tokens}")
else:
await message.reply(content=f"会话的token数: {get_user_usage_tokens(session_dict[session_id])}\n系统最大缓存token数: {max_tokens}")
return
if qq_msg == "/status":
chatgpt_cfg_str = ""
for k, v in gpt_config.items():
if k == "key":
continue
chatgpt_cfg_str += f"{k}: {v}"
await message.reply(content=f"ChatGPT配置:\n - {chatgpt_cfg_str}\n QQChannelChatGPT 版本: {version}")
return
if qq_msg == "/count":
try:
f = open("./configs/stat", "r", encoding="utf-8")
fjson = json.loads(f.read())
f.close()
guild_count = 0
guild_msg_count = 0
guild_direct_msg_count = 0
for k,v in fjson.items():
guild_count += 1
guild_msg_count += v['count']
guild_direct_msg_count += v['direct_count']
session_count = 0
f = open("./configs/session", "r", encoding="utf-8")
fjson = json.loads(f.read())
f.close()
for k,v in fjson.items():
session_count += 1
except:
pass
await message.reply(content=f"当前会话数: {len(session_dict)}\n共有频道数: {guild_count} \n共有消息数: {guild_msg_count}\n私信数: {guild_direct_msg_count}\n历史会话数: {session_count}")
return
if qq_msg == "/help":
await message.reply(content=f"请联系频道管理员或者前往github(仓库名: QQChannelChatGPT)提issue~")
return
if session_id not in session_dict:
session_dict[session_id] = []
fjson = {}
try:
f = open("./configs/session", "r", encoding="utf-8")
fjson = json.loads(f.read())
f.close()
except:
pass
finally:
fjson[session_id] = 'true'
f = open("./configs/session", "w", encoding="utf-8")
f.write(json.dumps(fjson))
f.flush()
f.close()
# 获取缓存
cache_prompt = ''
cache_data_list = session_dict[session_id]
cache_prompt = get_prompts_by_cache_list(cache_data_list)
cache_prompt += "Human: "+ qq_msg + "\nAI: "
# 请求chatGPT获得结果
try:
chatgpt_res, current_usage_tokens = await get_chatGPT_response(cache_prompt)
except (PromptExceededError) as e:
print("出现token超限, 清空对应缓存")
# 超过4097tokens错误,清空缓存
session_dict[session_id] = []
cache_data_list = []
cache_prompt = "Human: "+ qq_msg + "\nAI: "
chatgpt_res, current_usage_tokens = await get_chatGPT_response(cache_prompt)
except (BaseException) as e:
print("OpenAI API错误:(")
await message.reply(content=f"OpenAI API错误:( 原因如下:\n{str(e)} \n*前往github(仓库名: QQChannelChatGPT)反馈~")
# 超过指定tokens, 尽可能的保留最多的条目,直到小于max_tokens
# print("current_usage_tokens: ", current_usage_tokens)
# print("max_tokens: ", max_tokens)
if current_usage_tokens > max_tokens:
t = current_usage_tokens
cache_list = session_dict[session_id]
index = 0
while t > max_tokens:
if index >= len(cache_list):
break
t -= int(cache_list[index]['single_tokens'])
index += 1
session_dict[session_id] = cache_list[index:]
cache_data_list = session_dict[session_id]
cache_prompt = get_prompts_by_cache_list(cache_data_list)
# cache_prompt += chatgpt_res + "\n";
# 添加新条目进入缓存的prompt
if len(cache_data_list) > 0:
single_record = {
"prompt": f'Human: {qq_msg}\nAI: {chatgpt_res}\n',
"usage_tokens": current_usage_tokens,
"single_tokens": current_usage_tokens - int(cache_data_list[-1]['usage_tokens'])
}
else:
single_record = {
"prompt": f'Human: {qq_msg}\nAI: {chatgpt_res}\n',
"usage_tokens": current_usage_tokens,
"single_tokens": current_usage_tokens
}
# print(single_record)
cache_data_list.append(single_record)
# 写入数据库
# try:
# data = {
# "data": cache_data_list
# }
# data_str = json.dumps(data)
# if len(cache_data_list) > 1:
# db.update_session(session_id, data_str)
# else:
# db.insert_session(session_id, data_str)
# except Exception as e:
# print(e)
# print("数据库写入失败")
session_dict[session_id] = cache_data_list
# #检测是否存在url,如果存在,则去除url 防止被qq频道过滤
chatgpt_res = chatgpt_res.replace(".", " . ")
# 发送qq信息
try:
await message.reply(content=f"[ChatGPT]{chatgpt_res}")
except BaseException as e:
print("QQ频道API错误: \n"+str(e))
f_res = ""
for t in chatgpt_res:
f_res += t + ' '
await message.reply(content=f"QQ频道API错误:{str(e)}\n下面是格式化后的回答:\n{f_res}")