Browse Source

Add an AI generated logo
Set logo as room/space avatar
Stay admin when creating a room
Add a settings table to the database

Kumi 1 year ago
parent
commit
cfeaae3fac
6 changed files with 99 additions and 13 deletions
  1. BIN
      assets/logo.png
  2. 59 10
      classes/bot.py
  3. 7 2
      commands/newroom.py
  4. 10 0
      config.dist.ini
  5. 1 1
      migrations/__init__.py
  6. 22 0
      migrations/migration_8.py

BIN
assets/logo.png


+ 59 - 10
classes/bot.py

@@ -36,6 +36,7 @@ from typing import Optional, List, Dict, Tuple
 from configparser import ConfigParser
 from configparser import ConfigParser
 from datetime import datetime
 from datetime import datetime
 from io import BytesIO
 from io import BytesIO
+from pathlib import Path
 
 
 import uuid
 import uuid
 import traceback
 import traceback
@@ -53,7 +54,7 @@ from .trackingmore import TrackingMore
 class GPTBot:
 class GPTBot:
     # Default values
     # Default values
     database: Optional[duckdb.DuckDBPyConnection] = None
     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."
     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 default system message to be included even if a custom room message is set
     force_system_message: bool = False
     force_system_message: bool = False
@@ -69,6 +70,8 @@ class GPTBot:
     operator: Optional[str] = None
     operator: Optional[str] = None
     room_ignore_list: List[str] = []  # List of rooms to ignore invites from
     room_ignore_list: List[str] = []  # List of rooms to ignore invites from
     debug: bool = False
     debug: bool = False
+    logo: Optional[Image.Image] = None
+    logo_uri: Optional[str] = None
 
 
     @classmethod
     @classmethod
     def from_config(cls, config: ConfigParser):
     def from_config(cls, config: ConfigParser):
@@ -99,6 +102,15 @@ class GPTBot:
                 "ForceSystemMessage", bot.force_system_message)
                 "ForceSystemMessage", bot.force_system_message)
             bot.debug = config["GPTBot"].getboolean("Debug", bot.debug)
             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(
         bot.chat_api = bot.image_api = bot.classification_api = OpenAI(
             config["OpenAI"]["APIKey"], config["OpenAI"].get("Model"), bot.logger)
             config["OpenAI"]["APIKey"], config["OpenAI"].get("Model"), bot.logger)
         bot.max_tokens = config["OpenAI"].getint("MaxTokens", bot.max_tokens)
         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")
                         f"Error leaving room {invite}: {leave_response.message}", "error")
                     self.room_ignore_list.append(invite)
                     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):
     async def send_image(self, room: MatrixRoom, image: bytes, message: Optional[str] = None):
         """Send an image to a room.
         """Send an image to a room.
 
 
@@ -361,14 +397,7 @@ class GPTBot:
         self.logger.log(
         self.logger.log(
             f"Uploading - Image size: {width}x{height} pixels, MIME type: {mime}")
             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...")
         self.logger.log("Uploaded image - sending message...")
 
 
@@ -381,7 +410,7 @@ class GPTBot:
                 "h": height,
                 "h": height,
             },
             },
             "msgtype": "m.image",
             "msgtype": "m.image",
-            "url": response.content_uri
+            "url": content_uri
         }
         }
 
 
         status = await self.matrix_client.room_send(
         status = await self.matrix_client.room_send(
@@ -547,6 +576,26 @@ class GPTBot:
         self.matrix_client.add_response_callback(
         self.matrix_client.add_response_callback(
             self.response_callback, Response)
             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
         # Start syncing events
         self.logger.log("Starting sync loop...")
         self.logger.log("Starting sync loop...")
         try:
         try:

+ 7 - 2
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]}...")
         bot.logger.log(f"Adding new room to space {space[0]}...")
         await bot.add_rooms_to_space(space[0], [new_room.room_id])
         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(
     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.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(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?")
+    await bot.send_message(bot.rooms[new_room.room_id], f"Welcome to the new room! What can I do for you?")

+ 10 - 0
config.dist.ini

@@ -88,6 +88,16 @@ Operator = Contact details not set
 #
 #
 # ForceSystemMessage = 0
 # 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]
 [Database]
 
 
 # Settings for the DuckDB database.
 # Settings for the DuckDB database.

+ 1 - 1
migrations/__init__.py

@@ -4,7 +4,7 @@ from importlib import import_module
 
 
 from duckdb import DuckDBPyConnection
 from duckdb import DuckDBPyConnection
 
 
-MAX_MIGRATION = 7
+MAX_MIGRATION = 8
 
 
 MIGRATIONS = OrderedDict()
 MIGRATIONS = OrderedDict()
 
 

+ 22 - 0
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()