From 3d32343e54e433e4c2a0c8e82eb89d1108d5967e Mon Sep 17 00:00:00 2001 From: Kumi Date: Tue, 9 May 2023 10:27:03 +0000 Subject: [PATCH] Add spaces feature --- callbacks/join.py | 9 +++ classes/bot.py | 50 +++++++++++++- commands/__init__.py | 1 + commands/help.py | 1 + commands/newroom.py | 9 +++ commands/space.py | 136 ++++++++++++++++++++++++++++++++++++++ migrations/__init__.py | 2 +- migrations/migration_7.py | 23 +++++++ 8 files changed, 229 insertions(+), 2 deletions(-) create mode 100644 commands/space.py create mode 100644 migrations/migration_7.py diff --git a/callbacks/join.py b/callbacks/join.py index b424c13..2341c89 100644 --- a/callbacks/join.py +++ b/callbacks/join.py @@ -4,4 +4,13 @@ async def join_callback(response, bot): bot.matrix_client.joined_rooms() + with bot.database.cursor() as cursor: + cursor.execute( + "SELECT space_id FROM user_spaces WHERE user_id = ? AND active = TRUE", (event.sender,)) + space = cursor.fetchone() + + if space: + 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.send_message(bot.matrix_client.rooms[response.room_id], "Hello! Thanks for inviting me! How can I help you today?") \ No newline at end of file diff --git a/classes/bot.py b/classes/bot.py index f670d8e..b88a9a2 100644 --- a/classes/bot.py +++ b/classes/bot.py @@ -26,6 +26,9 @@ from nio import ( JoinError, RoomLeaveError, RoomSendError, + RoomVisibility, + RoomCreateResponse, + RoomCreateError, ) from nio.crypto import Olm @@ -64,7 +67,7 @@ class GPTBot: classification_api: Optional[OpenAI] = None parcel_api: Optional[TrackingMore] = 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 @classmethod @@ -552,6 +555,51 @@ class GPTBot: self.logger.log("Syncing one last time...") await self.matrix_client.sync(timeout=30000) + async def create_space(self, name, visibility=RoomVisibility.private) -> str: + """Create a space. + + Args: + name (str): The name of the space. + visibility (RoomVisibility, optional): The visibility of the space. Defaults to RoomVisibility.private. + + Returns: + MatrixRoom: The created space. + """ + + response = await self.matrix_client.room_create( + name=name, visibility=visibility, space=True) + + if isinstance(response, RoomCreateError): + self.logger.log( + f"Error creating space: {response.message}", "error") + return + + return response.room_id + + async def add_rooms_to_space(self, space: MatrixRoom | str, rooms: List[MatrixRoom | str]): + """Add rooms to a space. + + Args: + space (MatrixRoom | str): The space to add the rooms to. + rooms (List[MatrixRoom | str]): The rooms to add to the space. + """ + + if isinstance(space, MatrixRoom): + space = space.room_id + + for room in rooms: + if isinstance(room, MatrixRoom): + room = room.room_id + + await self.matrix_client.room_put_state(space, "m.space.child", { + "via": [room.split(":")[1], space.split(":")[1]], + }, room) + + await self.matrix_client.room_put_state(room, "m.room.parent", { + "via": [space.split(":")[1], room.split(":")[1]], + "canonical": True + }, space) + def respond_to_room_messages(self, room: MatrixRoom | str) -> bool: """Check whether the bot should respond to all messages sent in a room. diff --git a/commands/__init__.py b/commands/__init__.py index c638982..e6ec6b1 100644 --- a/commands/__init__.py +++ b/commands/__init__.py @@ -21,6 +21,7 @@ for command in [ "roomsettings", "dice", "parcel", + "space", ]: function = getattr(import_module( "commands." + command), "command_" + command) diff --git a/commands/help.py b/commands/help.py index 0b2c8a1..db1a317 100644 --- a/commands/help.py +++ b/commands/help.py @@ -11,6 +11,7 @@ async def command_help(room: MatrixRoom, event: RoomMessageText, bot): - !gptbot newroom \ - Create a new room and invite yourself to it - !gptbot stats - Show usage statistics for this room - !gptbot systemmessage \ - Get or set the system message for this room +- !gptbot space [enable|disable|update|invite] - Enable, disable, force update, or invite yourself to your space - !gptbot coin - Flip a coin (heads or tails) - !gptbot dice [number] - Roll a dice with the specified number of sides (default: 6) - !gptbot imagine \ - Generate an image from a prompt diff --git a/commands/newroom.py b/commands/newroom.py index 169cc61..9a0d19d 100644 --- a/commands/newroom.py +++ b/commands/newroom.py @@ -23,6 +23,15 @@ async def command_newroom(room: MatrixRoom, event: RoomMessageText, bot): await bot.send_message(room, f"Sorry, I was unable to invite you to the new room. Please try again later, or create a room manually.", True) return + with bot.database.cursor() as cursor: + cursor.execute( + "SELECT space_id FROM user_spaces WHERE user_id = ? AND active = TRUE", (event.sender,)) + space = cursor.fetchone() + + if space: + 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.matrix_client.room_put_state( new_room.room_id, "m.room.power_levels", {"users": {event.sender: 100}}) diff --git a/commands/space.py b/commands/space.py new file mode 100644 index 0000000..68d58ee --- /dev/null +++ b/commands/space.py @@ -0,0 +1,136 @@ +from nio.events.room_events import RoomMessageText +from nio.rooms import MatrixRoom +from nio.responses import RoomInviteError + + +async def command_space(room: MatrixRoom, event: RoomMessageText, bot): + if len(event.body.split()) == 3: + request = event.body.split()[2] + + if request.lower() == "enable": + bot.logger.log("Enabling space...") + + with bot.database.cursor() as cursor: + cursor.execute( + "SELECT space_id FROM user_spaces WHERE user_id = ? AND active = TRUE", (event.sender,)) + space = cursor.fetchone() + + if not space: + space = await bot.create_space("GPTBot") + bot.logger.log( + f"Created space {space} for user {event.sender}") + + with bot.database.cursor() as cursor: + cursor.execute( + "INSERT INTO user_spaces (space_id, user_id) VALUES (?, ?)", (space, event.sender)) + + else: + space = space[0] + + response = await bot.matrix_client.room_invite(space, event.sender) + + if isinstance(response, RoomInviteError): + bot.logger.log( + f"Failed to invite user {event.sender} to space {space}", "error") + await bot.send_message( + room, "Sorry, I couldn't invite you to the space. Please try again later.", True) + return + + bot.database.commit() + await bot.send_message(room, "Space enabled.", True) + request = "update" + + elif request.lower() == "disable": + bot.logger.log("Disabling space...") + + with bot.database.cursor() as cursor: + cursor.execute( + "SELECT space_id FROM user_spaces WHERE user_id = ? AND active = TRUE", (event.sender,)) + space = cursor.fetchone()[0] + + if not space: + bot.logger.log(f"User {event.sender} does not have a space") + await bot.send_message(room, "You don't have a space enabled.", True) + return + + with bot.database.cursor() as cursor: + cursor.execute( + "UPDATE user_spaces SET active = FALSE WHERE user_id = ?", (event.sender,)) + + bot.database.commit() + await bot.send_message(room, "Space disabled.", True) + return + + if request.lower() == "update": + bot.logger.log("Updating space...") + + with bot.database.cursor() as cursor: + cursor.execute( + "SELECT space_id FROM user_spaces WHERE user_id = ? AND active = TRUE", (event.sender,)) + space = cursor.fetchone()[0] + + if not space: + bot.logger.log(f"User {event.sender} does not have a space") + await bot.send_message( + room, "You don't have a space enabled. Create one first using `!gptbot space enable`.", True) + return + + rooms = bot.matrix_client.rooms + + join_rooms = [] + + for room in rooms.values(): + if event.sender in room.users.keys(): + bot.logger.log( + f"Adding room {room.room_id} to space {space}") + join_rooms.append(room.room_id) + + await bot.add_rooms_to_space(space, join_rooms) + + await bot.send_message(room, "Space updated.", True) + return + + if request.lower() == "invite": + bot.logger.log("Inviting user to space...") + + with bot.database.cursor() as cursor: + cursor.execute( + "SELECT space_id FROM user_spaces WHERE user_id = ?", (event.sender,)) + space = cursor.fetchone()[0] + + if not space: + bot.logger.log(f"User {event.sender} does not have a space") + await bot.send_message( + room, "You don't have a space enabled. Create one first using `!gptbot space enable`.", True) + return + + response = await bot.matrix_client.room_invite(space, event.sender) + + if isinstance(response, RoomInviteError): + bot.logger.log( + f"Failed to invite user {user} to space {space}", "error") + await bot.send_message( + room, "Sorry, I couldn't invite you to the space. Please try again later.", True) + return + + await bot.send_message(room, "Invited you to the space.", True) + return + + with bot.database.cursor() as cursor: + cursor.execute( + "SELECT active FROM user_spaces WHERE user_id = ?", (event.sender,)) + status = cursor.fetchone() + + if not status: + await bot.send_message( + room, "You don't have a space enabled. Create one using `!gptbot space enable`.", True) + return + + if not status[0]: + await bot.send_message( + room, "Your space is disabled. Enable it using `!gptbot space enable`.", True) + return + + await bot.send_message( + room, "Your space is enabled. Rooms will be added to it automatically.", True) + return diff --git a/migrations/__init__.py b/migrations/__init__.py index 53b0ecd..1c0fa51 100644 --- a/migrations/__init__.py +++ b/migrations/__init__.py @@ -4,7 +4,7 @@ from importlib import import_module from duckdb import DuckDBPyConnection -MAX_MIGRATION = 6 +MAX_MIGRATION = 7 MIGRATIONS = OrderedDict() diff --git a/migrations/migration_7.py b/migrations/migration_7.py new file mode 100644 index 0000000..56d30df --- /dev/null +++ b/migrations/migration_7.py @@ -0,0 +1,23 @@ +# Migration to add user_spaces table + +from datetime import datetime + +def migration(conn): + with conn.cursor() as cursor: + cursor.execute( + """ + CREATE TABLE user_spaces ( + space_id TEXT NOT NULL, + user_id TEXT NOT NULL, + active BOOLEAN NOT NULL DEFAULT TRUE, + PRIMARY KEY (space_id, user_id) + ) + """ + ) + + cursor.execute( + "INSERT INTO migrations (id, timestamp) VALUES (7, ?)", + (datetime.now(),) + ) + + conn.commit() \ No newline at end of file