formatting code

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):

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)
+        return self.bots["text_to_voice"].textToVoice(text)

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
-