Implement chat message classification

!gptbot roomsettings command
Permit custom commands (!gptbot custom ...)

classes/bot.py

@@ -235,6 +235,24 @@ class GPTBot:
 
         await COMMANDS.get(command, COMMANDS[None])(room, event, self)
 
+    def room_uses_classification(self, room: MatrixRoom | int) -> bool:
+        """Check if a room uses classification.
+
+        Args:
+            room (MatrixRoom): The room to check.
+
+        Returns:
+            bool: Whether the room uses classification.
+        """
+        room_id = room.room_id if isinstance(room, MatrixRoom) else room
+
+        with self.database.cursor() as cursor:
+            cursor.execute(
+                "SELECT value FROM room_settings WHERE room_id = ? AND setting = ?", (room_id, "use_classification"))
+            result = cursor.fetchone()
+
+        return False if not result else bool(int(result[0]))
+
     async def event_callback(self, room: MatrixRoom, event: Event):
         self.logger.log("Received event: " + str(event.event_id), "debug")
         try:
@@ -456,11 +474,21 @@ class GPTBot:
             self.logger.log("Syncing one last time...")
             await self.matrix_client.sync(timeout=30000)
 
-    async def process_query(self, room: MatrixRoom, event: RoomMessageText):
+    async def process_query(self, room: MatrixRoom, event: RoomMessageText, allow_classify: bool = True):
         await self.matrix_client.room_typing(room.room_id, True)
 
         await self.matrix_client.room_read_markers(room.room_id, event.event_id)
 
+        if allow_classify and self.room_uses_classification(room):
+            classification, tokens = self.classification_api.classify_message(event.body, room.room_id)
+
+            self.log_api_usage(event, room, f"{self.classification_api.api_code}-{self.classification_api.classification_api}", tokens)
+
+            if not classification["type"] == "chat":
+                event.body = f"!gptbot {classification['type']} {classification['prompt']}"
+                await self.process_command(room, event)
+                return
+
         try:
             last_messages = await self._last_n_messages(room.room_id, 20)
         except Exception as e:
@@ -520,8 +548,8 @@ class GPTBot:
 
         with self.database.cursor() as cur:
             cur.execute(
-                "SELECT body FROM system_messages WHERE room_id = ? ORDER BY timestamp DESC LIMIT 1",
-                (room_id,)
+                "SELECT value FROM room_settings WHERE room_id = ? AND setting = ?",
+                (room_id, "system_message")
             )
             system_message = cur.fetchone()
 

commands/__init__.py

@@ -1,27 +1,27 @@
-from .help import command_help
-from .newroom import command_newroom
-from .stats import command_stats
-from .botinfo import command_botinfo
+from importlib import import_module
+
 from .unknown import command_unknown
-from .coin import command_coin
-from .ignoreolder import command_ignoreolder
-from .systemmessage import command_systemmessage
-from .imagine import command_imagine
-from .calculate import command_calculate
-from .classify import command_classify
-from .chat import command_chat
 
-COMMANDS = {
-    "help": command_help,
-    "newroom": command_newroom,
-    "stats": command_stats,
-    "botinfo": command_botinfo,
-    "coin": command_coin,
-    "ignoreolder": command_ignoreolder,
-    "systemmessage": command_systemmessage,
-    "imagine": command_imagine,
-    "calculate": command_calculate,
-    "classify": command_classify,
-    "chat": command_chat,
-    None: command_unknown,
-}
+COMMANDS = {}
+
+for command in [
+    "help",
+    "newroom",
+    "stats",
+    "botinfo",
+    "coin",
+    "ignoreolder",
+    "systemmessage",
+    "imagine",
+    "calculate",
+    "classify",
+    "chat",
+    "custom",
+    "privacy",
+    "roomsettings",
+]:
+    function = getattr(import_module(
+        "commands." + command), "command_" + command)
+    COMMANDS[command] = function
+
+COMMANDS[None] = command_unknown

commands/chat.py

@@ -8,7 +8,7 @@ async def command_chat(room: MatrixRoom, event: RoomMessageText, bot):
     if prompt:
         bot.logger.log("Sending chat message...")
         event.body = prompt
-        await bot.process_query(room, event)
+        await bot.process_query(room, event, allow_classify=False)
 
         return
 

commands/custom.py

@@ -0,0 +1,9 @@
+from nio.events.room_events import RoomMessageText
+from nio.rooms import MatrixRoom
+
+
+async def command_custom(room: MatrixRoom, event: RoomMessageText, bot):
+    bot.logger.log("Forwarding custom command to room...")
+    await bot.process_query(room, event)
+
+    return

commands/help.py

@@ -17,6 +17,7 @@ async def command_help(room: MatrixRoom, event: RoomMessageText, bot):
 - !gptbot privacy - Show privacy information
 - !gptbot chat \<message\> - Send a message to the chat API
 - !gptbot classify \<message\> - Classify a message using the classification API
+- !gptbot custom \<message\> - Used for custom commands handled by the chat model and defined through the room's system message
 """
 
     await bot.send_message(room, body, True)

commands/imagine.py

@@ -14,7 +14,7 @@ async def command_imagine(room: MatrixRoom, event: RoomMessageText, bot):
             bot.logger.log(f"Sending image...")
             await bot.send_image(room, image)
 
-        bot.log_api_usage(event, room, f"{self.image_api.api_code}-{self.image_api.image_api}", tokens_used)
+        bot.log_api_usage(event, room, f"{bot.image_api.api_code}-{bot.image_api.image_api}", tokens_used)
 
         return
 

commands/roomsettings.py

@@ -0,0 +1,65 @@
+from nio.events.room_events import RoomMessageText
+from nio.rooms import MatrixRoom
+
+
+async def command_roomsettings(room: MatrixRoom, event: RoomMessageText, bot):
+    setting = event.body.split()[2]
+    value = " ".join(event.body.split()[3:]) if len(
+        event.body.split()) > 3 else None
+
+    if setting == "system_message":
+        if value:
+            bot.logger.log("Adding system message...")
+
+            with bot.database.cursor() as cur:
+                cur.execute(
+                    """INSERT INTO room_settings (room_id, setting, value) VALUES (?, ?, ?)
+                    ON CONFLICT (room_id, setting) DO UPDATE SET value = ?;""",
+                    (room.room_id, "system_message", value, value)
+                )
+
+            await bot.send_message(room, f"Alright, I've stored the system message: '{value}'.", True)
+            return
+
+        bot.logger.log("Retrieving system message...")
+
+        system_message = bot.get_system_message(room)
+
+        await bot.send_message(room, f"The current system message is: '{system_message}'.", True)
+        return
+
+    if setting == "classification":
+        if value:
+            if value.lower() in ["true", "false"]:
+                value = value.lower() == "true"
+
+                bot.logger.log("Setting classification status...")
+
+                with bot.database.cursor() as cur:
+                    cur.execute(
+                        """INSERT INTO room_settings (room_id, setting, value) VALUES (?, ?, ?)
+                        ON CONFLICT (room_id, setting) DO UPDATE SET value = ?;""",
+                        (room.room_id, "use_classification", "1" if value else "0", "1" if value else "0")
+                    )
+
+                await bot.send_message(room, f"Alright, I've set use_classification to: '{value}'.", True)
+                return
+
+            await bot.send_message(room, "You need to provide a boolean value (true/false).", True)
+            return
+
+        bot.logger.log("Retrieving classification status...")
+
+        use_classification = await bot.room_uses_classification(room)
+
+        await bot.send_message(room, f"The current classification status is: '{use_classification}'.", True)
+        return
+
+    message = f"""
+    The following settings are available:
+
+    - system_message [message]: Get or set the system message to be sent to the chat model
+    - classification [true/false]: Get or set whether the room uses classification
+    """
+
+    await bot.send_message(room, message, True)

commands/systemmessage.py

@@ -10,9 +10,11 @@ async def command_systemmessage(room: MatrixRoom, event: RoomMessageText, bot):
 
         with bot.database.cursor() as cur:
             cur.execute(
-                "INSERT INTO system_messages (room_id, message_id, user_id, body, timestamp) VALUES (?, ?, ?, ?, ?)",
-                (room.room_id, event.event_id, event.sender,
-                 system_message, event.server_timestamp)
+                """
+                INSERT INTO room_settings (room_id, setting, value) VALUES (?, ?, ?)
+                ON CONFLICT (room_id, setting) DO UPDATE SET value = ?;
+                """,
+                (room.room_id, "system_message", system_message, system_message)
             )
 
         await bot.send_message(room, f"Alright, I've stored the system message: '{system_message}'.", True)
@@ -22,4 +24,4 @@ async def command_systemmessage(room: MatrixRoom, event: RoomMessageText, bot):
 
     system_message = bot.get_system_message(room)
 
-    bot.send_message(room, f"The current system message is: '{system_message}'.", True)
+    await bot.send_message(room, f"The current system message is: '{system_message}'.", True)

migrations/__init__.py

@@ -1,19 +1,15 @@
 from collections import OrderedDict
 from typing import Optional
+from importlib import import_module
 
 from duckdb import DuckDBPyConnection
 
-from .migration_1 import migration as migration_1
-from .migration_2 import migration as migration_2
-from .migration_3 import migration as migration_3
-from .migration_4 import migration as migration_4
+MAX_MIGRATION = 6
 
 MIGRATIONS = OrderedDict()
 
-MIGRATIONS[1] = migration_1
-MIGRATIONS[2] = migration_2
-MIGRATIONS[3] = migration_3
-MIGRATIONS[4] = migration_4
+for i in range(1, MAX_MIGRATION + 1):
+    MIGRATIONS[i] = import_module(f".migration_{i}", __package__).migration
 
 def get_version(db: DuckDBPyConnection) -> int:
     """Get the current database version.

migrations/migration_5.py

@@ -0,0 +1,48 @@
+# Migration to add room settings table
+
+from datetime import datetime
+
+def migration(conn):
+    with conn.cursor() as cursor:
+        cursor.execute(
+            """
+            CREATE TABLE IF NOT EXISTS room_settings (
+                room_id TEXT NOT NULL,
+                setting TEXT NOT NULL,
+                value TEXT NOT NULL,
+                PRIMARY KEY (room_id, setting)
+            )
+            """
+        )
+
+        cursor.execute("SELECT * FROM system_messages")
+        system_messages = cursor.fetchall()
+
+        # Get latest system message for each room
+
+        cursor.execute(
+            """
+            SELECT system_messages.room_id, system_messages.message_id, system_messages.user_id, system_messages.body, system_messages.timestamp
+            FROM system_messages
+            INNER JOIN (
+                SELECT room_id, MAX(timestamp) AS timestamp FROM system_messages GROUP BY room_id
+            ) AS latest_system_message ON system_messages.room_id = latest_system_message.room_id AND system_messages.timestamp = latest_system_message.timestamp
+            """
+        )
+
+        system_messages = cursor.fetchall()
+
+        for message in system_messages:
+            cursor.execute(
+                "INSERT INTO room_settings (room_id, setting, value) VALUES (?, ?, ?)",
+                (message[0], "system_message", message[1])
+            )
+
+        cursor.execute("DROP TABLE system_messages")
+
+        cursor.execute(
+            "INSERT INTO migrations (id, timestamp) VALUES (5, ?)",
+            (datetime.now(),)
+        )
+
+        conn.commit()

migrations/migration_6.py

@@ -0,0 +1,32 @@
+# Migration to drop primary key constraint from token_usage table
+
+from datetime import datetime
+
+def migration(conn):
+    with conn.cursor() as cursor:
+        cursor.execute(
+            """
+            CREATE TABLE token_usage_temp (
+                message_id TEXT NOT NULL,
+                room_id TEXT NOT NULL,
+                api TEXT NOT NULL,
+                tokens INTEGER NOT NULL,
+                timestamp TIMESTAMP NOT NULL
+            )
+            """
+        )
+
+        cursor.execute(
+            "INSERT INTO token_usage_temp SELECT message_id, room_id, api, tokens, timestamp FROM token_usage"
+        )
+
+        cursor.execute("DROP TABLE token_usage")
+
+        cursor.execute("ALTER TABLE token_usage_temp RENAME TO token_usage")
+
+        cursor.execute(
+            "INSERT INTO migrations (id, timestamp) VALUES (6, ?)",
+            (datetime.now(),)
+        )
+
+        conn.commit()