diff --git a/assets/logo.png b/assets/logo.png new file mode 100644 index 0000000..d47437c Binary files /dev/null and b/assets/logo.png differ diff --git a/classes/bot.py b/classes/bot.py index b88a9a2..c728633 100644 --- a/classes/bot.py +++ b/classes/bot.py @@ -36,6 +36,7 @@ from typing import Optional, List, Dict, Tuple from configparser import ConfigParser from datetime import datetime from io import BytesIO +from pathlib import Path import uuid import traceback @@ -53,7 +54,7 @@ from .trackingmore import TrackingMore class GPTBot: # Default values database: Optional[duckdb.DuckDBPyConnection] = None - default_room_name: str = "GPTBot" # Default name of rooms created by the bot + display_name = default_room_name = "GPTBot" # Default name of rooms created by the bot default_system_message: str = "You are a helpful assistant." # Force default system message to be included even if a custom room message is set force_system_message: bool = False @@ -69,6 +70,8 @@ class GPTBot: operator: Optional[str] = None room_ignore_list: List[str] = [] # List of rooms to ignore invites from debug: bool = False + logo: Optional[Image.Image] = None + logo_uri: Optional[str] = None @classmethod def from_config(cls, config: ConfigParser): @@ -99,6 +102,15 @@ class GPTBot: "ForceSystemMessage", bot.force_system_message) bot.debug = config["GPTBot"].getboolean("Debug", bot.debug) + logo_path = config["GPTBot"].get("Logo", str(Path(__file__).parent.parent / "assets/logo.png")) + + bot.logger.log(f"Loading logo from {logo_path}") + + if Path(logo_path).exists() and Path(logo_path).is_file(): + bot.logo = Image.open(logo_path) + + bot.display_name = config["GPTBot"].get("DisplayName", bot.display_name) + bot.chat_api = bot.image_api = bot.classification_api = OpenAI( config["OpenAI"]["APIKey"], config["OpenAI"].get("Model"), bot.logger) bot.max_tokens = config["OpenAI"].getint("MaxTokens", bot.max_tokens) @@ -340,6 +352,30 @@ class GPTBot: f"Error leaving room {invite}: {leave_response.message}", "error") self.room_ignore_list.append(invite) + async def upload_file(self, file: bytes, filename: str = "file", mime: str = "application/octet-stream") -> str: + """Upload a file to the homeserver. + + Args: + file (bytes): The file to upload. + filename (str, optional): The name of the file. Defaults to "file". + mime (str, optional): The MIME type of the file. Defaults to "application/octet-stream". + + Returns: + str: The MXC URI of the uploaded file. + """ + + bio = BytesIO(file) + bio.seek(0) + + response, _ = await self.matrix_client.upload( + bio, + content_type=mime, + filename=filename, + filesize=len(file) + ) + + return response.content_uri + async def send_image(self, room: MatrixRoom, image: bytes, message: Optional[str] = None): """Send an image to a room. @@ -361,14 +397,7 @@ class GPTBot: self.logger.log( f"Uploading - Image size: {width}x{height} pixels, MIME type: {mime}") - bio.seek(0) - - response, _ = await self.matrix_client.upload( - bio, - content_type=mime, - filename="image", - filesize=len(image) - ) + content_uri = await self.upload_file(image, "image", mime) self.logger.log("Uploaded image - sending message...") @@ -381,7 +410,7 @@ class GPTBot: "h": height, }, "msgtype": "m.image", - "url": response.content_uri + "url": content_uri } status = await self.matrix_client.room_send( @@ -547,6 +576,26 @@ class GPTBot: self.matrix_client.add_response_callback( self.response_callback, Response) + # Set custom name / logo + + if self.display_name: + self.logger.log(f"Setting display name to {self.display_name}") + await self.matrix_client.set_displayname(self.display_name) + if self.logo: + self.logger.log("Setting avatar...") + logo_bio = BytesIO() + self.logo.save(logo_bio, format=self.logo.format) + uri = await self.upload_file(logo_bio.getvalue(), "logo", Image.MIME[self.logo.format]) + self.logo_uri = uri + + asyncio.create_task(self.matrix_client.set_avatar(uri)) + + for room in self.matrix_client.rooms.keys(): + self.logger.log(f"Setting avatar for {room}...", "debug") + asyncio.create_task(self.matrix_client.room_put_state(room, "m.room.avatar", { + "url": uri + }, "")) + # Start syncing events self.logger.log("Starting sync loop...") try: diff --git a/commands/newroom.py b/commands/newroom.py index 9a0d19d..f7e1514 100644 --- a/commands/newroom.py +++ b/commands/newroom.py @@ -32,9 +32,14 @@ async def command_newroom(room: MatrixRoom, event: RoomMessageText, bot): bot.logger.log(f"Adding new room to space {space[0]}...") await bot.add_rooms_to_space(space[0], [new_room.room_id]) + if bot.logo_uri: + await bot.matrix_client.room_put_state(room, "m.room.avatar", { + "url": bot.logo_uri + }, "") + await bot.matrix_client.room_put_state( - new_room.room_id, "m.room.power_levels", {"users": {event.sender: 100}}) + new_room.room_id, "m.room.power_levels", {"users": {event.sender: 100, bot.matrix_client.user_id: 100}}) await bot.matrix_client.joined_rooms() await bot.send_message(room, f"Alright, I've created a new room called '{room_name}' and invited you to it. You can find it at {new_room.room_id}", True) - await bot.send_message(bot.rooms[new_room.room_id], f"Welcome to the new room! What can I do for you?") \ No newline at end of file + await bot.send_message(bot.rooms[new_room.room_id], f"Welcome to the new room! What can I do for you?") diff --git a/config.dist.ini b/config.dist.ini index e3d412a..e2f92e6 100644 --- a/config.dist.ini +++ b/config.dist.ini @@ -88,6 +88,16 @@ Operator = Contact details not set # # ForceSystemMessage = 0 +# Path to a custom logo +# Used as room/space image and profile picture +# Defaults to logo.png in assets directory +# +# Logo = assets/logo.png + +# Display name for the bot +# +# DisplayName = GPTBot + [Database] # Settings for the DuckDB database. diff --git a/migrations/__init__.py b/migrations/__init__.py index 1c0fa51..1cf8d53 100644 --- a/migrations/__init__.py +++ b/migrations/__init__.py @@ -4,7 +4,7 @@ from importlib import import_module from duckdb import DuckDBPyConnection -MAX_MIGRATION = 7 +MAX_MIGRATION = 8 MIGRATIONS = OrderedDict() diff --git a/migrations/migration_8.py b/migrations/migration_8.py new file mode 100644 index 0000000..86308fe --- /dev/null +++ b/migrations/migration_8.py @@ -0,0 +1,22 @@ +# Migration to add settings table + +from datetime import datetime + +def migration(conn): + with conn.cursor() as cursor: + cursor.execute( + """ + CREATE TABLE IF NOT EXISTS settings ( + setting TEXT NOT NULL, + value TEXT NOT NULL, + PRIMARY KEY (setting) + ) + """ + ) + + cursor.execute( + "INSERT INTO migrations (id, timestamp) VALUES (8, ?)", + (datetime.now(),) + ) + + conn.commit() \ No newline at end of file