Add spaces feature

This commit is contained in:
Kumi 2023-05-09 10:27:03 +00:00
parent 7fbeff9d9c
commit 3d32343e54
Signed by: kumi
GPG key ID: ECBCC9082395383F
8 changed files with 229 additions and 2 deletions

View file

@ -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?")

View file

@ -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.

View file

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

View file

@ -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

View file

@ -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
commands/space.py Normal file
View file

@ -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

View file

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

23
migrations/migration_7.py Normal file
View file

@ -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()