Implement chat message classification

!gptbot roomsettings command
Permit custom commands (!gptbot custom ...)
This commit is contained in:
Kumi 2023-05-01 08:12:50 +00:00
parent 2fb607310d
commit 5997ee8ab1
Signed by: kumi
GPG key ID: ECBCC9082395383F
11 changed files with 224 additions and 43 deletions

View file

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

View file

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

View file

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

9
commands/custom.py Normal file
View file

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

View file

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

View file

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

65
commands/roomsettings.py Normal file
View file

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

View file

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

View file

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

48
migrations/migration_5.py Normal file
View file

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

32
migrations/migration_6.py Normal file
View file

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