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

View File

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

37
bot.py
View File

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

5
log.py
View File

@ -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

51
main.py
View File

@ -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())

View File

@ -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

View File

@ -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'<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
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'</a><br>' + user_message + r'</blockquote></mx-reply>' + reply_message
content={"msgtype": "m.text", "body": body, "format": format, "formatted_body": formatted_body,