Support E2E encrypted room

This commit is contained in:
hibobmaster 2023-04-10 19:37:43 +08:00
parent 4832d6f00b
commit 4755653e7e
7 changed files with 65 additions and 55 deletions

View file

@ -18,11 +18,11 @@ class askGPT:
}, },
], ],
} }
max_try = 5 max_try = 3
while max_try > 0: while max_try > 0:
try: try:
async with self.session.post(url=api_endpoint, 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 status_code = response.status
if not status_code == 200: if not status_code == 200:
# print failed reason # print failed reason

View file

@ -22,10 +22,10 @@ class BingBot:
async def ask_bing(self, prompt) -> str: async def ask_bing(self, prompt) -> str:
self.data['message'] = prompt self.data['message'] = prompt
max_try = 5 max_try = 3
while max_try > 0: while max_try > 0:
try: 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 status_code = resp.status
body = await resp.read() body = await resp.read()
if not status_code == 200: if not status_code == 200:

35
bot.py
View file

@ -71,11 +71,6 @@ class Bot:
if self.access_token is not None: if self.access_token is not None:
self.client.access_token = self.access_token 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}] # regular expression to match keyword [!gpt {prompt}] [!chat {prompt}]
self.gpt_prog = re.compile(r"^\s*!gpt\s*(.+)$") self.gpt_prog = re.compile(r"^\s*!gpt\s*(.+)$")
self.chat_prog = re.compile(r"^\s*!chat\s*(.+)$") self.chat_prog = re.compile(r"^\s*!chat\s*(.+)$")
@ -85,7 +80,7 @@ class Bot:
# initialize chatbot and chatgpt_api_endpoint # initialize chatbot and chatgpt_api_endpoint
if self.api_key != '': 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 self.chatgpt_api_endpoint = self.chatgpt_api_endpoint
# request header for !gpt command # request header for !gpt command
@ -153,7 +148,10 @@ class Bot:
m = self.gpt_prog.match(content_body) m = self.gpt_prog.match(content_body)
if m: if m:
prompt = m.group(1) prompt = m.group(1)
try:
await self.gpt(room_id, reply_to_event_id, prompt, sender_id, raw_user_message) await self.gpt(room_id, reply_to_event_id, prompt, sender_id, raw_user_message)
except Exception as e:
logger.error(e)
# bing ai # bing ai
if self.bing_api_endpoint != '': if self.bing_api_endpoint != '':
@ -407,9 +405,9 @@ class Bot:
# !chat command # !chat command
async def chat(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message): 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: 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: except TimeoutError as e:
logger.error("timeoutException", exc_info=True) logger.error("timeoutException", exc_info=True)
text = "Timeout error" 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): async def gpt(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message):
try: try:
# sending typing state # sending typing state
await self.client.room_typing(room_id, timeout=120000) await self.client.room_typing(room_id, timeout=180000)
# timeout 120s # 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: except TimeoutError:
logger.error("timeoutException", exc_info=True) logger.error("timeoutException", exc_info=True)
text = "Timeout error" 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): async def bing(self, room_id, reply_to_event_id, prompt, sender_id, raw_content_body):
try: try:
# sending typing state # sending typing state
await self.client.room_typing(room_id, timeout=120000) await self.client.room_typing(room_id, timeout=180000)
# timeout 120s # 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: except TimeoutError:
logger.error("timeoutException", exc_info=True) logger.error("timeoutException", exc_info=True)
text = "Timeout error" text = "Timeout error"
@ -481,7 +479,8 @@ class Bot:
help_info = "!gpt [content], generate response without context conversation\n" + \ help_info = "!gpt [content], generate response without context conversation\n" + \
"!chat [content], chat with context conversation\n" + \ "!chat [content], chat with context conversation\n" + \
"!bing [content], chat with context conversation powered by Bing AI\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) await send_room_message(self.client, room_id, reply_message=help_info)
except Exception as e: except Exception as e:
@ -499,16 +498,20 @@ class Bot:
logger.error(f"Error: {e}", exc_info=True) logger.error(f"Error: {e}", exc_info=True)
# sync messages in the room # 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) await self.client.sync_forever(timeout=timeout, full_state=full_state)
# Sync encryption keys with the server # 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: if self.client.should_upload_keys:
await self.client.keys_upload() await self.client.keys_upload()
# Trust own devices # 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) await self.client.sync(timeout=30000, full_state=True)
for device_id, olm_device in self.client.device_store[ for device_id, olm_device in self.client.device_store[
self.user_id].items(): self.user_id].items():

5
log.py
View file

@ -7,20 +7,25 @@ def getlogger():
# create handlers # create handlers
warn_handler = logging.StreamHandler() warn_handler = logging.StreamHandler()
info_handler = logging.StreamHandler()
error_handler = logging.FileHandler('bot.log', mode='a') error_handler = logging.FileHandler('bot.log', mode='a')
warn_handler.setLevel(logging.WARNING) warn_handler.setLevel(logging.WARNING)
error_handler.setLevel(logging.ERROR) error_handler.setLevel(logging.ERROR)
info_handler.setLevel(logging.INFO)
# create formatters # create formatters
warn_format = logging.Formatter('%(name)s - %(funcName)s - %(levelname)s - %(message)s') 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') error_format = logging.Formatter('%(asctime)s - %(name)s - %(funcName)s - %(levelname)s - %(message)s')
info_format = logging.Formatter('%(message)s')
# set formatter # set formatter
warn_handler.setFormatter(warn_format) warn_handler.setFormatter(warn_format)
error_handler.setFormatter(error_format) error_handler.setFormatter(error_format)
info_handler.setFormatter(info_format)
# add handlers to logger # add handlers to logger
logger.addHandler(warn_handler) logger.addHandler(warn_handler)
logger.addHandler(error_handler) logger.addHandler(error_handler)
logger.addHandler(info_handler)
return logger return logger

47
main.py
View file

@ -1,28 +1,26 @@
#!/usr/bin/env python3
import json import json
import os import os
import asyncio import asyncio
from bot import Bot from bot import Bot
from nio import Api, SyncResponse
from log import getlogger from log import getlogger
logger = getlogger() logger = getlogger()
async def main(): async def main():
if os.path.exists('config.json'):
fp = open('config.json', 'r', encoding="utf8") fp = open('config.json', 'r', encoding="utf8")
config = json.load(fp) config = json.load(fp)
matrix_bot = Bot(homeserver=os.environ.get("HOMESERVER") or config.get('homeserver'), matrix_bot = Bot(homeserver=config.get('homeserver'),
user_id=os.environ.get("USER_ID") or config.get('user_id') , user_id=config.get('user_id') ,
password=os.environ.get("PASSWORD") or config.get('password'), password=config.get('password'),
device_id=os.environ.get("DEVICE_ID") or config.get('device_id'), device_id=config.get('device_id'),
room_id=os.environ.get("ROOM_ID") or config.get('room_id'), room_id=config.get('room_id'),
api_key=os.environ.get("OPENAI_API_KEY") or config.get('api_key'), api_key=config.get('api_key'),
bing_api_endpoint=os.environ.get("BING_API_ENDPOINT") or config.get('bing_api_endpoint'), bing_api_endpoint=config.get('bing_api_endpoint'),
access_token=os.environ.get("ACCESS_TOKEN") or config.get('access_token'), access_token=config.get('access_token'),
jailbreakEnabled=os.environ.get("JAILBREAKENABLED", "False").lower() in ('true', '1') or config.get('jailbreakEnabled'), jailbreakEnabled=config.get('jailbreakEnabled'),
bing_auth_cookie=os.environ.get("BING_AUTH_COOKIE") or config.get('bing_auth_cookie'), bing_auth_cookie=config.get('bing_auth_cookie'),
) )
# if not set access_token, then login via password # if not set access_token, then login via password
# if os.path.exists('config.json'): # if os.path.exists('config.json'):
@ -36,21 +34,20 @@ async def main():
await matrix_bot.login() 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=30000, full_state=True)
await matrix_bot.sync_forever(timeout=3000, full_state=True)
finally:
await matrix_bot.client.close()
if __name__ == "__main__": if __name__ == "__main__":
logger.debug("matrix chatgpt bot start.....") print("matrix chatgpt bot start.....")
try: # try:
loop = asyncio.get_running_loop() # loop = asyncio.get_running_loop()
except RuntimeError: # except RuntimeError:
loop = asyncio.new_event_loop() # loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop) # asyncio.set_event_loop(loop)
asyncio.run(main()) asyncio.run(main())

View file

@ -7,11 +7,12 @@ async-timeout==4.0.2
atomicwrites==1.4.1 atomicwrites==1.4.1
attrs==22.2.0 attrs==22.2.0
blobfile==2.0.1 blobfile==2.0.1
cachetools==5.3.0 cachetools==4.2.4
certifi==2022.12.7 certifi==2022.12.7
cffi==1.15.1 cffi==1.15.1
charset-normalizer==3.1.0 charset-normalizer==3.1.0
filelock==3.10.7 cryptography==40.0.1
filelock==3.11.0
frozenlist==1.3.3 frozenlist==1.3.3
future==0.18.3 future==0.18.3
h11==0.14.0 h11==0.14.0
@ -24,7 +25,8 @@ idna==3.4
jsonschema==4.17.3 jsonschema==4.17.3
Logbook==1.5.3 Logbook==1.5.3
lxml==4.9.2 lxml==4.9.2
matrix-nio[e2e] Markdown==3.4.3
matrix-nio[e2e]==0.20.2
multidict==6.0.4 multidict==6.0.4
peewee==3.16.0 peewee==3.16.0
Pillow==9.5.0 Pillow==9.5.0
@ -32,15 +34,18 @@ pycparser==2.21
pycryptodome==3.17 pycryptodome==3.17
pycryptodomex==3.17 pycryptodomex==3.17
pyrsistent==0.19.3 pyrsistent==0.19.3
python-cryptography-fernet-wrapper==1.0.4
python-magic==0.4.27 python-magic==0.4.27
python-olm==3.1.3
python-socks==2.2.0 python-socks==2.2.0
regex==2023.3.23 regex==2023.3.23
requests==2.28.2 requests==2.28.2
rfc3986==1.5.0 rfc3986==1.5.0
six==1.16.0
sniffio==1.3.0 sniffio==1.3.0
tiktoken==0.3.3
toml==0.10.2
unpaddedbase64==2.1.0 unpaddedbase64==2.1.0
urllib3==1.26.15 urllib3==1.26.15
wcwidth==0.2.6 wcwidth==0.2.6
yarl==1.8.2 yarl==1.8.2
python-olm >= '3.1.0'
tiktoken==0.3.3

View file

@ -11,8 +11,8 @@ async def send_room_message(client: AsyncClient,
else: else:
body = r'> <' + sender_id + r'> ' + user_message + r'\n\n' + reply_message body = r'> <' + sender_id + r'> ' + user_message + r'\n\n' + reply_message
format = r'org.matrix.custom.html' format = r'org.matrix.custom.html'
formatted_body = r'<mx-reply><blockquote><a href=\"https://matrix.to/#/' + room_id + r'/' + reply_to_event_id formatted_body = r'<mx-reply><blockquote><a href="https://matrix.to/#/' + room_id + r'/' + reply_to_event_id \
+ r'\">In reply to</a> <a href=\"https://matrix.to/#/' + sender_id + r'\">' + sender_id + r'">In reply to</a> <a href="https://matrix.to/#/' + sender_id + r'">' + sender_id \
+ r'</a><br>' + user_message + r'</blockquote></mx-reply>' + reply_message + r'</a><br>' + user_message + r'</blockquote></mx-reply>' + reply_message
content={"msgtype": "m.text", "body": body, "format": format, "formatted_body": formatted_body, content={"msgtype": "m.text", "body": body, "format": format, "formatted_body": formatted_body,