Browse Source

TrackingMore integration
Fix newroom command

Kumi 1 year ago
parent
commit
7fbeff9d9c

+ 3 - 0
callbacks/__init__.py

@@ -9,6 +9,7 @@ from nio import (
     OlmEvent,
     MegolmEvent,
     RoomMemberEvent,
+    Response,
 )
 
 from .test import test_callback
@@ -17,8 +18,10 @@ from .invite import room_invite_callback
 from .join import join_callback
 from .message import message_callback
 from .roommember import roommember_callback
+from .test_response import test_response_callback
 
 RESPONSE_CALLBACKS = {
+    Response: test_response_callback,
     SyncResponse: sync_callback,
     JoinResponse: join_callback,
 }

+ 9 - 2
callbacks/message.py

@@ -1,4 +1,4 @@
-from nio import MatrixRoom, RoomMessageText, MegolmEvent
+from nio import MatrixRoom, RoomMessageText, MegolmEvent, RoomKeyRequestError, RoomKeyRequestResponse
 
 from datetime import datetime
 
@@ -15,7 +15,14 @@ async def message_callback(room: MatrixRoom | str, event: RoomMessageText | Mego
         except Exception as e:
             try:
                 bot.logger.log("Requesting new encryption keys...")
-                await bot.matrix_client.request_room_key(event)
+                response = await bot.matrix_client.request_room_key(event)
+
+                if isinstance(response, RoomKeyRequestError):
+                    bot.logger.log(f"Error requesting encryption keys: {response}", "error")
+                elif isinstance(response, RoomKeyRequestResponse):
+                    bot.logger.log(f"Encryption keys received: {response}", "debug")
+                    bot.matrix_bot.olm.handle_response(response)
+                    event = await bot.matrix_client.decrypt_event(event)
             except:
                 pass
 

+ 4 - 0
callbacks/test_response.py

@@ -0,0 +1,4 @@
+async def test_response_callback(response, bot):
+    bot.logger.log(
+        f"{response.__class__} response received", "debug")
+    

+ 20 - 1
classes/bot.py

@@ -25,6 +25,7 @@ from nio import (
     RoomMessageNotice,
     JoinError,
     RoomLeaveError,
+    RoomSendError,
 )
 from nio.crypto import Olm
 
@@ -34,6 +35,7 @@ from datetime import datetime
 from io import BytesIO
 
 import uuid
+import traceback
 
 from .logging import Logger
 from migrations import migrate
@@ -42,6 +44,7 @@ from commands import COMMANDS
 from .store import DuckDBStore
 from .openai import OpenAI
 from .wolframalpha import WolframAlpha
+from .trackingmore import TrackingMore
 
 
 class GPTBot:
@@ -59,8 +62,10 @@ class GPTBot:
     chat_api: Optional[OpenAI] = None
     image_api: Optional[OpenAI] = None
     classification_api: Optional[OpenAI] = None
+    parcel_api: Optional[TrackingMore] = None
     operator: Optional[str] = None
     room_ignore_list: List[str] = [] # List of rooms to ignore invites from
+    debug: bool = False
 
     @classmethod
     def from_config(cls, config: ConfigParser):
@@ -89,6 +94,7 @@ class GPTBot:
                 "SystemMessage", bot.default_system_message)
             bot.force_system_message = config["GPTBot"].getboolean(
                 "ForceSystemMessage", bot.force_system_message)
+            bot.debug = config["GPTBot"].getboolean("Debug", bot.debug)
 
         bot.chat_api = bot.image_api = bot.classification_api = OpenAI(
             config["OpenAI"]["APIKey"], config["OpenAI"].get("Model"), bot.logger)
@@ -101,6 +107,11 @@ class GPTBot:
             bot.calculation_api = WolframAlpha(
                 config["WolframAlpha"]["APIKey"], bot.logger)
 
+        # Set up TrackingMore
+        if "TrackingMore" in config:
+            bot.parcel_api = TrackingMore(
+                config["TrackingMore"]["APIKey"], bot.logger)
+
         # Set up the Matrix client
 
         assert "Matrix" in config, "Matrix config not found"
@@ -266,6 +277,9 @@ class GPTBot:
             self.logger.log(
                 f"Error in event callback for {event.__class__}: {e}", "error")
 
+            if self.debug:
+                await self.send_message(room, f"Error: {e}\n\n```\n{traceback.format_exc()}\n```", True)
+
     async def event_callback(self, room: MatrixRoom, event: Event):
         task = asyncio.create_task(self._event_callback(room, event))
 
@@ -433,7 +447,12 @@ class GPTBot:
             self.matrix_client.access_token, room.room_id, msgtype, content, uuid.uuid4()
         )
 
-        return await self.matrix_client._send(RoomSendResponse, method, path, data, (room.room_id,))
+        response = await self.matrix_client._send(RoomSendResponse, method, path, data, (room.room_id,))
+
+        if isinstance(response, RoomSendError):
+            self.logger.log(
+                f"Error sending message: {response.message}", "error")
+            return
 
     def log_api_usage(self, message: Event | str, room: MatrixRoom | str, api: str, tokens: int):
         """Log API usage to the database.

+ 32 - 0
classes/trackingmore.py

@@ -0,0 +1,32 @@
+import trackingmore
+import requests
+
+from .logging import Logger
+
+from typing import Dict, List, Tuple, Generator, Optional
+
+class TrackingMore:
+    api_key: str
+    logger: Logger
+    client: trackingmore.TrackingMore
+
+    api_code: str = "trackingmore"
+    parcel_api: str = "trackingmore"
+
+    operator: str = "TrackingMore ([https://www.trackingmore.com](https://www.trackingmore.com))"
+
+    def __init__(self, api_key: str, logger: Optional[Logger] = None):
+        self.api_key: str = api_key
+        self.logger: Logger = logger or Logger()
+        self.client = trackingmore.TrackingMore(self.api_key)
+
+    def lookup_parcel(self, query: str, carrier: Optional[str] = None, user: Optional[str] = None) -> Tuple[str, int]:
+        self.logger.log(f"Querying TrackingMore for {query}")
+
+        if query == "carriers":
+            response = "\n".join(f"* {carrier['courier_name']} - {carrier['courier_code']}" for carrier in self.client.get_carriers())
+            return response, 1
+
+        response = self.client.track_shipment(query)
+
+        return response, 1

+ 1 - 0
commands/__init__.py

@@ -20,6 +20,7 @@ for command in [
     "privacy",
     "roomsettings",
     "dice",
+    "parcel",
 ]:
     function = getattr(import_module(
         "commands." + command), "command_" + command)

+ 1 - 1
commands/newroom.py

@@ -28,4 +28,4 @@ async def command_newroom(room: MatrixRoom, event: RoomMessageText, bot):
 
     await bot.matrix_client.joined_rooms()
     await bot.send_message(room, f"Alright, I've created a new room called '{room_name}' and invited you to it. You can find it at {new_room.room_id}", True)
-    await bot.send_message(new_room, f"Welcome to the new room! What can I do for you?")
+    await bot.send_message(bot.rooms[new_room.room_id], f"Welcome to the new room! What can I do for you?")

+ 19 - 0
commands/parcel.py

@@ -0,0 +1,19 @@
+from nio.events.room_events import RoomMessageText
+from nio.rooms import MatrixRoom
+
+
+async def command_parcel(room: MatrixRoom, event: RoomMessageText, bot):
+    prompt = event.body.split()[2:]
+
+    if prompt:
+        bot.logger.log("Looking up parcels...")
+
+        for parcel in prompt:
+            status, tokens_used = bot.parcel_api.lookup_parcel(parcel, user=room.room_id)
+
+            await bot.send_message(room, status, True)
+
+            bot.log_api_usage(event, room, f"{bot.parcel_api.api_code}-{bot.parcel_api.parcel_api}", tokens_used)
+            return
+
+    await bot.send_message(room, "You need to provide tracking numbers.", True)

+ 7 - 0
config.dist.ini

@@ -95,3 +95,10 @@ Operator = Contact details not set
 # N.B.: Encryption doesn't work as it is supposed to anyway.
 
 Path = database.db
+
+[TrackingMore]
+
+# API key for TrackingMore
+# If not defined, the bot will not be able to provide parcel tracking
+#
+# APIKey = abcde-fghij-klmnop

+ 1 - 0
requirements.txt

@@ -6,3 +6,4 @@ duckdb
 python-magic
 pillow
 wolframalpha
+git+https://kumig.it/kumitterer/trackingmore-api-tool.git