Added support for DALL-E and WolframAlpha

New "imagine" and "calculate" commands
Implemented image sending
Moved OpenAI specific code to OpenAI class
Abstracted away OpenAI API in bot class
Minor fixes

README.md

@@ -1,13 +1,23 @@
 # GPTbot
 
-GPTbot is a simple bot that uses the [OpenAI Chat Completion API](https://platform.openai.com/docs/guides/chat)
-to generate responses to messages in a Matrix room.
+GPTbot is a simple bot that uses different APIs to generate responses to 
+messages in a Matrix room.
 
-It will also save a log of the spent tokens to a DuckDB database
-(database.db in the working directory, by default).
+It is called GPTbot because it was originally intended to only use GPT-3 to
+generate responses. However, it supports other services/APIs, and I will 
+probably add more in the future, so the name is a bit misleading.
 
-Note that this bot does not yet support encryption - this is still work in
-progress.
+## Features
+
+- AI-generated responses to all messages in a Matrix room (chatbot)
+  - Currently supports OpenAI (tested with `gpt-3.5-turbo`)
+- AI-generated pictures via the `!gptbot imagine` command
+  - Currently supports OpenAI (DALL-E)
+- DuckDB database to store spent tokens
+
+## Planned features
+
+- End-to-end encryption support (partly implemented, but not yet working)
 
 ## Installation
 

classes/bot.py

@@ -1,10 +1,11 @@
-import openai
 import markdown2
 import duckdb
 import tiktoken
-
+import magic
 import asyncio
 
+from PIL import Image
+
 from nio import (
     AsyncClient,
     AsyncClientConfig,
@@ -24,9 +25,10 @@ from nio import (
 )
 from nio.crypto import Olm
 
-from typing import Optional, List, Dict
+from typing import Optional, List, Dict, Tuple
 from configparser import ConfigParser
 from datetime import datetime
+from io import BytesIO
 
 import uuid
 
@@ -35,6 +37,8 @@ from migrations import migrate
 from callbacks import RESPONSE_CALLBACKS, EVENT_CALLBACKS
 from commands import COMMANDS
 from .store import DuckDBStore
+from .openai import OpenAI
+from .wolframalpha import WolframAlpha
 
 
 class GPTBot:
@@ -46,11 +50,11 @@ class GPTBot:
     force_system_message: bool = False
     max_tokens: int = 3000  # Maximum number of input tokens
     max_messages: int = 30  # Maximum number of messages to consider as input
-    model: str = "gpt-3.5-turbo"  # OpenAI chat model to use
     matrix_client: Optional[AsyncClient] = None
     sync_token: Optional[str] = None
     logger: Optional[Logger] = Logger()
-    openai_api_key: Optional[str] = None
+    chat_api: Optional[OpenAI] = None
+    image_api: Optional[OpenAI] = None
 
     @classmethod
     def from_config(cls, config: ConfigParser):
@@ -79,12 +83,15 @@ 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.max_tokens = config["OpenAI"].getint("MaxTokens", bot.max_tokens)
         bot.max_messages = config["OpenAI"].getint(
             "MaxMessages", bot.max_messages)
-        bot.model = config["OpenAI"].get("Model", bot.model)
 
-        bot.openai_api_key = config["OpenAI"]["APIKey"]
+        # Set up WolframAlpha
+        if "WolframAlpha" in config:
+            bot.calculation_api = WolframAlpha(
+                config["WolframAlpha"]["APIKey"], bot.logger)
 
         # Set up the Matrix client
 
@@ -165,7 +172,7 @@ class GPTBot:
     def _truncate(self, messages: list, max_tokens: Optional[int] = None,
                   model: Optional[str] = None, system_message: Optional[str] = None):
         max_tokens = max_tokens or self.max_tokens
-        model = model or self.model
+        model = model or self.chat_api.chat_model
         system_message = self.default_system_message if system_message is None else system_message
 
         encoding = tiktoken.encoding_for_model(model)
@@ -224,10 +231,13 @@ class GPTBot:
         await COMMANDS.get(command, COMMANDS[None])(room, event, self)
 
     async def event_callback(self,room: MatrixRoom, event: Event):
-        self.logger.log("Received event: " + str(event), "debug")
+        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")
 
     async def response_callback(self, response: Response):
         for response_type, callback in RESPONSE_CALLBACKS.items():
@@ -244,6 +254,52 @@ class GPTBot:
         for invite in invites.keys():
             await self.matrix_client.join(invite)
 
+    async def send_image(self, room: MatrixRoom, image: bytes, message: Optional[str] = None):
+        self.logger.log(
+            f"Sending image of size {len(image)} bytes to room {room.room_id}")
+
+        bio = BytesIO(image)
+        img = Image.open(bio)
+        mime = Image.MIME[img.format]
+
+        (width, height) = img.size
+
+        self.logger.log(
+            f"Uploading - Image size: {width}x{height} pixels, MIME type: {mime}")
+
+        bio.seek(0)
+
+        response, _ = await self.matrix_client.upload(
+            bio,
+            content_type=mime,
+            filename="image",
+            filesize=len(image)
+        )
+
+        self.logger.log("Uploaded image - sending message...")
+
+        content = {
+            "body": message or "",
+            "info": {
+                "mimetype": mime,
+                "size": len(image),
+                "w": width,
+                "h": height,
+            },
+            "msgtype": "m.image",
+            "url": response.content_uri
+        }
+
+        status = await self.matrix_client.room_send(
+            room.room_id,
+            "m.room.message",
+            content
+        )
+
+        self.logger.log(str(status), "debug")
+
+        self.logger.log("Sent image")
+
     async def send_message(self, room: MatrixRoom, message: str, notice: bool = False):
         markdowner = markdown2.Markdown(extras=["fenced-code-blocks"])
         formatted_body = markdowner.convert(message)
@@ -428,28 +484,17 @@ class GPTBot:
 
         await self.matrix_client.room_typing(room.room_id, False)
 
-    async def generate_chat_response(self, messages: List[Dict[str, str]]) -> str:
+    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:
-            str: The response to the chat.
+            Tuple[str, int]: The response text and the number of tokens used.
         """
 
-        self.logger.log(f"Generating response to {len(messages)} messages...")
-
-        response = openai.ChatCompletion.create(
-            model=self.model,
-            messages=messages,
-            api_key=self.openai_api_key
-        )
-
-        result_text = response.choices[0].message['content']
-        tokens_used = response.usage["total_tokens"]
-        self.logger.log(f"Generated response with {tokens_used} tokens.")
-        return result_text, 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

@@ -0,0 +1,62 @@
+import openai
+import requests
+
+from .logging import Logger
+
+from typing import Dict, List, Tuple, Generator
+
+class OpenAI:
+    api_key: str
+    chat_model: str = "gpt-3.5-turbo"
+    logger: Logger
+
+    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]:
+        """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.
+        """
+
+        self.logger.log(f"Generating response to {len(messages)} messages using {self.chat_model}...")
+
+        response = openai.ChatCompletion.create(
+            model=self.chat_model,
+            messages=messages,
+            api_key=self.api_key
+        )
+
+        result_text = response.choices[0].message['content']
+        tokens_used = response.usage["total_tokens"]
+        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]:
+        """Generate an image from a prompt.
+
+        Args:
+            prompt (str): The prompt to use.
+
+        Yields:
+            bytes: The image data.
+        """
+
+        self.logger.log(f"Generating image from prompt '{prompt}'...")
+
+        response = openai.Image.create(
+            prompt=prompt,
+            n=1,
+            api_key=self.api_key,
+            size="1024x1024"
+        )
+
+        for image in response.data:
+            image = requests.get(image.url).content
+            yield image

classes/wolframalpha.py

@@ -0,0 +1,35 @@
+import wolframalpha
+import requests
+
+from .logging import Logger
+
+from typing import Dict, List, Tuple, Generator, Optional
+
+class WolframAlpha:
+    api_key: str
+    logger: Logger
+    client: wolframalpha.Client
+
+    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]:
+        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:
+            self.logger.log(pod.keys(), "debug")
+            if pod.title:
+                yield "*" + pod.title + "*"
+            for subpod in pod.subpods:
+                self.logger.log(subpod.keys(), "debug")
+                if subpod.title:
+                    yield "*" + subpod.title + "*"
+                if subpod.img and not text:
+                    image = requests.get(subpod.img.src).content
+                    yield image
+                elif subpod.plaintext:
+                    yield "```\n" + subpod.plaintext + "\n```"

commands/__init__.py

@@ -6,6 +6,8 @@ from .unknown import command_unknown
 from .coin import command_coin
 from .ignoreolder import command_ignoreolder
 from .systemmessage import command_systemmessage
+from .imagine import command_imagine
+from .calculate import command_calculate
 
 COMMANDS = {
     "help": command_help,
@@ -15,5 +17,7 @@ COMMANDS = {
     "coin": command_coin,
     "ignoreolder": command_ignoreolder,
     "systemmessage": command_systemmessage,
+    "imagine": command_imagine,
+    "calculate": command_calculate,
     None: command_unknown,
 }

commands/calculate.py

@@ -0,0 +1,34 @@
+from nio.events.room_events import RoomMessageText
+from nio.rooms import MatrixRoom
+
+
+async def command_calculate(room: MatrixRoom, event: RoomMessageText, bot):
+    prompt = event.body.split()[2:]
+    text = False
+    results_only = True
+
+    if "--text" in prompt:
+        text = True
+        delete = prompt.index("--text")
+        del prompt[delete]
+
+    if "--details" in prompt:
+        results_only = False
+        delete = prompt.index("--details")
+        del prompt[delete]
+
+    prompt = " ".join(prompt)
+
+    if prompt:
+        bot.logger.log("Querying WolframAlpha")
+
+        for subpod in bot.calculation_api.generate_calculation_response(prompt, text, results_only):
+            bot.logger.log(f"Sending subpod...")
+            if isinstance(subpod, bytes):
+                await bot.send_image(room, subpod)
+            else:
+                await bot.send_message(room, subpod, True)
+
+        return
+
+    await bot.send_message(room, "You need to provide a prompt.", True)

commands/help.py

@@ -5,13 +5,15 @@ from nio.rooms import MatrixRoom
 async def command_help(room: MatrixRoom, event: RoomMessageText, bot):
     body = """Available commands:
 
-!gptbot help - Show this message
-!gptbot newroom <room name> - Create a new room and invite yourself to it
-!gptbot stats - Show usage statistics for this room
-!gptbot botinfo - Show information about the bot
-!gptbot coin - Flip a coin (heads or tails)
-!gptbot ignoreolder - Ignore messages before this point as context
-!gptbot systemmessage - Get or set the system message for this room
+- !gptbot help - Show this message
+- !gptbot newroom \<room name\> - Create a new room and invite yourself to it
+- !gptbot stats - Show usage statistics for this room
+- !gptbot botinfo - Show information about the bot
+- !gptbot coin - Flip a coin (heads or tails)
+- !gptbot ignoreolder - Ignore messages before this point as context
+- !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
 """
 
     await bot.send_message(room, body, True)

commands/imagine.py

@@ -0,0 +1,17 @@
+from nio.events.room_events import RoomMessageText
+from nio.rooms import MatrixRoom
+
+
+async def command_imagine(room: MatrixRoom, event: RoomMessageText, bot):
+    prompt = " ".join(event.body.split()[2:])
+
+    if prompt:
+        bot.logger.log("Generating image...")
+
+        for image in bot.image_api.generate_image(prompt):
+            bot.logger.log(f"Sending image...")
+            await bot.send_image(room, image)
+
+        return
+
+    await bot.send_message(room, "You need to provide a prompt.", True)

commands/newroom.py

@@ -23,7 +23,7 @@ async def command_newroom(room: MatrixRoom, event: RoomMessageText, bot):
         await bot.send_message(room, f"Sorry, I was unable to invite you to the new room. Please try again later, or create a room manually.", True)
         return
 
-    await context["client"].room_put_state(
+    await bot.matrix_client.room_put_state(
         new_room.room_id, "m.room.power_levels", {"users": {event.sender: 100}})
 
     await bot.matrix_client.joined_rooms()

commands/unknown.py

@@ -2,7 +2,7 @@ from nio.events.room_events import RoomMessageText
 from nio.rooms import MatrixRoom
 
 
-async def command_unknown(room: MatrixRoom, event: RoomMessageText, context: dict):
+async def command_unknown(room: MatrixRoom, event: RoomMessageText, bot):
     bot.logger.log("Unknown command")
 
-    bot.send_message(room, "Unknown command - try !gptbot help", True)
+    await bot.send_message(room, "Unknown command - try !gptbot help", True)

config.dist.ini

@@ -38,6 +38,15 @@ APIKey = sk-yoursecretkey
 #
 # MaxMessages = 20
 
+[WolframAlpha]
+
+# An API key for Wolfram|Alpha
+# Request one at https://developer.wolframalpha.com
+#
+# Leave unset to disable Wolfram|Alpha integration (`!gptbot calculate`)
+#
+#APIKey = YOUR-APIKEY
+
 [Matrix]
 
 # The URL to your Matrix homeserver

requirements.txt

@@ -3,3 +3,6 @@ matrix-nio[e2e]
 markdown2[all]
 tiktoken
 duckdb
+python-magic
+pillow
+wolframalpha