diff --git a/callbacks/__init__.py b/callbacks/__init__.py index 047bebe..76575d0 100644 --- a/callbacks/__init__.py +++ b/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, } diff --git a/callbacks/message.py b/callbacks/message.py index 36685a6..5f5a176 100644 --- a/callbacks/message.py +++ b/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 diff --git a/callbacks/test_response.py b/callbacks/test_response.py new file mode 100644 index 0000000..51a6302 --- /dev/null +++ b/callbacks/test_response.py @@ -0,0 +1,4 @@ +async def test_response_callback(response, bot): + bot.logger.log( + f"{response.__class__} response received", "debug") + \ No newline at end of file diff --git a/classes/bot.py b/classes/bot.py index 4d1baf2..f670d8e 100644 --- a/classes/bot.py +++ b/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. diff --git a/classes/trackingmore.py b/classes/trackingmore.py new file mode 100644 index 0000000..248f03c --- /dev/null +++ b/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 \ No newline at end of file diff --git a/commands/__init__.py b/commands/__init__.py index cd3d2f9..c638982 100644 --- a/commands/__init__.py +++ b/commands/__init__.py @@ -20,6 +20,7 @@ for command in [ "privacy", "roomsettings", "dice", + "parcel", ]: function = getattr(import_module( "commands." + command), "command_" + command) diff --git a/commands/newroom.py b/commands/newroom.py index 11af1f8..169cc61 100644 --- a/commands/newroom.py +++ b/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?") \ No newline at end of file + await bot.send_message(bot.rooms[new_room.room_id], f"Welcome to the new room! What can I do for you?") \ No newline at end of file diff --git a/commands/parcel.py b/commands/parcel.py new file mode 100644 index 0000000..e85395f --- /dev/null +++ b/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) \ No newline at end of file diff --git a/config.dist.ini b/config.dist.ini index 753d9c6..e3d412a 100644 --- a/config.dist.ini +++ b/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 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 451d377..392defb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,3 +6,4 @@ duckdb python-magic pillow wolframalpha +git+https://kumig.it/kumitterer/trackingmore-api-tool.git \ No newline at end of file