From 4755653e7e1537b95216eeff3abe92a7a395829e Mon Sep 17 00:00:00 2001 From: hibobmaster Date: Mon, 10 Apr 2023 19:37:43 +0800 Subject: [PATCH] Support E2E encrypted room --- askgpt.py | 4 ++-- bing.py | 4 ++-- bot.py | 37 +++++++++++++++++++---------------- log.py | 5 +++++ main.py | 51 +++++++++++++++++++++++------------------------- requirements.txt | 15 +++++++++----- send_message.py | 4 ++-- 7 files changed, 65 insertions(+), 55 deletions(-) diff --git a/askgpt.py b/askgpt.py index dd25ef7..9cca2d3 100644 --- a/askgpt.py +++ b/askgpt.py @@ -18,11 +18,11 @@ class askGPT: }, ], } - max_try = 5 + max_try = 3 while max_try > 0: try: async with self.session.post(url=api_endpoint, - json=jsons, headers=headers, timeout=30) as response: + json=jsons, headers=headers, timeout=60) as response: status_code = response.status if not status_code == 200: # print failed reason diff --git a/bing.py b/bing.py index 643b648..15350cc 100644 --- a/bing.py +++ b/bing.py @@ -22,10 +22,10 @@ class BingBot: async def ask_bing(self, prompt) -> str: self.data['message'] = prompt - max_try = 5 + max_try = 3 while max_try > 0: try: - resp = await self.session.post(url=self.bing_api_endpoint, json=self.data) + resp = await self.session.post(url=self.bing_api_endpoint, json=self.data, timeout=60) status_code = resp.status body = await resp.read() if not status_code == 200: diff --git a/bot.py b/bot.py index 1a95f27..030d4b1 100644 --- a/bot.py +++ b/bot.py @@ -71,11 +71,6 @@ class Bot: if self.access_token is not None: self.client.access_token = self.access_token - # setup event callbacks - self.client.add_event_callback(self.message_callback, (RoomMessageText, )) - self.client.add_event_callback(self.invite_callback, (InviteMemberEvent, )) - self.client.add_to_device_callback(self.to_device_callback, (KeyVerificationEvent, )) - # regular expression to match keyword [!gpt {prompt}] [!chat {prompt}] self.gpt_prog = re.compile(r"^\s*!gpt\s*(.+)$") self.chat_prog = re.compile(r"^\s*!chat\s*(.+)$") @@ -85,7 +80,7 @@ class Bot: # initialize chatbot and chatgpt_api_endpoint if self.api_key != '': - self.chatbot = Chatbot(api_key=self.api_key) + self.chatbot = Chatbot(api_key=self.api_key, timeout=60) self.chatgpt_api_endpoint = self.chatgpt_api_endpoint # request header for !gpt command @@ -153,7 +148,10 @@ class Bot: m = self.gpt_prog.match(content_body) if m: prompt = m.group(1) - await self.gpt(room_id, reply_to_event_id, prompt, sender_id, raw_user_message) + try: + await self.gpt(room_id, reply_to_event_id, prompt, sender_id, raw_user_message) + except Exception as e: + logger.error(e) # bing ai if self.bing_api_endpoint != '': @@ -407,9 +405,9 @@ class Bot: # !chat command async def chat(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message): - await self.client.room_typing(room_id, timeout=120000) + await self.client.room_typing(room_id, timeout=180000) try: - text = await asyncio.wait_for(self.chatbot.ask_async(prompt), timeout=120) + text = await asyncio.wait_for(self.chatbot.ask_async(prompt), timeout=180) except TimeoutError as e: logger.error("timeoutException", exc_info=True) text = "Timeout error" @@ -428,9 +426,9 @@ class Bot: async def gpt(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message): try: # sending typing state - await self.client.room_typing(room_id, timeout=120000) + await self.client.room_typing(room_id, timeout=180000) # timeout 120s - text = await asyncio.wait_for(self.askgpt.oneTimeAsk(prompt, self.chatgpt_api_endpoint, self.headers), timeout=120) + text = await asyncio.wait_for(self.askgpt.oneTimeAsk(prompt, self.chatgpt_api_endpoint, self.headers), timeout=180) except TimeoutError: logger.error("timeoutException", exc_info=True) text = "Timeout error" @@ -446,9 +444,9 @@ class Bot: async def bing(self, room_id, reply_to_event_id, prompt, sender_id, raw_content_body): try: # sending typing state - await self.client.room_typing(room_id, timeout=120000) + await self.client.room_typing(room_id, timeout=180000) # timeout 120s - text = await asyncio.wait_for(self.bingbot.ask_bing(prompt), timeout=120) + text = await asyncio.wait_for(self.bingbot.ask_bing(prompt), timeout=180) except TimeoutError: logger.error("timeoutException", exc_info=True) text = "Timeout error" @@ -481,7 +479,8 @@ class Bot: help_info = "!gpt [content], generate response without context conversation\n" + \ "!chat [content], chat with context conversation\n" + \ "!bing [content], chat with context conversation powered by Bing AI\n" + \ - "!pic [prompt], Image generation by Microsoft Bing" + "!pic [prompt], Image generation by Microsoft Bing\n" + \ + "!help, help message" await send_room_message(self.client, room_id, reply_message=help_info) except Exception as e: @@ -499,16 +498,20 @@ class Bot: logger.error(f"Error: {e}", exc_info=True) # sync messages in the room - async def sync_forever(self, timeout=30000, full_state=True): + async def sync_forever(self, timeout=30000, full_state=True) -> None: + # setup event callbacks + self.client.add_event_callback(self.message_callback, RoomMessageText) + self.client.add_event_callback(self.invite_callback, InviteMemberEvent) + self.client.add_to_device_callback(self.to_device_callback, KeyVerificationEvent) await self.client.sync_forever(timeout=timeout, full_state=full_state) # Sync encryption keys with the server - async def sync_encryption_key(self): + async def sync_encryption_key(self) -> None: if self.client.should_upload_keys: await self.client.keys_upload() # Trust own devices - async def trust_own_devices(self): + async def trust_own_devices(self) -> None: await self.client.sync(timeout=30000, full_state=True) for device_id, olm_device in self.client.device_store[ self.user_id].items(): diff --git a/log.py b/log.py index eb66833..f56a19a 100644 --- a/log.py +++ b/log.py @@ -7,20 +7,25 @@ def getlogger(): # create handlers warn_handler = logging.StreamHandler() + info_handler = logging.StreamHandler() error_handler = logging.FileHandler('bot.log', mode='a') warn_handler.setLevel(logging.WARNING) error_handler.setLevel(logging.ERROR) + info_handler.setLevel(logging.INFO) # create formatters warn_format = logging.Formatter('%(name)s - %(funcName)s - %(levelname)s - %(message)s') error_format = logging.Formatter('%(asctime)s - %(name)s - %(funcName)s - %(levelname)s - %(message)s') + info_format = logging.Formatter('%(message)s') # set formatter warn_handler.setFormatter(warn_format) error_handler.setFormatter(error_format) + info_handler.setFormatter(info_format) # add handlers to logger logger.addHandler(warn_handler) logger.addHandler(error_handler) + logger.addHandler(info_handler) return logger diff --git a/main.py b/main.py index 477e301..8253769 100644 --- a/main.py +++ b/main.py @@ -1,28 +1,26 @@ -#!/usr/bin/env python3 import json import os import asyncio from bot import Bot -from nio import Api, SyncResponse from log import getlogger logger = getlogger() async def main(): - if os.path.exists('config.json'): - fp = open('config.json', 'r', encoding="utf8") - config = json.load(fp) - matrix_bot = Bot(homeserver=os.environ.get("HOMESERVER") or config.get('homeserver'), - user_id=os.environ.get("USER_ID") or config.get('user_id') , - password=os.environ.get("PASSWORD") or config.get('password'), - device_id=os.environ.get("DEVICE_ID") or config.get('device_id'), - room_id=os.environ.get("ROOM_ID") or config.get('room_id'), - api_key=os.environ.get("OPENAI_API_KEY") or config.get('api_key'), - bing_api_endpoint=os.environ.get("BING_API_ENDPOINT") or config.get('bing_api_endpoint'), - access_token=os.environ.get("ACCESS_TOKEN") or config.get('access_token'), - jailbreakEnabled=os.environ.get("JAILBREAKENABLED", "False").lower() in ('true', '1') or config.get('jailbreakEnabled'), - bing_auth_cookie=os.environ.get("BING_AUTH_COOKIE") or config.get('bing_auth_cookie'), + fp = open('config.json', 'r', encoding="utf8") + config = json.load(fp) + + matrix_bot = Bot(homeserver=config.get('homeserver'), + user_id=config.get('user_id') , + password=config.get('password'), + device_id=config.get('device_id'), + room_id=config.get('room_id'), + api_key=config.get('api_key'), + bing_api_endpoint=config.get('bing_api_endpoint'), + access_token=config.get('access_token'), + jailbreakEnabled=config.get('jailbreakEnabled'), + bing_auth_cookie=config.get('bing_auth_cookie'), ) # if not set access_token, then login via password # if os.path.exists('config.json'): @@ -36,21 +34,20 @@ async def main(): await matrix_bot.login() - # await matrix_bot.sync_encryption_key() + await matrix_bot.sync_encryption_key() - # await matrix_bot.trust_own_devices() + await matrix_bot.trust_own_devices() - try: - await matrix_bot.sync_forever(timeout=3000, full_state=True) - finally: - await matrix_bot.client.close() + await matrix_bot.sync_forever(timeout=30000, full_state=True) if __name__ == "__main__": - logger.debug("matrix chatgpt bot start.....") - try: - loop = asyncio.get_running_loop() - except RuntimeError: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) + print("matrix chatgpt bot start.....") + # try: + # loop = asyncio.get_running_loop() + # except RuntimeError: + # loop = asyncio.new_event_loop() + # asyncio.set_event_loop(loop) asyncio.run(main()) + + diff --git a/requirements.txt b/requirements.txt index 47c538f..84ee0c1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,11 +7,12 @@ async-timeout==4.0.2 atomicwrites==1.4.1 attrs==22.2.0 blobfile==2.0.1 -cachetools==5.3.0 +cachetools==4.2.4 certifi==2022.12.7 cffi==1.15.1 charset-normalizer==3.1.0 -filelock==3.10.7 +cryptography==40.0.1 +filelock==3.11.0 frozenlist==1.3.3 future==0.18.3 h11==0.14.0 @@ -24,7 +25,8 @@ idna==3.4 jsonschema==4.17.3 Logbook==1.5.3 lxml==4.9.2 -matrix-nio[e2e] +Markdown==3.4.3 +matrix-nio[e2e]==0.20.2 multidict==6.0.4 peewee==3.16.0 Pillow==9.5.0 @@ -32,15 +34,18 @@ pycparser==2.21 pycryptodome==3.17 pycryptodomex==3.17 pyrsistent==0.19.3 +python-cryptography-fernet-wrapper==1.0.4 python-magic==0.4.27 +python-olm==3.1.3 python-socks==2.2.0 regex==2023.3.23 requests==2.28.2 rfc3986==1.5.0 +six==1.16.0 sniffio==1.3.0 +tiktoken==0.3.3 +toml==0.10.2 unpaddedbase64==2.1.0 urllib3==1.26.15 wcwidth==0.2.6 yarl==1.8.2 -python-olm >= '3.1.0' -tiktoken==0.3.3 diff --git a/send_message.py b/send_message.py index 58f4aa5..76ac1fd 100644 --- a/send_message.py +++ b/send_message.py @@ -11,8 +11,8 @@ async def send_room_message(client: AsyncClient, else: body = r'> <' + sender_id + r'> ' + user_message + r'\n\n' + reply_message format = r'org.matrix.custom.html' - formatted_body = r'
In reply to ' + sender_id + formatted_body = r'
In reply to ' + sender_id \ + r'
' + user_message + r'
' + reply_message content={"msgtype": "m.text", "body": body, "format": format, "formatted_body": formatted_body,