From 9e07703eb1264367d147d50c679b40c580278cee Mon Sep 17 00:00:00 2001 From: lanvent Date: Sun, 12 Mar 2023 01:22:49 +0800 Subject: [PATCH] formatting code --- bot/chatgpt/chat_gpt_bot.py | 58 +++++++++--------- bridge/bridge.py | 4 +- channel/wechat/wechat_channel.py | 101 ++++++++++++++++--------------- 3 files changed, 83 insertions(+), 80 deletions(-) diff --git a/bot/chatgpt/chat_gpt_bot.py b/bot/chatgpt/chat_gpt_bot.py index 327cc95..9bfea5b 100644 --- a/bot/chatgpt/chat_gpt_bot.py +++ b/bot/chatgpt/chat_gpt_bot.py @@ -13,28 +13,28 @@ class ChatGPTBot(Bot): def __init__(self): openai.api_key = conf().get('open_ai_api_key') proxy = conf().get('proxy') - self.sessions=SessionManager() + self.sessions = SessionManager() if proxy: openai.proxy = proxy def reply(self, query, context=None): # acquire reply content - if context['type']=='TEXT': + if context['type'] == 'TEXT': logger.info("[OPEN_AI] query={}".format(query)) session_id = context['session_id'] - reply=None + reply = None if query == '#清除记忆': self.sessions.clear_session(session_id) - reply={'type':'INFO', 'content':'记忆已清除'} + reply = {'type': 'INFO', 'content': '记忆已清除'} elif query == '#清除所有': self.sessions.clear_all_session() - reply={'type':'INFO', 'content':'所有人记忆已清除'} + reply = {'type': 'INFO', 'content': '所有人记忆已清除'} elif query == '#更新配置': load_config() - reply={'type':'INFO', 'content':'配置已更新'} + reply = {'type': 'INFO', 'content': '配置已更新'} elif query == '#DEBUG': logger.setLevel('DEBUG') - reply={'type':'INFO', 'content':'DEBUG模式已开启'} + reply = {'type': 'INFO', 'content': 'DEBUG模式已开启'} if reply: return reply session = self.sessions.build_session_query(query, session_id) @@ -46,28 +46,28 @@ class ChatGPTBot(Bot): reply_content = self.reply_text(session, session_id, 0) logger.debug("[OPEN_AI] new_query={}, session_id={}, reply_cont={}".format(session, session_id, reply_content["content"])) - if reply_content['completion_tokens']==0 and len(reply_content['content'])>0: - reply={'type':'ERROR', 'content':reply_content['content']} + if reply_content['completion_tokens'] == 0 and len(reply_content['content']) > 0: + reply = {'type': 'ERROR', 'content': reply_content['content']} elif reply_content["completion_tokens"] > 0: self.sessions.save_session(reply_content["content"], session_id, reply_content["total_tokens"]) reply={'type':'TEXT', 'content':reply_content["content"]} else: - reply={'type':'ERROR', 'content':reply_content['content']} + reply = {'type': 'ERROR', 'content': reply_content['content']} logger.debug("[OPEN_AI] reply {} used 0 tokens.".format(reply_content)) return reply elif context['type'] == 'IMAGE_CREATE': - ok, retstring=self.create_img(query, 0) - reply=None + ok, retstring = self.create_img(query, 0) + reply = None if ok: - reply = {'type':'IMAGE', 'content':retstring} + reply = {'type': 'IMAGE', 'content': retstring} else: - reply = {'type':'ERROR', 'content':retstring} + reply = {'type': 'ERROR', 'content': retstring} return reply else: reply= {'type':'ERROR', 'content':'Bot不支持处理{}类型的消息'.format(context['type'])} - def reply_text(self, session, session_id, retry_count=0) ->dict: + def reply_text(self, session, session_id, retry_count=0) -> dict: ''' call openai's ChatCompletion to get the answer :param session: a conversation session @@ -86,8 +86,8 @@ class ChatGPTBot(Bot): presence_penalty=0.0, # [-2,2]之间，该值越大则更倾向于产生不同的内容 ) # logger.info("[ChatGPT] reply={}, total_tokens={}".format(response.choices[0]['message']['content'], response["usage"]["total_tokens"])) - return {"total_tokens": response["usage"]["total_tokens"], - "completion_tokens": response["usage"]["completion_tokens"], + return {"total_tokens": response["usage"]["total_tokens"], + "completion_tokens": response["usage"]["completion_tokens"], "content": response.choices[0]['message']['content']} except openai.error.RateLimitError as e: # rate limit exception @@ -102,11 +102,11 @@ class ChatGPTBot(Bot): # api connection exception logger.warn(e) logger.warn("[OPEN_AI] APIConnection failed") - return {"completion_tokens": 0, "content":"我连接不到你的网络"} + return {"completion_tokens": 0, "content": "我连接不到你的网络"} except openai.error.Timeout as e: logger.warn(e) logger.warn("[OPEN_AI] Timeout") - return {"completion_tokens": 0, "content":"我没有收到你的消息"} + return {"completion_tokens": 0, "content": "我没有收到你的消息"} except Exception as e: # unknown exception logger.exception(e) @@ -123,7 +123,7 @@ class ChatGPTBot(Bot): ) image_url = response['data'][0]['url'] logger.info("[OPEN_AI] image_url={}".format(image_url)) - return True,image_url + return True, image_url except openai.error.RateLimitError as e: logger.warn(e) if retry_count < 1: @@ -131,15 +131,17 @@ class ChatGPTBot(Bot): logger.warn("[OPEN_AI] ImgCreate RateLimit exceed, 第{}次重试".format(retry_count+1)) return self.create_img(query, retry_count+1) else: - return False,"提问太快啦，请休息一下再问我吧" + return False, "提问太快啦，请休息一下再问我吧" except Exception as e: logger.exception(e) - return False,str(e) - + return False, str(e) + + class SessionManager(object): def __init__(self): self.sessions = {} - def build_session_query(self,query, session_id): + + def build_session_query(self, query, session_id): ''' build query with conversation history e.g. [ @@ -167,7 +169,7 @@ class SessionManager(object): if not max_tokens: # default 3000 max_tokens = 1000 - max_tokens=int(max_tokens) + max_tokens = int(max_tokens) session = self.sessions.get(session_id) if session: @@ -177,7 +179,7 @@ class SessionManager(object): # discard exceed limit conversation self.discard_exceed_conversation(session, max_tokens, total_tokens) - + def discard_exceed_conversation(self, session, max_tokens, total_tokens): dec_tokens = int(total_tokens) # logger.info("prompt tokens used={},max_tokens={}".format(used_tokens,max_tokens)) @@ -187,10 +189,10 @@ class SessionManager(object): session.pop(1) session.pop(1) else: - break + break dec_tokens = dec_tokens - max_tokens - def clear_session(self,session_id): + def clear_session(self, session_id): self.sessions[session_id] = [] def clear_all_session(self): diff --git a/bridge/bridge.py b/bridge/bridge.py index 392d9e8..81e5d73 100644 --- a/bridge/bridge.py +++ b/bridge/bridge.py @@ -2,6 +2,7 @@ from bot import bot_factory from common.singleton import singleton from voice import voice_factory + @singleton class Bridge(object): def __init__(self): @@ -15,7 +16,6 @@ class Bridge(object): except ModuleNotFoundError as e: print(e) - # 以下所有函数需要得到一个reply字典，格式如下： # reply["type"] = "ERROR" / "TEXT" / "VOICE" / ... # reply["content"] = reply的内容 @@ -27,4 +27,4 @@ class Bridge(object): return self.bots["voice_to_text"].voiceToText(voiceFile) def fetch_text_to_voice(self, text): - return self.bots["text_to_voice"].textToVoice(text) \ No newline at end of file + return self.bots["text_to_voice"].textToVoice(text) diff --git a/channel/wechat/wechat_channel.py b/channel/wechat/wechat_channel.py index ddf38a4..e8be17e 100644 --- a/channel/wechat/wechat_channel.py +++ b/channel/wechat/wechat_channel.py @@ -46,7 +46,7 @@ class WechatChannel(Channel): # start message listener itchat.run() - + # handle_* 系列函数处理收到的消息后构造context，然后调用handle函数处理context # context是一个字典，包含了消息的所有信息，包括以下key # type: 消息类型，包括TEXT、VOICE、CMD_IMAGE_CREATE @@ -57,18 +57,17 @@ class WechatChannel(Channel): # receiver: 需要回复的对象 def handle_voice(self, msg): - if conf().get('speech_recognition') != True : + if conf().get('speech_recognition') != True: return logger.debug("[WX]receive voice msg: " + msg['FileName']) from_user_id = msg['FromUserName'] other_user_id = msg['User']['UserName'] if from_user_id == other_user_id: - context = { 'isgroup': False, 'msg': msg, 'receiver': other_user_id} - context['type']='VOICE' - context['session_id']=other_user_id + context = {'isgroup': False, 'msg': msg, 'receiver': other_user_id} + context['type'] = 'VOICE' + context['session_id'] = other_user_id thread_pool.submit(self.handle, context) - def handle_text(self, msg): logger.debug("[WX]receive text msg: " + json.dumps(msg, ensure_ascii=False)) content = msg['Text'] @@ -80,22 +79,21 @@ class WechatChannel(Channel): logger.debug("[WX]reference query skipped") return if match_prefix: - content=content.replace(match_prefix,'',1).strip() + content = content.replace(match_prefix, '', 1).strip() else: return - context = { 'isgroup': False, 'msg': msg, 'receiver': other_user_id} - context['session_id']=other_user_id - + context = {'isgroup': False, 'msg': msg, 'receiver': other_user_id} + context['session_id'] = other_user_id + img_match_prefix = check_prefix(content, conf().get('image_create_prefix')) if img_match_prefix: - content=content.replace(img_match_prefix,'',1).strip() - context['type']='CMD_IMAGE_CREATE' + content = content.replace(img_match_prefix, '', 1).strip() + context['type'] = 'CMD_IMAGE_CREATE' else: - context['type']='TEXT' - - context['content']=content - thread_pool.submit(self.handle, context) + context['type'] = 'TEXT' + context['content'] = content + thread_pool.submit(self.handle, context) def handle_group(self, msg): logger.debug("[WX]receive group msg: " + json.dumps(msg, ensure_ascii=False)) @@ -122,32 +120,32 @@ class WechatChannel(Channel): img_match_prefix = check_prefix(content, conf().get('image_create_prefix')) if img_match_prefix: - content=content.replace(img_match_prefix,'',1).strip() - context['type']='CMD_IMAGE_CREATE' + content = content.replace(img_match_prefix, '', 1).strip() + context['type'] = 'CMD_IMAGE_CREATE' else: - context['type']='TEXT' - context['content']=content + context['type'] = 'TEXT' + context['content'] = content group_chat_in_one_session = conf().get('group_chat_in_one_session', []) - if ('ALL_GROUP' in group_chat_in_one_session or \ - group_name in group_chat_in_one_session or \ + if ('ALL_GROUP' in group_chat_in_one_session or + group_name in group_chat_in_one_session or check_contain(group_name, group_chat_in_one_session)): context['session_id'] = group_id else: context['session_id'] = msg['ActualUserName'] - + thread_pool.submit(self.handle, context) - + # 统一的发送函数，根据reply的type字段发送不同类型的消息 def send(self, reply, receiver): - if reply['type']=='TEXT': + if reply['type'] == 'TEXT': itchat.send(reply['content'], toUserName=receiver) logger.info('[WX] sendMsg={}, receiver={}'.format(reply, receiver)) - elif reply['type']=='ERROR' or reply['type']=='INFO': + elif reply['type'] == 'ERROR' or reply['type'] == 'INFO': itchat.send(reply['content'], toUserName=receiver) logger.info('[WX] sendMsg={}, receiver={}'.format(reply, receiver)) - elif reply['type']=='VOICE': + elif reply['type'] == 'VOICE': itchat.send_file(reply['content'], toUserName=receiver) logger.info('[WX] sendFile={}, receiver={}'.format(reply['content'], receiver)) elif reply['type']=='IMAGE_URL': # 从网络下载图片 @@ -164,52 +162,56 @@ class WechatChannel(Channel): image_storage.seek(0) itchat.send_image(image_storage, toUserName=receiver) logger.info('[WX] sendImage, receiver={}'.format(receiver)) - + # 处理消息 def handle(self, context): - content=context['content'] - reply=None + content = context['content'] + reply = None logger.debug('[WX] ready to handle context: {}'.format(context)) # reply的构建步骤 - if context['type']=='TEXT' or context['type']=='CMD_IMAGE_CREATE': - reply = super().build_reply_content(content,context) - elif context['type']=='VOICE': - msg=context['msg'] + if context['type'] == 'TEXT' or context['type'] == 'CMD_IMAGE_CREATE': + reply = super().build_reply_content(content, context) + elif context['type'] == 'VOICE': + msg = context['msg'] file_name = TmpDir().path() + msg['FileName'] msg.download(file_name) reply = super().build_voice_to_text(file_name) if reply['type'] != 'ERROR' and reply['type'] != 'INFO': - reply = super().build_reply_content(reply['content'],context) - if reply['type']=='TEXT': + reply = super().build_reply_content(reply['content'], context) + if reply['type'] == 'TEXT': if conf().get('voice_reply_voice'): reply = super().build_text_to_voice(reply['content']) else: logger.error('[WX] unknown context type: {}'.format(context['type'])) return - + logger.debug('[WX] ready to decorate reply: {}'.format(reply)) # reply的包装步骤 if reply: - if reply['type']=='TEXT': - reply_text=reply['content'] + if reply['type'] == 'TEXT': + reply_text = reply['content'] if context['isgroup']: - reply_text = '@' + context['msg']['ActualNickName'] + ' ' + reply_text.strip() - reply_text=conf().get("group_chat_reply_prefix","")+reply_text + reply_text = '@' + \ + context['msg']['ActualNickName'] + \ + ' ' + reply_text.strip() + reply_text = conf().get("group_chat_reply_prefix", "")+reply_text else: - reply_text=conf().get("single_chat_reply_prefix","")+reply_text - reply['content']=reply_text - elif reply['type']=='ERROR' or reply['type']=='INFO': - reply['content']=reply['type']+": "+ reply['content'] - elif reply['type']=='IMAGE_URL' or reply['type']=='VOICE': + reply_text = conf().get("single_chat_reply_prefix", "")+reply_text + reply['content'] = reply_text + elif reply['type'] == 'ERROR' or reply['type'] == 'INFO': + reply['content'] = reply['type']+": " + reply['content'] + elif reply['type'] == 'IMAGE_URL' or reply['type'] == 'VOICE': pass else: - logger.error('[WX] unknown reply type: {}'.format(reply['type'])) + logger.error( + '[WX] unknown reply type: {}'.format(reply['type'])) return if reply: - logger.debug('[WX] ready to send reply: {} to {}'.format(reply,context['receiver'])) + logger.debug('[WX] ready to send reply: {} to {}'.format( + reply, context['receiver'])) self.send(reply, context['receiver']) - + def check_prefix(content, prefix_list): for prefix in prefix_list: @@ -225,4 +227,3 @@ def check_contain(content, keyword_list): if content.find(ky) != -1: return True return None -