Browse Source

Add option for timing message processing
Fixes

Kumi 1 year ago
parent
commit
85cfecf88a
6 changed files with 86 additions and 14 deletions
  1. 16 2
      callbacks/message.py
  2. 1 1
      callbacks/roommember.py
  3. 30 3
      classes/bot.py
  4. 8 0
      classes/dict.py
  5. 27 4
      classes/store.py
  6. 4 4
      commands/roomsettings.py

+ 16 - 2
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)
+        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)

+ 1 - 1
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

+ 30 - 3
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)
 

+ 8 - 0
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

+ 27 - 4
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,

+ 4 - 4
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