From 5997ee8ab18aaeb201e74d19835b6de80942c876 Mon Sep 17 00:00:00 2001 From: Kumi Date: Mon, 1 May 2023 08:12:50 +0000 Subject: [PATCH] Implement chat message classification !gptbot roomsettings command Permit custom commands (!gptbot custom ...) --- classes/bot.py | 34 ++++++++++++++++++-- commands/__init__.py | 52 +++++++++++++++---------------- commands/chat.py | 2 +- commands/custom.py | 9 ++++++ commands/help.py | 1 + commands/imagine.py | 2 +- commands/roomsettings.py | 65 +++++++++++++++++++++++++++++++++++++++ commands/systemmessage.py | 10 +++--- migrations/__init__.py | 12 +++----- migrations/migration_5.py | 48 +++++++++++++++++++++++++++++ migrations/migration_6.py | 32 +++++++++++++++++++ 11 files changed, 224 insertions(+), 43 deletions(-) create mode 100644 commands/custom.py create mode 100644 commands/roomsettings.py create mode 100644 migrations/migration_5.py create mode 100644 migrations/migration_6.py diff --git a/classes/bot.py b/classes/bot.py index 5b841a0..2ffc175 100644 --- a/classes/bot.py +++ b/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() diff --git a/commands/__init__.py b/commands/__init__.py index b75b567..d842211 100644 --- a/commands/__init__.py +++ b/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 .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 +from importlib import import_module -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, -} +from .unknown import 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 diff --git a/commands/chat.py b/commands/chat.py index 83f9654..3ecb534 100644 --- a/commands/chat.py +++ b/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 diff --git a/commands/custom.py b/commands/custom.py new file mode 100644 index 0000000..16272eb --- /dev/null +++ b/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 \ No newline at end of file diff --git a/commands/help.py b/commands/help.py index a715055..5c2de33 100644 --- a/commands/help.py +++ b/commands/help.py @@ -17,6 +17,7 @@ async def command_help(room: MatrixRoom, event: RoomMessageText, bot): - !gptbot privacy - Show privacy information - !gptbot chat \ - Send a message to the chat API - !gptbot classify \ - Classify a message using the classification API +- !gptbot custom \ - Used for custom commands handled by the chat model and defined through the room's system message """ await bot.send_message(room, body, True) \ No newline at end of file diff --git a/commands/imagine.py b/commands/imagine.py index c7aa099..54e6c71 100644 --- a/commands/imagine.py +++ b/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 diff --git a/commands/roomsettings.py b/commands/roomsettings.py new file mode 100644 index 0000000..267674b --- /dev/null +++ b/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) diff --git a/commands/systemmessage.py b/commands/systemmessage.py index 39f58e8..5309ff2 100644 --- a/commands/systemmessage.py +++ b/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) diff --git a/migrations/__init__.py b/migrations/__init__.py index c753bc6..53b0ecd 100644 --- a/migrations/__init__.py +++ b/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. diff --git a/migrations/migration_5.py b/migrations/migration_5.py new file mode 100644 index 0000000..9b4eeb8 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/migrations/migration_6.py b/migrations/migration_6.py new file mode 100644 index 0000000..791f167 --- /dev/null +++ b/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() \ No newline at end of file