matrix_chatgpt_bot/bot.py

561 lines
24 KiB
Python
Raw Normal View History

2023-03-05 14:07:25 +00:00
import sys
import asyncio
2023-03-05 14:07:25 +00:00
import re
import os
2023-04-10 02:52:18 +00:00
import traceback
from typing import Optional, Union
from nio import (
AsyncClient,
MatrixRoom,
RoomMessageText,
InviteMemberEvent,
2023-04-10 13:00:22 +00:00
MegolmEvent,
2023-04-10 02:52:18 +00:00
LoginResponse,
JoinError,
ToDeviceError,
LocalProtocolError,
KeyVerificationEvent,
KeyVerificationStart,
KeyVerificationCancel,
KeyVerificationKey,
KeyVerificationMac,
AsyncClientConfig
)
2023-03-05 14:07:25 +00:00
from nio.store.database import SqliteStore
2023-04-10 02:52:18 +00:00
from askgpt import askGPT
2023-03-05 14:07:25 +00:00
from send_message import send_room_message
from v3 import Chatbot
2023-03-10 13:43:18 +00:00
from log import getlogger
2023-03-10 15:45:38 +00:00
from bing import BingBot
2023-03-22 14:28:22 +00:00
from BingImageGen import ImageGen
from send_image import send_room_image
2023-04-10 02:52:18 +00:00
2023-03-10 13:43:18 +00:00
logger = getlogger()
2023-03-05 14:07:25 +00:00
class Bot:
def __init__(
2023-03-10 15:45:38 +00:00
self,
homeserver: str,
user_id: str,
device_id: str,
2023-04-10 02:52:18 +00:00
chatgpt_api_endpoint: str = os.environ.get("CHATGPT_API_ENDPOINT") or "https://api.openai.com/v1/chat/completions",
api_key: Optional[str] = os.environ.get("OPENAI_API_KEY") or "",
2023-04-10 13:00:22 +00:00
room_id: Union[str, None] = None,
2023-03-10 15:45:38 +00:00
bing_api_endpoint: Optional[str] = '',
2023-04-10 02:52:18 +00:00
password: Union[str, None] = None,
access_token: Union[str, None] = None,
jailbreakEnabled: Optional[bool] = True,
2023-03-22 14:28:22 +00:00
bing_auth_cookie: Optional[str] = '',
2023-03-05 14:07:25 +00:00
):
2023-04-10 13:00:22 +00:00
if (homeserver is None or user_id is None \
or device_id is None):
logger.warning("homeserver && user_id && device_id is required")
sys.exit(1)
if (password is None and access_token is None):
logger.warning("password and access_toekn is required")
sys.exit(1)
2023-03-05 14:07:25 +00:00
self.homeserver = homeserver
self.user_id = user_id
self.password = password
2023-04-10 02:52:18 +00:00
self.access_token = access_token
2023-03-05 14:07:25 +00:00
self.device_id = device_id
self.room_id = room_id
self.api_key = api_key
2023-04-10 02:52:18 +00:00
self.chatgpt_api_endpoint = chatgpt_api_endpoint
2023-03-10 15:45:38 +00:00
self.bing_api_endpoint = bing_api_endpoint
self.jailbreakEnabled = jailbreakEnabled
2023-03-22 14:28:22 +00:00
self.bing_auth_cookie = bing_auth_cookie
2023-03-05 14:07:25 +00:00
# initialize AsyncClient object
self.store_path = os.getcwd()
self.config = AsyncClientConfig(store=SqliteStore,
2023-04-10 02:52:18 +00:00
store_name="db",
2023-03-05 14:07:25 +00:00
store_sync_tokens=True,
2023-03-14 14:37:30 +00:00
encryption_enabled=True,
2023-03-05 14:07:25 +00:00
)
2023-04-10 02:52:18 +00:00
self.client = AsyncClient(homeserver=self.homeserver, user=self.user_id, device_id=self.device_id,
2023-03-14 14:37:30 +00:00
config=self.config, store_path=self.store_path,)
2023-04-10 02:52:18 +00:00
if self.access_token is not None:
self.client.access_token = self.access_token
2023-04-10 13:00:22 +00:00
# setup event callbacks
self.client.add_event_callback(self.message_callback, (RoomMessageText, ))
self.client.add_event_callback(self.decryption_failure, (MegolmEvent, ))
self.client.add_event_callback(self.invite_callback, (InviteMemberEvent, ))
self.client.add_to_device_callback(self.to_device_callback, (KeyVerificationEvent, ))
2023-03-05 14:07:25 +00:00
# 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*(.+)$")
2023-03-10 15:45:38 +00:00
self.bing_prog = re.compile(r"^\s*!bing\s*(.+)$")
2023-03-22 14:28:22 +00:00
self.pic_prog = re.compile(r"^\s*!pic\s*(.+)$")
self.help_prog = re.compile(r"^\s*!help\s*.*$")
2023-03-10 15:45:38 +00:00
# initialize chatbot and chatgpt_api_endpoint
if self.api_key != '':
2023-04-10 11:37:43 +00:00
self.chatbot = Chatbot(api_key=self.api_key, timeout=60)
2023-03-05 14:07:25 +00:00
2023-04-10 02:52:18 +00:00
self.chatgpt_api_endpoint = self.chatgpt_api_endpoint
# request header for !gpt command
self.headers = {
"Content-Type": "application/json",
2023-04-10 02:52:18 +00:00
"Authorization": f"Bearer {self.api_key}",
}
else:
2023-04-10 02:52:18 +00:00
self.chatgpt_api_endpoint = self.chatgpt_api_endpoint
self.headers = {
"Content-Type": "application/json",
}
2023-04-10 02:52:18 +00:00
# initialize askGPT class
self.askgpt = askGPT()
2023-03-10 15:45:38 +00:00
# initialize bingbot
if self.bing_api_endpoint != '':
self.bingbot = BingBot(bing_api_endpoint, jailbreakEnabled=self.jailbreakEnabled)
2023-03-10 15:45:38 +00:00
2023-03-22 14:28:22 +00:00
# initialize BingImageGen
if self.bing_auth_cookie != '':
self.imageGen = ImageGen(self.bing_auth_cookie)
2023-04-10 13:00:22 +00:00
# message_callback RoomMessageText event
2023-03-05 14:07:25 +00:00
async def message_callback(self, room: MatrixRoom, event: RoomMessageText) -> None:
2023-04-10 13:00:22 +00:00
if self.room_id is None:
2023-03-14 14:37:30 +00:00
room_id = room.room_id
else:
# if event room id does not match the room id in config, return
if room.room_id != self.room_id:
return
room_id = self.room_id
# reply event_id
reply_to_event_id = event.event_id
2023-04-10 02:52:18 +00:00
# sender_id
sender_id = event.sender
# user_message
raw_user_message = event.body
2023-03-10 13:43:18 +00:00
# print info to console
print(
f"Message received in room {room.display_name}\n"
2023-04-10 02:52:18 +00:00
f"{room.user_name(event.sender)} | {raw_user_message}"
2023-03-10 13:43:18 +00:00
)
2023-04-10 02:52:18 +00:00
# prevent command trigger loop
2023-03-22 14:28:22 +00:00
if self.user_id != event.sender:
# remove newline character from event.body
2023-04-10 02:52:18 +00:00
content_body = re.sub("\r\n|\r|\n", " ", raw_user_message)
2023-03-10 15:45:38 +00:00
2023-03-22 14:28:22 +00:00
# chatgpt
2023-04-10 02:52:18 +00:00
n = self.chat_prog.match(content_body)
2023-03-22 14:28:22 +00:00
if n:
prompt = n.group(1)
2023-03-22 14:28:22 +00:00
if self.api_key != '':
2023-04-10 13:00:22 +00:00
try:
await self.chat(room_id, reply_to_event_id, prompt, sender_id, raw_user_message)
except Exception as e:
logger.error(e)
await send_room_message(self.client, room_id, reply_message=str(e))
2023-03-22 14:28:22 +00:00
else:
logger.warning("No API_KEY provided")
2023-04-10 13:00:22 +00:00
await send_room_message(self.client, room_id, reply_message="API_KEY not provided")
2023-03-22 14:28:22 +00:00
2023-04-10 02:52:18 +00:00
m = self.gpt_prog.match(content_body)
2023-03-22 14:28:22 +00:00
if m:
prompt = m.group(1)
2023-04-10 11:37:43 +00:00
try:
await self.gpt(room_id, reply_to_event_id, prompt, sender_id, raw_user_message)
except Exception as e:
logger.error(e)
2023-04-10 13:00:22 +00:00
await send_room_message(self.client, room_id, reply_message=str(e))
2023-03-22 14:28:22 +00:00
# bing ai
if self.bing_api_endpoint != '':
2023-04-10 02:52:18 +00:00
b = self.bing_prog.match(content_body)
2023-03-22 14:28:22 +00:00
if b:
prompt = b.group(1)
2023-04-10 02:52:18 +00:00
# raw_content_body used for construct formatted_body
2023-04-10 13:00:22 +00:00
try:
await self.bing(room_id, reply_to_event_id, prompt, sender_id, raw_user_message)
except Exception as e:
await send_room_message(self.client, room_id, reply_message=str(e))
2023-03-22 14:28:22 +00:00
# Image Generation by Microsoft Bing
if self.bing_auth_cookie != '':
2023-04-10 02:52:18 +00:00
i = self.pic_prog.match(content_body)
2023-03-22 14:28:22 +00:00
if i:
prompt = i.group(1)
2023-04-10 13:00:22 +00:00
try:
await self.pic(room_id, prompt)
except Exception as e:
await send_room_message(self.client, room_id, reply_message=str(e))
2023-03-22 14:28:22 +00:00
# help command
2023-04-10 02:52:18 +00:00
h = self.help_prog.match(content_body)
2023-03-22 14:28:22 +00:00
if h:
await self.help(room_id)
2023-04-10 13:00:22 +00:00
# message_callback decryption_failure event
async def decryption_failure(self, room: MatrixRoom, event: MegolmEvent) -> None:
if not isinstance(event, MegolmEvent):
return
logger.error(
f"Failed to decrypt message: {event.event_id} from {event.sender} in {room.room_id}\n" + \
"Please make sure the bot current session is verified"
)
2023-04-10 02:52:18 +00:00
# invite_callback event
async def invite_callback(self, room: MatrixRoom, event: InviteMemberEvent) -> None:
"""Handle an incoming invite event.
https://github.com/8go/matrix-eno-bot/blob/ad037e02bd2960941109e9526c1033dd157bb212/callbacks.py#L104
If an invite is received, then join the room specified in the invite.
code copied from:
"""
logger.debug(f"Got invite to {room.room_id} from {event.sender}.")
# Attempt to join 3 times before giving up
for attempt in range(3):
result = await self.client.join(room.room_id)
if type(result) == JoinError:
logger.error(
f"Error joining room {room.room_id} (attempt %d): %s",
attempt, result.message,
)
else:
break
else:
logger.error("Unable to join room: %s", room.room_id)
# Successfully joined room
logger.info(f"Joined {room.room_id}")
# to_device_callback event
async def to_device_callback(self, event: KeyVerificationEvent) -> None:
"""Handle events sent to device.
Specifically this will perform Emoji verification.
It will accept an incoming Emoji verification requests
and follow the verification protocol.
code copied from: https://github.com/8go/matrix-eno-bot/blob/ad037e02bd2960941109e9526c1033dd157bb212/callbacks.py#L127
"""
2023-03-22 14:28:22 +00:00
try:
2023-04-10 02:52:18 +00:00
client = self.client
logger.debug(
f"Device Event of type {type(event)} received in "
"to_device_cb().")
if isinstance(event, KeyVerificationStart): # first step
""" first step: receive KeyVerificationStart
KeyVerificationStart(
source={'content':
{'method': 'm.sas.v1',
'from_device': 'DEVICEIDXY',
'key_agreement_protocols':
['curve25519-hkdf-sha256', 'curve25519'],
'hashes': ['sha256'],
'message_authentication_codes':
['hkdf-hmac-sha256', 'hmac-sha256'],
'short_authentication_string':
['decimal', 'emoji'],
'transaction_id': 'SomeTxId'
},
'type': 'm.key.verification.start',
'sender': '@user2:example.org'
},
sender='@user2:example.org',
transaction_id='SomeTxId',
from_device='DEVICEIDXY',
method='m.sas.v1',
key_agreement_protocols=[
'curve25519-hkdf-sha256', 'curve25519'],
hashes=['sha256'],
message_authentication_codes=[
'hkdf-hmac-sha256', 'hmac-sha256'],
short_authentication_string=['decimal', 'emoji'])
"""
if "emoji" not in event.short_authentication_string:
estr = ("Other device does not support emoji verification "
f"{event.short_authentication_string}. Aborting.")
print(estr)
logger.info(estr)
return
resp = await client.accept_key_verification(
event.transaction_id)
if isinstance(resp, ToDeviceError):
estr = f"accept_key_verification() failed with {resp}"
print(estr)
logger.info(estr)
sas = client.key_verifications[event.transaction_id]
todevice_msg = sas.share_key()
resp = await client.to_device(todevice_msg)
if isinstance(resp, ToDeviceError):
estr = f"to_device() failed with {resp}"
print(estr)
logger.info(estr)
elif isinstance(event, KeyVerificationCancel): # anytime
""" at any time: receive KeyVerificationCancel
KeyVerificationCancel(source={
'content': {'code': 'm.mismatched_sas',
'reason': 'Mismatched authentication string',
'transaction_id': 'SomeTxId'},
'type': 'm.key.verification.cancel',
'sender': '@user2:example.org'},
sender='@user2:example.org',
transaction_id='SomeTxId',
code='m.mismatched_sas',
reason='Mismatched short authentication string')
"""
# There is no need to issue a
# client.cancel_key_verification(tx_id, reject=False)
# here. The SAS flow is already cancelled.
# We only need to inform the user.
estr = (f"Verification has been cancelled by {event.sender} "
f"for reason \"{event.reason}\".")
print(estr)
logger.info(estr)
elif isinstance(event, KeyVerificationKey): # second step
""" Second step is to receive KeyVerificationKey
KeyVerificationKey(
source={'content': {
'key': 'SomeCryptoKey',
'transaction_id': 'SomeTxId'},
'type': 'm.key.verification.key',
'sender': '@user2:example.org'
},
sender='@user2:example.org',
transaction_id='SomeTxId',
key='SomeCryptoKey')
"""
sas = client.key_verifications[event.transaction_id]
print(f"{sas.get_emoji()}")
# don't log the emojis
# The bot process must run in forground with a screen and
# keyboard so that user can accept/reject via keyboard.
# For emoji verification bot must not run as service or
# in background.
yn = input("Do the emojis match? (Y/N) (C for Cancel) ")
if yn.lower() == "y":
estr = ("Match! The verification for this "
"device will be accepted.")
print(estr)
logger.info(estr)
resp = await client.confirm_short_auth_string(
event.transaction_id)
if isinstance(resp, ToDeviceError):
estr = ("confirm_short_auth_string() "
f"failed with {resp}")
print(estr)
logger.info(estr)
elif yn.lower() == "n": # no, don't match, reject
estr = ("No match! Device will NOT be verified "
"by rejecting verification.")
print(estr)
logger.info(estr)
resp = await client.cancel_key_verification(
event.transaction_id, reject=True)
if isinstance(resp, ToDeviceError):
estr = (f"cancel_key_verification failed with {resp}")
print(estr)
logger.info(estr)
else: # C or anything for cancel
estr = ("Cancelled by user! Verification will be "
"cancelled.")
print(estr)
logger.info(estr)
resp = await client.cancel_key_verification(
event.transaction_id, reject=False)
if isinstance(resp, ToDeviceError):
estr = (f"cancel_key_verification failed with {resp}")
print(estr)
logger.info(estr)
elif isinstance(event, KeyVerificationMac): # third step
""" Third step is to receive KeyVerificationMac
KeyVerificationMac(
source={'content': {
'mac': {'ed25519:DEVICEIDXY': 'SomeKey1',
'ed25519:SomeKey2': 'SomeKey3'},
'keys': 'SomeCryptoKey4',
'transaction_id': 'SomeTxId'},
'type': 'm.key.verification.mac',
'sender': '@user2:example.org'},
sender='@user2:example.org',
transaction_id='SomeTxId',
mac={'ed25519:DEVICEIDXY': 'SomeKey1',
'ed25519:SomeKey2': 'SomeKey3'},
keys='SomeCryptoKey4')
"""
sas = client.key_verifications[event.transaction_id]
try:
todevice_msg = sas.get_mac()
except LocalProtocolError as e:
# e.g. it might have been cancelled by ourselves
estr = (f"Cancelled or protocol error: Reason: {e}.\n"
f"Verification with {event.sender} not concluded. "
"Try again?")
print(estr)
logger.info(estr)
else:
resp = await client.to_device(todevice_msg)
if isinstance(resp, ToDeviceError):
estr = f"to_device failed with {resp}"
print(estr)
logger.info(estr)
estr = (f"sas.we_started_it = {sas.we_started_it}\n"
f"sas.sas_accepted = {sas.sas_accepted}\n"
f"sas.canceled = {sas.canceled}\n"
f"sas.timed_out = {sas.timed_out}\n"
f"sas.verified = {sas.verified}\n"
f"sas.verified_devices = {sas.verified_devices}\n")
print(estr)
logger.info(estr)
estr = ("Emoji verification was successful!\n"
"Initiate another Emoji verification from "
"another device or room if desired. "
"Or if done verifying, hit Control-C to stop the "
"bot in order to restart it as a service or to "
"run it in the background.")
print(estr)
logger.info(estr)
else:
estr = (f"Received unexpected event type {type(event)}. "
f"Event is {event}. Event will be ignored.")
print(estr)
logger.info(estr)
except BaseException:
estr = traceback.format_exc()
print(estr)
logger.info(estr)
# !chat command
async def chat(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message):
2023-04-10 11:37:43 +00:00
await self.client.room_typing(room_id, timeout=180000)
2023-04-10 02:52:18 +00:00
try:
2023-04-10 11:37:43 +00:00
text = await asyncio.wait_for(self.chatbot.ask_async(prompt), timeout=180)
2023-04-10 02:52:18 +00:00
except TimeoutError as e:
logger.error("timeoutException", exc_info=True)
text = "Timeout error"
2023-03-22 14:28:22 +00:00
except Exception as e:
logger.error("Error", exc_info=True)
print(f"Error: {e}")
2023-04-10 02:52:18 +00:00
text = text.strip()
try:
await send_room_message(self.client, room_id, reply_message=text,
reply_to_event_id=reply_to_event_id, sender_id=sender_id, user_message=raw_user_message)
except Exception as e:
logger.error(f"Error: {e}", exc_info=True)
# !gpt command
async def gpt(self, room_id, reply_to_event_id, prompt, sender_id, raw_user_message):
2023-03-22 14:28:22 +00:00
try:
2023-03-10 11:19:49 +00:00
# sending typing state
2023-04-10 11:37:43 +00:00
await self.client.room_typing(room_id, timeout=180000)
2023-03-22 14:28:22 +00:00
# timeout 120s
2023-04-10 11:37:43 +00:00
text = await asyncio.wait_for(self.askgpt.oneTimeAsk(prompt, self.chatgpt_api_endpoint, self.headers), timeout=180)
2023-03-22 14:28:22 +00:00
except TimeoutError:
logger.error("timeoutException", exc_info=True)
text = "Timeout error"
2023-03-10 11:19:49 +00:00
2023-03-22 14:28:22 +00:00
text = text.strip()
try:
2023-04-10 02:52:18 +00:00
await send_room_message(self.client, room_id, reply_message=text,
reply_to_event_id=reply_to_event_id, sender_id=sender_id, user_message=raw_user_message)
2023-03-22 14:28:22 +00:00
except Exception as e:
logger.error(f"Error: {e}", exc_info=True)
# !bing command
2023-04-10 02:52:18 +00:00
async def bing(self, room_id, reply_to_event_id, prompt, sender_id, raw_content_body):
2023-03-22 14:28:22 +00:00
try:
# sending typing state
2023-04-10 11:37:43 +00:00
await self.client.room_typing(room_id, timeout=180000)
2023-03-22 14:28:22 +00:00
# timeout 120s
2023-04-10 11:37:43 +00:00
text = await asyncio.wait_for(self.bingbot.ask_bing(prompt), timeout=180)
2023-03-22 14:28:22 +00:00
except TimeoutError:
logger.error("timeoutException", exc_info=True)
text = "Timeout error"
text = text.strip()
try:
2023-04-10 02:52:18 +00:00
await send_room_message(self.client, room_id, reply_message=text,
reply_to_event_id=reply_to_event_id, sender=sender_id, raw_content_body=raw_content_body)
2023-03-22 14:28:22 +00:00
except Exception as e:
logger.error(f"Error: {e}", exc_info=True)
2023-03-10 11:19:49 +00:00
2023-03-22 14:28:22 +00:00
# !pic command
async def pic(self, room_id, prompt):
try:
2023-04-10 13:00:22 +00:00
await self.client.room_typing(room_id, timeout=180000)
2023-03-22 14:28:22 +00:00
# generate image
generated_image_path = self.imageGen.save_images(
self.imageGen.get_images(prompt),
"images",
)
# send image
if generated_image_path != "":
await send_room_image(self.client, room_id, generated_image_path)
2023-04-10 13:00:22 +00:00
await self.client.room_typing(room_id, bool=False)
2023-03-22 14:28:22 +00:00
except Exception as e:
logger.error(f"Error: {e}", exc_info=True)
# !help command
async def help(self, room_id):
try:
# sending typing state
await self.client.room_typing(room_id)
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" + \
2023-04-10 11:37:43 +00:00
"!pic [prompt], Image generation by Microsoft Bing\n" + \
"!help, help message"
2023-03-22 14:28:22 +00:00
2023-04-10 02:52:18 +00:00
await send_room_message(self.client, room_id, reply_message=help_info)
2023-03-22 14:28:22 +00:00
except Exception as e:
logger.error(f"Error: {e}", exc_info=True)
2023-03-10 15:45:38 +00:00
2023-03-05 14:07:25 +00:00
# bot login
async def login(self) -> None:
2023-03-12 15:24:05 +00:00
try:
resp = await self.client.login(password=self.password)
if not isinstance(resp, LoginResponse):
logger.error("Login Failed")
print(f"Login Failed: {resp}")
sys.exit(1)
except Exception as e:
2023-03-22 14:28:22 +00:00
logger.error(f"Error: {e}", exc_info=True)
2023-03-05 14:07:25 +00:00
# sync messages in the room
2023-04-10 11:37:43 +00:00
async def sync_forever(self, timeout=30000, full_state=True) -> None:
2023-04-10 13:00:22 +00:00
2023-04-10 02:52:18 +00:00
await self.client.sync_forever(timeout=timeout, full_state=full_state)
# Sync encryption keys with the server
2023-04-10 11:37:43 +00:00
async def sync_encryption_key(self) -> None:
2023-04-10 02:52:18 +00:00
if self.client.should_upload_keys:
await self.client.keys_upload()
# Trust own devices
2023-04-10 11:37:43 +00:00
async def trust_own_devices(self) -> None:
2023-04-10 02:52:18 +00:00
await self.client.sync(timeout=30000, full_state=True)
for device_id, olm_device in self.client.device_store[
self.user_id].items():
logger.debug("My other devices are: "
f"device_id={device_id}, "
f"olm_device={olm_device}.")
logger.info("Setting up trust for my own "
f"device {device_id} and session key "
f"{olm_device.keys['ed25519']}.")
self.client.verify_device(olm_device)