Browse Source

Add spaces feature

Kumi 1 year ago
parent
commit
3d32343e54
8 changed files with 229 additions and 2 deletions
  1. 9 0
      callbacks/join.py
  2. 49 1
      classes/bot.py
  3. 1 0
      commands/__init__.py
  4. 1 0
      commands/help.py
  5. 9 0
      commands/newroom.py
  6. 136 0
      commands/space.py
  7. 1 1
      migrations/__init__.py
  8. 23 0
      migrations/migration_7.py

+ 9 - 0
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?")

+ 49 - 1
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.
 

+ 1 - 0
commands/__init__.py

@@ -21,6 +21,7 @@ for command in [
     "roomsettings",
     "dice",
     "parcel",
+    "space",
 ]:
     function = getattr(import_module(
         "commands." + command), "command_" + command)

+ 1 - 0
commands/help.py

@@ -11,6 +11,7 @@ async def command_help(room: MatrixRoom, event: RoomMessageText, bot):
 - !gptbot newroom \<room name\> - Create a new room and invite yourself to it
 - !gptbot stats - Show usage statistics for this room
 - !gptbot systemmessage \<message\> - 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 \<prompt\> - Generate an image from a prompt

+ 9 - 0
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}})
 

+ 136 - 0
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

+ 1 - 1
migrations/__init__.py

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

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