""" ollamarama-matrix: An AI chatbot for the Matrix chat protocol with infinite personalities. Author: Dustin Whyte Date: December 2023 """ from nio import AsyncClient, MatrixRoom, RoomMessageText import json import datetime import asyncio import requests import markdown class ollamarama: def __init__(self): #load config file self.config_file = "config.json" with open(self.config_file, "r") as f: config = json.load(f) f.close() self.server, self.username, self.password, self.channels, self.admins = config["matrix"].values() self.api_url = config["ollama"]["api_base"] + "/api/chat" self.default_personality = config["ollama"]["personality"] self.personality = self.default_personality self.client = AsyncClient(self.server, self.username) # time program started and joined channels self.join_time = datetime.datetime.now() # store chat history self.messages = {} #prompt parts self.prompt = config["ollama"]["prompt"] self.models = config["ollama"]["models"] self.default_model = self.models[config["ollama"]["default_model"]] self.model = self.default_model self.temperature, self.top_p, self.repeat_penalty = config["ollama"]["options"].values() self.defaults = { "temperature": self.temperature, "top_p": self.top_p, "repeat_penalty": self.repeat_penalty } #load help menu with open("help.txt", "r") as f: self.help, self.help_admin = f.read().split("~~~") f.close() # get the display name for a user async def display_name(self, user): try: name = await self.client.get_displayname(user) return name.displayname except Exception as e: print(e) # simplifies sending messages to the channel async def send_message(self, channel, message): await self.client.room_send( room_id=channel, message_type="m.room.message", content={ "msgtype": "m.text", "body": message, "format": "org.matrix.custom.html", "formatted_body": markdown.markdown(message, extensions=['fenced_code', 'nl2br'])}, ) # add messages to the history dictionary async def add_history(self, role, channel, sender, message): #check if channel is in the history yet if channel in self.messages: #check if user is in channel history if sender in self.messages[channel]: self.messages[channel][sender].append({"role": role, "content": message}) else: self.messages[channel][sender] = [ {"role": "system", "content": self.prompt[0] + self.personality + self.prompt[1]}, {"role": role, "content": message}] else: #set up channel in history self.messages[channel]= {} self.messages[channel][sender] = {} if role == "system": self.messages[channel][sender] = [{"role": role, "content": message}] else: #add personality to the new user entry self.messages[channel][sender] = [ {"role": "system", "content": self.prompt[0] + self.personality + self.prompt[1]}, {"role": role, "content": message}] #generate Ollama model response async def respond(self, channel, sender, message, sender2=None): try: data = { "model": self.model, "messages": message, "stream": False, "options": { "top_p": self.top_p, "temperature": self.temperature, "repeat_penalty": self.repeat_penalty } } response = requests.post(self.api_url, json=data, timeout=120) response.raise_for_status() data = response.json() except Exception as e: await self.send_message(channel, "Something went wrong") print(e) else: #Extract response text response_text = data["message"]["content"] #add to history await self.add_history("assistant", channel, sender, response_text) # .x function was used if sender2: display_name = await self.display_name(sender2) # .ai was used else: display_name = await self.display_name(sender) response_text = display_name + ":\n" + response_text.strip() #Send response to channel try: await self.send_message(channel, response_text) except Exception as e: print(e) #Shrink history list for token size management if len(self.messages[channel][sender]) > 24: if self.messages[channel][sender][0]["role"] == "system": del self.messages[channel][sender][1:3] else: del self.messages[channel][sender][0:2] # change the personality of the bot async def persona(self, channel, sender, persona): #clear existing history try: await self.messages[channel][sender].clear() except: pass personality = self.prompt[0] + persona + self.prompt[1] #set system prompt await self.add_history("system", channel, sender, personality) await self.add_history("user", channel, sender, "introduce yourself") # use a custom prompt async def custom(self, channel, sender, prompt): try: await self.messages[channel][sender].clear() except: pass await self.add_history("system", channel, sender, prompt) await self.add_history("user", channel, sender, "introduce yourself") # tracks the messages in channels async def message_callback(self, room: MatrixRoom, event: RoomMessageText): # Main bot functionality if isinstance(event, RoomMessageText): # convert timestamp message_time = event.server_timestamp / 1000 message_time = datetime.datetime.fromtimestamp(message_time) # assign parts of event to variables message = event.body sender = event.sender sender_display = await self.display_name(sender) room_id = room.room_id #check if the message was sent after joining and not by the bot if message_time > self.join_time and sender != self.username: user = await self.display_name(event.sender) #admin commands if message == ".admins": await self.send_message(room_id, f"Bot admins: {', '.join(self.admins)}") if sender_display in self.admins: #model switching if message.startswith(".model"): with open(self.config_file, "r") as f: config = json.load(f) f.close() self.models = config["ollama"]["models"] if message == ".models": current_model = f"Current model: {self.model}\nAvailable models: {', '.join(sorted(list(self.models)))}" await self.send_message(room_id, current_model) if message.startswith(".model "): m = message.split(" ", 1)[1] if m != None: if m in self.models: self.model = self.models[m] elif m == 'reset': self.model = self.default_model await self.send_message(room_id, f"Model set to {self.model}") #bot owner commands if sender_display == self.admins[0]: #add admins if message.startswith(".auth "): nick = message.split(" ", 1)[1].strip() if nick != None: self.admins.append(nick) await self.send_message(room_id, f"{nick} added to admins") #remove admins if message.startswith(".deauth "): nick = message.split(" ", 1)[1].strip() if nick != None: self.admins.remove(nick) await self.send_message(room_id, f"{nick} removed from admins") #set new global personality if message.startswith(".gpersona "): self.messages.clear() m = message.split(" ", 1)[1] if m != None: if m == 'reset': self.personality = self.default_personality else: self.personality = m.strip() await self.send_message(room_id, f"Global personality set to {self.personality}") #remove personality globally if message == ".gstock": pass #i'll figure this out later #reset history for all users if message == ".clear": self.messages.clear() self.model = self.default_model self.personality = self.default_personality self.temperature, self.top_p, self.repeat_penalty = self.defaults.values() await self.send_message(room_id, "Bot has been reset for everyone") if message.startswith((".temperature ", ".top_p ", ".repeat_penalty ")): attr_name = message.split()[0][1:] min_val, max_val, default_val = { "temperature": (0, 1, self.defaults["temperature"]), "top_p": (0, 1, self.defaults["top_p"]), "repeat_penalty": (0, 2, self.defaults["repeat_penalty"]) }[attr_name] if message.endswith(" reset"): setattr(self, attr_name, default_val) await self.send_message(room_id, f"{attr_name.capitalize()} set to {default_val}") else: try: value = float(message.split(" ", 1)[1]) if min_val <= value <= max_val: setattr(self, attr_name, value) await self.send_message(room_id, f"{attr_name.capitalize()} set to {value}") else: await self.send_message(room_id, f"Invalid input, {attr_name} is still {getattr(self, attr_name)}") except: await self.send_message(room_id, f"Invalid input, {attr_name} is still {getattr(self, attr_name)}") # main AI response functionality if message.startswith(".ai ") or message.startswith(self.bot_id): if message != ".ai reset": m = message.split(" ", 1) try: m = m[1] await self.add_history("user", room_id, sender, m) await self.respond(room_id, sender, self.messages[room_id][sender]) except: pass # collaborative functionality if message.startswith(".x "): m = message.split(" ", 2) m.pop(0) if len(m) > 1: disp_name = m[0] name_id = "" m = m[1] if room_id in self.messages: for user in self.messages[room_id]: try: username = await self.display_name(user) if disp_name == username: name_id = user except: name_id = disp_name await self.add_history("user", room_id, name_id, m) await self.respond(room_id, name_id, self.messages[room_id][name_id], sender) #change personality if message.startswith(".persona "): m = message.split(" ", 1) m = m[1] await self.persona(room_id, sender, m) await self.respond(room_id, sender, self.messages[room_id][sender]) #custom prompt use if message.startswith(".custom "): m = message.split(" ", 1) m = m[1] await self.custom(room_id, sender, m) await self.respond(room_id, sender, self.messages[room_id][sender]) # reset bot to default personality if message.startswith(".reset") or message == ".ai reset": #some users keep forgetting the correct command if room_id in self.messages: if sender in self.messages[room_id]: self.messages[room_id][sender].clear() await self.persona(room_id, sender, self.personality) del self.messages[room_id][sender][1] try: await self.send_message(room_id, f"{self.bot_id} reset to default for {sender_display}") except: await self.send_message(room_id, f"{self.bot_id} reset to default for {sender}") # Stock settings, no personality if message.startswith(".stock"): if room_id in self.messages: if sender in self.messages[room_id]: self.messages[room_id][sender].clear() else: self.messages[room_id] = {} self.messages[room_id][sender] = [] try: await self.send_message(room_id, f"Stock settings applied for {sender_display}") except: await self.send_message(room_id, f"Stock settings applied for {sender}") # help menu if message.startswith(".help"): await self.send_message(room_id, self.help) if sender_display in self.admins: await self.send_message(room_id, self.help_admin) # main loop async def main(self): # Login, print "Logged in as @alice:example.org device id: RANDOMDID" print(await self.client.login(self.password)) # get account display name self.bot_id = await self.display_name(self.username) # join channels for channel in self.channels: try: await self.client.join(channel) print(f"{self.bot_id} joined {channel}") except: print(f"Couldn't join {channel}") # start listening for messages self.client.add_event_callback(self.message_callback, RoomMessageText) await self.client.sync_forever(timeout=30000) if __name__ == "__main__": bot = ollamarama() asyncio.get_event_loop().run_until_complete(bot.main())