From 85cfecf88a1e0d7ac0f8aa06e608b75e97bff9c5 Mon Sep 17 00:00:00 2001 From: Kumi Date: Tue, 2 May 2023 06:58:49 +0000 Subject: [PATCH] Add option for timing message processing Fixes --- callbacks/message.py | 18 ++++++++++++++++-- callbacks/roommember.py | 2 +- classes/bot.py | 33 ++++++++++++++++++++++++++++++--- classes/dict.py | 8 ++++++++ classes/store.py | 31 +++++++++++++++++++++++++++---- commands/roomsettings.py | 8 ++++---- 6 files changed, 86 insertions(+), 14 deletions(-) create mode 100644 classes/dict.py diff --git a/callbacks/message.py b/callbacks/message.py index 68e0a3e..36685a6 100644 --- a/callbacks/message.py +++ b/callbacks/message.py @@ -1,8 +1,14 @@ from nio import MatrixRoom, RoomMessageText, MegolmEvent -async def message_callback(room: MatrixRoom, event: RoomMessageText | MegolmEvent, bot): +from datetime import datetime + +async def message_callback(room: MatrixRoom | str, event: RoomMessageText | MegolmEvent, bot): bot.logger.log(f"Received message from {event.sender} in room {room.room_id}") + sent = datetime.fromtimestamp(event.server_timestamp / 1000) + received = datetime.now() + latency = received - sent + if isinstance(event, MegolmEvent): try: event = await bot.matrix_client.decrypt_event(event) @@ -27,4 +33,12 @@ async def message_callback(room: MatrixRoom, event: RoomMessageText | MegolmEven bot.logger.log(f"Received {event.body} - might be a command, but not for this bot - ignoring") else: - await bot.process_query(room, event) \ No newline at end of file + await bot.process_query(room, event) + + processed = datetime.now() + processing_time = processed - received + + bot.logger.log(f"Message processing took {processing_time.total_seconds()} seconds (latency: {latency.total_seconds()} seconds)") + + if bot.room_uses_timing(room): + await bot.send_message(room, f"Message processing took {processing_time.total_seconds()} seconds (latency: {latency.total_seconds()} seconds)", True) \ No newline at end of file diff --git a/callbacks/roommember.py b/callbacks/roommember.py index 20c5c76..662de19 100644 --- a/callbacks/roommember.py +++ b/callbacks/roommember.py @@ -6,5 +6,5 @@ async def roommember_callback(room: MatrixRoom, event: RoomMemberEvent, bot): if len(room.users) == 1: bot.logger.log("Yes, I was abandoned - leaving...") - await bot.matrix_client.leave(room.room_id) + await bot.matrix_client.room_leave(room.room_id) return diff --git a/classes/bot.py b/classes/bot.py index b216c0c..4d1baf2 100644 --- a/classes/bot.py +++ b/classes/bot.py @@ -256,7 +256,7 @@ class GPTBot: return False if not result else bool(int(result[0])) - async def event_callback(self, room: MatrixRoom, event: Event): + async def _event_callback(self, room: MatrixRoom, event: Event): self.logger.log("Received event: " + str(event.event_id), "debug") try: for eventtype, callback in EVENT_CALLBACKS.items(): @@ -266,11 +266,35 @@ class GPTBot: self.logger.log( f"Error in event callback for {event.__class__}: {e}", "error") - async def response_callback(self, response: Response): + async def event_callback(self, room: MatrixRoom, event: Event): + task = asyncio.create_task(self._event_callback(room, event)) + + def room_uses_timing(self, room: MatrixRoom): + """Check if a room uses timing. + + Args: + room (MatrixRoom): The room to check. + + Returns: + bool: Whether the room uses timing. + """ + room_id = room.room_id + + with self.database.cursor() as cursor: + cursor.execute( + "SELECT value FROM room_settings WHERE room_id = ? AND setting = ?", (room_id, "use_timing")) + result = cursor.fetchone() + + return False if not result else bool(int(result[0])) + + async def _response_callback(self, response: Response): for response_type, callback in RESPONSE_CALLBACKS.items(): if isinstance(response, response_type): await callback(response, self) + async def response_callback(self, response: Response): + task = asyncio.create_task(self._response_callback(response)) + async def accept_pending_invites(self): """Accept all pending invites.""" @@ -353,7 +377,7 @@ class GPTBot: self.logger.log("Sent image") - async def send_message(self, room: MatrixRoom, message: str, notice: bool = False): + async def send_message(self, room: MatrixRoom | str, message: str, notice: bool = False): """Send a message to a room. Args: @@ -362,6 +386,9 @@ class GPTBot: notice (bool): Whether to send the message as a notice. Defaults to False. """ + if isinstance(room, str): + room = self.matrix_client.rooms[room] + markdowner = markdown2.Markdown(extras=["fenced-code-blocks"]) formatted_body = markdowner.convert(message) diff --git a/classes/dict.py b/classes/dict.py new file mode 100644 index 0000000..7dfd271 --- /dev/null +++ b/classes/dict.py @@ -0,0 +1,8 @@ +class AttrDict(dict): + def __getattr__(self, key): + if key in self: + return self[key] + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{key}'") + + def __setattr__(self, key, value): + self[key] = value \ No newline at end of file diff --git a/classes/store.py b/classes/store.py index 6d05707..e999abf 100644 --- a/classes/store.py +++ b/classes/store.py @@ -7,6 +7,8 @@ from random import SystemRandom from collections import defaultdict from typing import Dict, List, Optional, Tuple +from .dict import AttrDict + import json @@ -458,8 +460,14 @@ class DuckDBStore(MatrixStore): rows = cur.fetchall() return { - request.request_id: OutgoingKeyRequest.from_database(request) - for request in rows + row[1]: OutgoingKeyRequest.from_response(AttrDict({ + "id": row[0], + "account_id": row[1], + "request_id": row[2], + "session_id": row[3], + "room_id": row[4], + "algorithm": row[5], + })) for row in rows } def load_encrypted_rooms(self): @@ -571,15 +579,30 @@ class DuckDBStore(MatrixStore): insert_query, (account, session.sender_key, session.ed25519, session.room_id, chain)) def add_outgoing_key_request(self, key_request): + """Add a new outgoing key request to the database. + + Args: + key_request (OutgoingKeyRequest): The key request to add. + """ + account_id = self.account_id with self.conn.cursor() as cursor: cursor.execute( """ - INSERT INTO outgoing_key_requests (account_id, request_id, session_id, room_id, algorithm) - VALUES (?, ?, ?, ?, ?) + SELECT MAX(id) FROM outgoing_key_requests + """ + ) + row = cursor.fetchone() + request_id = row[0] + 1 if row[0] else 1 + + cursor.execute( + """ + INSERT INTO outgoing_key_requests (id, account_id, request_id, session_id, room_id, algorithm) + VALUES (?, ?, ?, ?, ?, ?) ON CONFLICT (account_id, request_id) DO NOTHING """, ( + request_id, account_id, key_request.request_id, key_request.session_id, diff --git a/commands/roomsettings.py b/commands/roomsettings.py index 999b226..ebe210e 100644 --- a/commands/roomsettings.py +++ b/commands/roomsettings.py @@ -7,8 +7,8 @@ async def command_roomsettings(room: MatrixRoom, event: RoomMessageText, bot): value = " ".join(event.body.split()[3:]) if len( event.body.split()) > 3 else None - if setting == "classification": - setting = "use_classification" + if setting in ("classification", "timing"): + setting = f"use_{setting}" if setting == "systemmessage": setting = "system_message" @@ -33,7 +33,7 @@ async def command_roomsettings(room: MatrixRoom, event: RoomMessageText, bot): await bot.send_message(room, f"The current system message is: '{system_message}'.", True) return - if setting in ("use_classification", "always_reply"): + if setting in ("use_classification", "always_reply", "use_timing"): if value: if value.lower() in ["true", "false"]: value = value.lower() == "true" @@ -64,7 +64,7 @@ async def command_roomsettings(room: MatrixRoom, event: RoomMessageText, bot): value = cur.fetchone()[0] if not value: - if setting == "use_classification": + if setting in ("use_classification", "use_timing"): value = False elif setting == "always_reply": value = True