initial commit

.coveragerc

@@ -0,0 +1,7 @@
+[run]
+omit =
+    setup.py
+    versioneer.py
+    dallebot/_version.py
+    dallebot/tests/*
+    */__init__.py

.gitattributes

@@ -0,0 +1 @@
+dallebot/_version.py export-subst

.gitignore

@@ -0,0 +1,4 @@
+venv_dallebot
+dallebot/env.json
+.idea
+*.csv

.pre-commit-config.yaml

@@ -0,0 +1,13 @@
+repos:
+  - repo: https://github.com/psf/black
+    rev: 22.3.0
+    hooks:
+    - id: black
+      language_version: python3
+      args: [--config, pyproject.toml]
+  - repo: https://gitlab.com/pycqa/flake8
+    rev: 3.9.2
+    hooks:
+    - id: flake8
+      language_version: python3
+      args: [--config=setup.cfg]

Dockerfile-linux

@@ -0,0 +1,10 @@
+FROM python:3.8-buster
+
+COPY requirements.txt requirements.txt
+RUN python3 -m pip install --upgrade pip
+RUN python3 -m pip install -r requirements.txt
+
+COPY . /dallebot
+RUN cd /dallebot && python3 -m pip install .
+
+WORKDIR /dallebot

Dockerfile-raspberry-pi

@@ -0,0 +1,11 @@
+FROM arm64v8/python:3.8-buster
+
+COPY pip.conf /etc/pip.conf
+COPY requirements.txt requirements.txt
+RUN python3 -m pip install --upgrade pip
+RUN python3 -m pip install -r requirements.txt
+
+COPY . /dallebot
+RUN cd /dallebot && python3 -m pip install .
+
+WORKDIR /dallebot

MANIFEST.in

@@ -0,0 +1,2 @@
+include versioneer.py
+include dallebot/_version.py

README.md

@@ -0,0 +1,56 @@
+# DALL·E Bot
+
+![screenshot](screenshot.png)
+
+This cool Telegram Bot can send images generated by OpenAI’s DALL·E.
+
+The bot is running on my Raspberry Pi and can be found here [![@dalle_mini_bot](https://img.shields.io/badge/Telegram%20Bot-@dalle_telegram_bot-blue?logo=telegram&style=plastic)](https://telegram.me/dalle_telegram_bot)
+
+## Development
+
+### Setup
+
+```shell
+source setup-local-venv.sh
+```
+
+## Deployment
+
+### Setup Environment
+
+- Create a Telegram Bot using the BotFather
+- Create a chat where the bot will send logs and errors
+- Create an OpenAI Api key [here](https://beta.openai.com/overview)
+- Create a file `env.json` in the `dallebot` subdirectory with the developer_chat_id, the bot_token, and the openai_api_key
+```json
+{
+  "developer_chat_id": "<REPLACE WITH DEVELOPER CHAT ID>",
+  "bot_token": "<REPLACE WITH BOT TOKEN>",
+  "openai_api_key": "<REPLACE WITH OPENAI API KEY>"
+}
+```
+
+### Build docker
+
+#### Raspberry Pi
+
+```shell
+./docker-build-raspberry-pi.sh
+```
+
+#### Linux/Mac
+
+```shell
+./docker-build-linux.sh
+```
+
+### Run docker
+
+```shell
+./docker-run.sh
+```
+
+## Sources
+
+- Using [OpenAI's DALL·E](https://beta.openai.com/docs/guides/images)
+- Inspired by https://github.com/python-telegram-bot/python-telegram-bot/wiki/InlineKeyboard-Example

dallebot/__init__.py

@@ -0,0 +1,3 @@
+from . import _version
+
+__version__ = _version.get_versions()["version"]

dallebot/__main__.py

@@ -0,0 +1,167 @@
+import datetime
+import html
+import json
+import logging
+import os
+import traceback
+
+import openai as openai
+import pandas as pd
+from telegram import Update, ParseMode, ChatAction
+from telegram.ext import Updater, CallbackContext, CommandHandler
+
+from tools import read_config
+
+logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+min_requests_delay = 60  # in seconds
+csv_file_name = "logs/dalle_bot_logs.csv"
+df_columns = ["group", "timestamp", "prompt", "size", "hashed_user"]
+
+config = read_config()
+developer_chat_id = config["developer_chat_id"]
+bot_token = config["bot_token"]
+openai_api_key = config["openai_api_key"]
+
+openai.api_key = openai_api_key
+
+try:
+    df = pd.read_csv(csv_file_name)
+    df = df.astype({"timestamp": "datetime64"})
+except Exception:
+    df = pd.DataFrame(columns=df_columns)
+    outdir = "logs"
+    if not os.path.exists(outdir):
+        os.mkdir(outdir)
+
+
+def start(update: Update, context: CallbackContext) -> None:
+    context.bot.send_message(
+        update.message.chat.id,
+        "Hi there! I’m DALL·E Bot.\n"
+        "Send me a prompt and I’ll send you an image generated by OpenAI’s DALL·E.\n"
+        f"Due to resource constraints it is only allowed to send "
+        f"one request per {min_requests_delay} seconds.\n"
+        f"In order to achieve this, I'm storing your anonymised hashed user id together with "
+        f"the timestamp and the prompt.",
+    )
+
+
+def generate(update: Update, context: CallbackContext, size=512) -> None:
+    """Sends a dalle image."""
+    chat_id = update.message.chat.id
+
+    context.bot.send_chat_action(chat_id=chat_id, action=ChatAction.TYPING)
+
+    prompt = " ".join(context.args)
+
+    global df
+
+    try:
+        if "group" in update.message.chat.type:
+            is_group = True
+        else:
+            is_group = False
+    except Exception as e:
+        logger.error(e)
+        is_group = False
+
+    hashed_user = hash(update.message.from_user.id)
+    datetime_now = datetime.datetime.now()
+
+    seconds_diff = (datetime_now - max(df[df.hashed_user == hashed_user]["timestamp"], default=datetime_now)).seconds
+
+    if seconds_diff < min_requests_delay:
+        context.bot.send_message(
+            chat_id,
+            f"Sorry, due to resource constraints, it's only allowed to send one request per "
+            f"{min_requests_delay} seconds.\n"
+            f"Please try again in {min_requests_delay - seconds_diff} seconds.",
+        )
+
+        return
+
+    df = pd.concat([df, pd.DataFrame([[is_group, datetime_now, prompt, size, hashed_user]], columns=df_columns)])
+    df.to_csv(csv_file_name, header=True, index=False)
+
+    if is_group:
+        is_group_text = "a group"
+    else:
+        is_group_text = "a single user"
+    context.bot.send_message(developer_chat_id, f"Sending a dalle image to {is_group_text}: {prompt}")
+
+    num_of_max_tries = 5
+    num_of_tries = 1
+    success = False
+
+    while num_of_tries <= num_of_max_tries and not success:
+        try:
+            num_of_tries += 1
+
+            moderation_response = openai.Moderation.create(prompt)
+
+            if not moderation_response["results"][0]["flagged"]:
+                response = openai.Image.create(prompt=prompt, n=1, size=f"{size}x{size}", user=str(hashed_user))
+                image_url = response["data"][0]["url"]
+
+                context.bot.send_photo(chat_id, image_url, caption=prompt)
+            else:
+                context.bot.send_message("This prompt doesn't comply with OpenAI's content policy.")
+
+            success = True
+        except Exception as e:
+            if num_of_tries == num_of_max_tries:
+                raise e
+            else:
+                logger.error(e)
+
+
+def error_handler(update: object, context: CallbackContext) -> None:
+    """Log the error and send a telegram message to notify the developer."""
+    # Log the error before we do anything else, so we can see it even if something breaks.
+    logger.error(msg="Exception while handling an update:", exc_info=context.error)
+
+    # traceback.format_exception returns the usual python message about an exception, but as a
+    # list of strings rather than a single string, so we have to join them together.
+    tb_list = traceback.format_exception(None, context.error, context.error.__traceback__)
+    tb_string = "".join(tb_list)
+
+    # Build the message with some markup and additional information about what happened.
+    # You might need to add some logic to deal with messages longer than the 4096 character limit.
+    update_str = update.to_dict() if isinstance(update, Update) else str(update)
+    message = (
+        f"An exception was raised while handling an update\n"
+        f"<pre>update = {html.escape(json.dumps(update_str, indent=2, ensure_ascii=False))}"
+        "</pre>\n\n"
+        f"<pre>context.chat_data = {html.escape(str(context.chat_data))}</pre>\n\n"
+        f"<pre>context.user_data = {html.escape(str(context.user_data))}</pre>\n\n"
+        f"<pre>{html.escape(tb_string)}"
+    )
+
+    message = message[:4090] + "</pre>"
+
+    # Finally, send the message
+    context.bot.send_message(chat_id=developer_chat_id, text=message, parse_mode=ParseMode.HTML)
+
+
+def main() -> None:
+    """Setup and run the bot."""
+    # Create the Updater and pass it your bot's token.
+    updater = Updater(bot_token)
+
+    updater.dispatcher.add_handler(CommandHandler("generate", generate, pass_args=True))
+    updater.dispatcher.add_handler(CommandHandler("start", start))
+
+    updater.dispatcher.add_error_handler(error_handler)
+
+    # Start the Bot
+    updater.start_polling(poll_interval=1)
+
+    # Run the bot until the user presses Ctrl-C or the process receives SIGINT,
+    # SIGTERM or SIGABRT
+    updater.idle()
+
+
+if __name__ == "__main__":
+    main()