Preparation for automatic message classification

classes/bot.py

@@ -21,7 +21,8 @@ from nio import (
     EncryptionError,
     RoomMessageText,
     RoomSendResponse,
-    SyncResponse
+    SyncResponse,
+    RoomMessageNotice
 )
 from nio.crypto import Olm
 
@@ -55,6 +56,8 @@ class GPTBot:
     logger: Optional[Logger] = Logger()
     chat_api: Optional[OpenAI] = None
     image_api: Optional[OpenAI] = None
+    classification_api: Optional[OpenAI] = None
+    operator: Optional[str] = None
 
     @classmethod
     def from_config(cls, config: ConfigParser):
@@ -76,6 +79,7 @@ class GPTBot:
 
         # Override default values
         if "GPTBot" in config:
+            bot.operator = config["GPTBot"].get("Operator", bot.operator)
             bot.default_room_name = config["GPTBot"].get(
                 "DefaultRoomName", bot.default_room_name)
             bot.default_system_message = config["GPTBot"].get(
@@ -83,7 +87,8 @@ class GPTBot:
             bot.force_system_message = config["GPTBot"].getboolean(
                 "ForceSystemMessage", bot.force_system_message)
 
-        bot.chat_api = bot.image_api = OpenAI(config["OpenAI"]["APIKey"], config["OpenAI"].get("Model"), bot.logger)
+        bot.chat_api = bot.image_api = bot.classification_api = OpenAI(
+            config["OpenAI"]["APIKey"], config["OpenAI"].get("Model"), bot.logger)
         bot.max_tokens = config["OpenAI"].getint("MaxTokens", bot.max_tokens)
         bot.max_messages = config["OpenAI"].getint(
             "MaxMessages", bot.max_messages)
@@ -158,10 +163,10 @@ class GPTBot:
                     self.logger.log(
                         f"Could not decrypt message {event.event_id} in room {room_id}", "error")
                     continue
-            if isinstance(event, RoomMessageText):
+            if isinstance(event, (RoomMessageText, RoomMessageNotice)):
                 if event.body.startswith("!gptbot ignoreolder"):
                     break
-                if not event.body.startswith("!"):
+                if (not event.body.startswith("!")) or (event.body.startswith("!gptbot")):
                     messages.append(event)
 
         self.logger.log(f"Found {len(messages)} messages (limit: {n})")
@@ -230,14 +235,15 @@ class GPTBot:
 
         await COMMANDS.get(command, COMMANDS[None])(room, event, self)
 
-    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():
                 if isinstance(event, eventtype):
                     await callback(room, event, self)
         except Exception as e:
-            self.logger.log(f"Error in event callback for {event.__class__}: {e}", "error")
+            self.logger.log(
+                f"Error in event callback for {event.__class__}: {e}", "error")
 
     async def response_callback(self, response: Response):
         for response_type, callback in RESPONSE_CALLBACKS.items():
@@ -347,6 +353,30 @@ class GPTBot:
 
         return await self.matrix_client._send(RoomSendResponse, method, path, data, (room.room_id,))
 
+    def log_api_usage(self, message: Event | str, room: MatrixRoom | int, api: str, tokens: int):
+        """Log API usage to the database.
+
+        Args:
+            message (Event): The event that triggered the API usage.
+            room (MatrixRoom | int): The room the event was sent in.
+            api (str): The API that was used.
+            tokens (int): The number of tokens used.
+        """
+
+        if not self.database:
+            return
+
+        if isinstance(message, Event):
+            message = message.event_id
+
+        if isinstance(room, MatrixRoom):
+            room = room.room_id
+
+        self.database.execute(
+            "INSERT INTO token_usage (message_id, room_id, tokens, api, timestamp) VALUES (?, ?, ?, ?, ?)",
+            (message, room, tokens, api, datetime.now())
+        )
+
     async def run(self):
         """Start the bot."""
 
@@ -410,7 +440,8 @@ class GPTBot:
         # Set up callbacks
 
         self.matrix_client.add_event_callback(self.event_callback, Event)
-        self.matrix_client.add_response_callback(self.response_callback, Response)
+        self.matrix_client.add_response_callback(
+            self.response_callback, Response)
 
         # Accept pending invites
 
@@ -454,7 +485,8 @@ class GPTBot:
             chat_messages, self.max_tokens - 1, system_message=system_message)
 
         try:
-            response, tokens_used = await self.generate_chat_response(truncated_messages)
+            response, tokens_used = self.chat_api.generate_chat_response(
+                chat_messages, user=room.room_id)
         except Exception as e:
             self.logger.log(f"Error generating response: {e}", "error")
             await self.send_message(
@@ -462,40 +494,22 @@ class GPTBot:
             return
 
         if response:
+            self.log_api_usage(event, room, f"{self.chat_api.api_code}-{self.chat_api.chat_api}", tokens_used)
+
             self.logger.log(f"Sending response to room {room.room_id}...")
 
             # Convert markdown to HTML
 
             message = await self.send_message(room, response)
 
-            if self.database:
-                self.logger.log("Storing record of tokens used...")
-
-                with self.database.cursor() as cursor:
-                    cursor.execute(
-                        "INSERT INTO token_usage (message_id, room_id, tokens, timestamp) VALUES (?, ?, ?, ?)",
-                        (message.event_id, room.room_id, tokens_used, datetime.now()))
-                    self.database.commit()
         else:
             # Send a notice to the room if there was an error
             self.logger.log("Didn't get a response from GPT API", "error")
-            send_message(
+            await send_message(
                 room, "Something went wrong. Please try again.", True)
 
         await self.matrix_client.room_typing(room.room_id, False)
 
-    async def generate_chat_response(self, messages: List[Dict[str, str]]) -> Tuple[str, int]:
-        """Generate a response to a chat message.
-
-        Args:
-            messages (List[Dict[str, str]]): A list of messages to use as context.
-
-        Returns:
-            Tuple[str, int]: The response text and the number of tokens used.
-        """
-
-        return self.chat_api.generate_chat_response(messages)
-
     def get_system_message(self, room: MatrixRoom | int) -> str:
         default = self.default_system_message
 

classes/openai.py

@@ -1,21 +1,34 @@
 import openai
 import requests
 
+import json
+
 from .logging import Logger
 
-from typing import Dict, List, Tuple, Generator
+from typing import Dict, List, Tuple, Generator, Optional
 
 class OpenAI:
     api_key: str
     chat_model: str = "gpt-3.5-turbo"
     logger: Logger
 
+    api_code: str = "openai"
+
+    @property
+    def chat_api(self) -> str:
+        return self.chat_model
+    
+    classification_api = chat_api
+    image_api: str = "dalle"
+
+    operator: str = "OpenAI ([https://openai.com](https://openai.com))"
+
     def __init__(self, api_key, chat_model=None, logger=None):
         self.api_key = api_key
         self.chat_model = chat_model or self.chat_model
         self.logger = logger or Logger()
 
-    def generate_chat_response(self, messages: List[Dict[str, str]]) -> Tuple[str, int]:
+    def generate_chat_response(self, messages: List[Dict[str, str]], user: Optional[str] = None) -> Tuple[str, int]:
         """Generate a response to a chat message.
 
         Args:
@@ -30,7 +43,8 @@ class OpenAI:
         response = openai.ChatCompletion.create(
             model=self.chat_model,
             messages=messages,
-            api_key=self.api_key
+            api_key=self.api_key,
+            user = user
         )
 
         result_text = response.choices[0].message['content']
@@ -38,7 +52,50 @@ class OpenAI:
         self.logger.log(f"Generated response with {tokens_used} tokens.")
         return result_text, tokens_used
 
-    def generate_image(self, prompt: str) -> Generator[bytes, None, None]:
+    def classify_message(self, query: str, user: Optional[str] = None) -> Tuple[Dict[str, str], int]:
+        system_message = """You are a classifier for different types of messages. You decide whether an incoming message is meant to be a prompt for an AI chat model, an image generation AI, or a calculation for WolframAlpha. You respond with a JSON object like this:
+
+{ "type": event_type, "prompt": prompt }
+
+- If the message you received is meant for the AI chat model, the event_type is "chat", and the prompt is the literal content of the message you received. This is also the default if none of the other options apply.
+- If it is a prompt for a calculation that can be answered better by WolframAlpha than an AI chat bot, the event_type is "calculate". Optimize the message you received for input to WolframAlpha, and return it as the prompt attribute.
+- If it is a prompt for an AI image generation, the event_type is "imagine". Optimize the message you received for use with DALL-E, and return it as the prompt attribute.
+- If for any reason you are unable to classify the message (for example, if it infringes on your terms of service), the event_type is "error", and the prompt is a message explaining why you are unable to process the message.
+
+Only the event_types mentioned above are allowed, you must not respond in any other way."""
+
+        messages = [
+            {
+                "role": "system", 
+                "content": system_message
+            },
+            {
+                "role": "user",
+                "content": query
+            }
+        ]
+
+        self.logger.log(f"Classifying message '{query}'...")
+
+        response = openai.ChatCompletion.create(
+            model=self.chat_model,
+            messages=messages,
+            api_key=self.api_key,
+            user = user
+        )
+
+        try:
+            result = json.loads(response.choices[0].message['content'])
+        except:
+            result = {"type": "chat", "prompt": query}
+
+        tokens_used = response.usage["total_tokens"]
+
+        self.logger.log(f"Classified message as {result['type']} with {tokens_used} tokens.")
+
+        return result, tokens_used
+
+    def generate_image(self, prompt: str, user: Optional[str] = None) -> Generator[bytes, None, None]:
         """Generate an image from a prompt.
 
         Args:
@@ -54,9 +111,14 @@ class OpenAI:
             prompt=prompt,
             n=1,
             api_key=self.api_key,
-            size="1024x1024"
+            size="1024x1024",
+            user = user
         )
 
+        images = []
+
         for image in response.data:
             image = requests.get(image.url).content
-            yield image
+            images.append(image)
+
+        return images, len(images)

classes/wolframalpha.py

@@ -10,17 +10,33 @@ class WolframAlpha:
     logger: Logger
     client: wolframalpha.Client
 
+    api_code: str = "wolfram"
+    calculation_api: str = "alpha"
+
+    operator: str = "Wolfram ([https://www.wolfram.com](https://www.wolfram.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 = wolframalpha.Client(self.api_key)
 
-    def generate_calculation_response(self, query: str, text: Optional[bool] = False, results_only: Optional[bool] = False) -> Generator[str | bytes, None, None]:
+    def generate_calculation_response(self, query: str, text: Optional[bool] = False, results_only: Optional[bool] = False, user: Optional[str] = None) -> Generator[str | bytes, None, None]:
         self.logger.log(f"Querying WolframAlpha for {query}")
 
         response: wolframalpha.Result = self.client.query(query)
 
-        for pod in response.pods if not results_only else response.results:
+        if not response.success:
+            yield "Could not process your query."
+
+            if response.didyoumeans:
+                yield "Did you mean: " + response.didyoumeans["didyoumean"][0]["#text"]
+
+            return
+
+        if response.error:
+            self.logger.log("Error in query to WolframAlpha: " + response.error, "error")
+
+        for pod in response.pods if not results_only else (response.results if len(list(response.results)) else response.pods):
             self.logger.log(pod.keys(), "debug")
             if pod.title:
                 yield "*" + pod.title + "*"

commands/__init__.py

@@ -8,6 +8,8 @@ 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
 
 COMMANDS = {
     "help": command_help,
@@ -19,5 +21,7 @@ COMMANDS = {
     "systemmessage": command_systemmessage,
     "imagine": command_imagine,
     "calculate": command_calculate,
+    "classify": command_classify,
+    "chat": command_chat,
     None: command_unknown,
 }

commands/calculate.py

@@ -20,15 +20,17 @@ async def command_calculate(room: MatrixRoom, event: RoomMessageText, bot):
     prompt = " ".join(prompt)
 
     if prompt:
-        bot.logger.log("Querying WolframAlpha")
+        bot.logger.log("Querying calculation API...")
 
-        for subpod in bot.calculation_api.generate_calculation_response(prompt, text, results_only):
+        for subpod in bot.calculation_api.generate_calculation_response(prompt, text, results_only, user=room.room_id):
             bot.logger.log(f"Sending subpod...")
             if isinstance(subpod, bytes):
                 await bot.send_image(room, subpod)
             else:
                 await bot.send_message(room, subpod, True)
 
+        bot.log_api_usage(event, room, f"{self.calculation_api.api_code}-{self.calculation_api.calculation_api}", tokens_used)
+
         return
 
     await bot.send_message(room, "You need to provide a prompt.", True)

commands/chat.py

@@ -0,0 +1,15 @@
+from nio.events.room_events import RoomMessageText
+from nio.rooms import MatrixRoom
+
+
+async def command_chat(room: MatrixRoom, event: RoomMessageText, bot):
+    prompt = " ".join(event.body.split()[2:])
+
+    if prompt:
+        bot.logger.log("Sending chat message...")
+        event.body = prompt
+        await bot.process_query(room, event)
+
+        return
+
+    await bot.send_message(room, "You need to provide a prompt.", True)

commands/classify.py

@@ -0,0 +1,24 @@
+from nio.events.room_events import RoomMessageText
+from nio.rooms import MatrixRoom
+
+
+async def command_classify(room: MatrixRoom, event: RoomMessageText, bot):
+    prompt = " ".join(event.body.split()[2:])
+
+    if prompt:
+        bot.logger.log("Classifying message...")
+
+        response, tokens_used = bot.classification_api.classify_message(prompt, user=room.room_id)
+
+        message = f"The message you provided seems to be of type: {response['type']}."
+
+        if not prompt == response["prompt"]:
+            message += f"\n\nPrompt: {response['prompt']}."
+
+        await bot.send_message(room, message, True)
+
+        bot.log_api_usage(event, room, f"{self.classification_api.api_code}-{self.classification_api.classification_api}", tokens_used)
+
+        return
+
+    await bot.send_message(room, "You need to provide a prompt.", True)

commands/help.py

@@ -14,6 +14,9 @@ async def command_help(room: MatrixRoom, event: RoomMessageText, bot):
 - !gptbot systemmessage \<message\> - Get or set the system message for this room
 - !gptbot imagine \<prompt\> - Generate an image from a prompt
 - !gptbot calculate [--text] [--details] \<query\> - Calculate a result to a calculation, optionally forcing text output instead of an image, and optionally showing additional details like the input interpretation
+- !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
 """
 
     await bot.send_message(room, body, True)

commands/imagine.py

@@ -8,10 +8,14 @@ async def command_imagine(room: MatrixRoom, event: RoomMessageText, bot):
     if prompt:
         bot.logger.log("Generating image...")
 
-        for image in bot.image_api.generate_image(prompt):
+        images, tokens_used = bot.image_api.generate_image(prompt, user=room.room_id)
+
+        for image in images:
             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)
+
         return
 
     await bot.send_message(room, "You need to provide a prompt.", True)

commands/privacy.py

@@ -0,0 +1,17 @@
+from nio.events.room_events import RoomMessageText
+from nio.rooms import MatrixRoom
+
+
+async def command_privacy(room: MatrixRoom, event: RoomMessageText, bot):
+    body = "**Privacy**\n\nIf you use this bot, note that your messages will be sent to the following recipients:\n\n"
+
+    body += "- The bot's operator" + (f"({bot.operator})" if bot.operator else "") + "\n"
+
+    if bot.chat_api:
+        body += "- For chat requests: " + f"{bot.chat_api.operator}" + "\n"
+    if bot.image_api:
+        body += "- For image generation requests (!gptbot imagine): " + f"{bot.image_api.operator}" + "\n"
+    if bot.calculate_api:
+        body += "- For calculation requests (!gptbot calculate): " + f"{bot.calculate_api.operator}" + "\n"
+
+    await bot.send_message(room, body, True)

commands/systemmessage.py

@@ -15,7 +15,7 @@ async def command_systemmessage(room: MatrixRoom, event: RoomMessageText, bot):
                  system_message, event.server_timestamp)
             )
 
-        bot.send_message(room, f"Alright, I've stored the system message: '{system_message}'.", True)
+        await bot.send_message(room, f"Alright, I've stored the system message: '{system_message}'.", True)
         return
 
     bot.logger.log("Retrieving system message...")

config.dist.ini

@@ -68,6 +68,11 @@ AccessToken = syt_yoursynapsetoken
 
 [GPTBot]
 
+# Some way for the user to contact you.
+# Ideally, either your personal user ID or a support room
+#
+Operator = Contact details not set
+
 # The default room name used by the !newroom command
 # Defaults to GPTBot if not set
 #

migrations/__init__.py

@@ -6,12 +6,14 @@ 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
 
 MIGRATIONS = OrderedDict()
 
 MIGRATIONS[1] = migration_1
 MIGRATIONS[2] = migration_2
 MIGRATIONS[3] = migration_3
+MIGRATIONS[4] = migration_4
 
 def get_version(db: DuckDBPyConnection) -> int:
     """Get the current database version.

migrations/migration_4.py

@@ -0,0 +1,18 @@
+# Migration to add API column to token usage table
+
+from datetime import datetime
+
+def migration(conn):
+    with conn.cursor() as cursor:
+        cursor.execute(
+            """
+            ALTER TABLE token_usage ADD COLUMN api TEXT DEFAULT 'openai'
+            """
+        )
+
+        cursor.execute(
+            "INSERT INTO migrations (id, timestamp) VALUES (4, ?)",
+            (datetime.now(),)
+        )
+
+        conn.commit()