├── .gitattributes ├── .github └── workflows │ ├── repo-tag_on_pr_merge.yaml │ ├── repo-update_license_year.yaml │ └── repo-verify_pr_release_tag.yaml ├── .gitignore ├── LICENSE.md ├── README.md ├── __main__.py ├── dalle_telegram_bot ├── __init__.py ├── entrypoint.py ├── logger.py ├── services │ ├── __init__.py │ ├── bot │ │ ├── __init__.py │ │ ├── bot.py │ │ ├── chatactions.py │ │ ├── constants.py │ │ ├── middlewares.py │ │ └── requester.py │ ├── dalle │ │ ├── __init__.py │ │ ├── dalle.py │ │ ├── exceptions.py │ │ └── models.py │ ├── logger_abc.py │ └── redis.py ├── settings.py └── utils.py ├── docs ├── Telegram-DalleMiniBot-screenshot.png └── bot-logo.jpg ├── requirements.txt └── sample.env /.gitattributes: -------------------------------------------------------------------------------- 1 | .github/* linguist-vendored 2 | -------------------------------------------------------------------------------- /.github/workflows/repo-tag_on_pr_merge.yaml: -------------------------------------------------------------------------------- 1 | name: "Tag Version on Pull Request merge" 2 | on: 3 | pull_request: 4 | types: 5 | - closed 6 | branches: 7 | - main 8 | 9 | jobs: 10 | TagOnPR: 11 | name: Tag on Pull Request merge 12 | runs-on: ubuntu-latest 13 | # Reference: https://github.com/David-Lor/action-tag-on-pr-merge/blob/develop/.github/workflows/tag_release.yaml 14 | if: github.event.pull_request.merged == true 15 | steps: 16 | - name: Tag on PR merge 17 | id: tag-on-pr-merge 18 | uses: David-Lor/action-tag-on-pr-merge@main 19 | with: 20 | push-tag: true 21 | - name: Print fetched tag 22 | run: echo "${{ steps.tag-on-pr-merge.outputs.tag }}" 23 | -------------------------------------------------------------------------------- /.github/workflows/repo-update_license_year.yaml: -------------------------------------------------------------------------------- 1 | name: Update year in license file 2 | on: 3 | workflow_dispatch: {} 4 | schedule: 5 | - cron: "0 3 1 1 *" # January 1st, 3:00 AM 6 | 7 | jobs: 8 | run: 9 | runs-on: ubuntu-latest 10 | steps: 11 | - name: Checkout 12 | uses: actions/checkout@v2 13 | with: 14 | fetch-depth: 0 15 | - name: Update year in license file 16 | uses: FantasticFiasco/action-update-license-year@e4432857e61361d140b4b5b02ece051ebfaa195b 17 | with: 18 | token: ${{ secrets.GITHUB_TOKEN }} 19 | path: LICENSE.md 20 | -------------------------------------------------------------------------------- /.github/workflows/repo-verify_pr_release_tag.yaml: -------------------------------------------------------------------------------- 1 | # When a PR to 'main' is opened/edited, validate that it has a valid version tag on its description, 2 | # that will create a tag when merged, using the "repo-tag_on_pr_merge" workflow. 3 | 4 | name: Verify PR Version Release Tag 5 | on: 6 | pull_request: 7 | types: 8 | - opened 9 | - edited 10 | branches: 11 | - main 12 | 13 | jobs: 14 | Verify: 15 | name: Verify PR Version Release Tag 16 | runs-on: ubuntu-latest 17 | env: 18 | NOTAG_COMMENT: | 19 | :warning: **No valid version tag detected!** :warning: 20 | Pull Requests to main **must** include a line on their body like `Tags x.y.z`, for tagging the version when merging. 21 | 22 | steps: 23 | - name: Extract tag from PR 24 | id: extract-tag 25 | uses: David-Lor/action-tag-on-pr-merge@feature/validation 26 | with: 27 | push-tag: false 28 | pr-sha: ${{ github.event.pull_request.head.sha }} 29 | 30 | - name: Find latest old fail comment 31 | id: last-fail-comment 32 | continue-on-error: true 33 | uses: sandeshjangam/comment-actions@v1 34 | with: 35 | type: find 36 | number: ${{ github.event.pull_request.number }} 37 | search_term: ${{ env.NOTAG_COMMENT }} 38 | direction: newer 39 | 40 | - name: Delete latest old fail comment 41 | if: ${{ steps.last-fail-comment.outputs.comment_id != '' }} 42 | continue-on-error: true 43 | uses: sandeshjangam/comment-actions@v1 44 | with: 45 | type: delete 46 | comment_id: ${{ steps.last-fail-comment.outputs.comment_id }} 47 | 48 | - name: Comment if not tag supplied 49 | if: ${{ steps.extract-tag.outputs.tag == '' }} 50 | continue-on-error: true 51 | uses: sandeshjangam/comment-actions@v1 52 | with: 53 | type: create 54 | number: ${{ github.event.pull_request.number }} 55 | body: ${{ env.NOTAG_COMMENT }} 56 | 57 | - name: Verify detected tag 58 | run: | 59 | TAG="${{ steps.extract-tag.outputs.tag }}" 60 | test "$TAG" && echo "::notice title=Detected tag::$TAG" || { echo "::error title=No detected tag::Pull Request body does not have a Tag"; exit 1; } 61 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .idea/ 2 | .vscode/ 3 | .env 4 | *pycache* 5 | .pytest_cache 6 | *.pyc 7 | local/ 8 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | Copyright 2022-2023 David Lorenzo @ github.com/David-Lor 2 | 3 | Permission to use, copy, modify, and/or distribute this software for any purpose with or without fee is hereby granted, provided that the above copyright notice and this permission notice appear in all copies. 4 | 5 | THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. 6 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # DALL·E mini Telegram bot 2 | 3 | [![@dalle_mini_bot](https://img.shields.io/badge/Telegram%20Bot-@dalle_mini_bot-blue?logo=telegram&style=plastic)](https://telegram.me/dalle_mini_bot) 4 | 5 | A Telegram bot interface for [DALL·E mini](https://github.com/borisdayma/dalle-mini). 6 | Request 9 AI-generated images from any prompt you give, directly from Telegram. 7 | 8 | ![Bot logo, generated by DALL·E mini with the prompt "a cat playing with a paper plane"](docs/bot-logo.jpg) 9 | 10 | [![Telegram Bot screenshot](docs/Telegram-DalleMiniBot-screenshot.png)](https://telegram.me/dalle_mini_bot) 11 | 12 | ## Features 13 | 14 | - Request from Telegram, return the 9 pictures result as an album 15 | - Status report while the images are being generated (the bot sends a 'typing-like' status to the user, until all its requests are completed) 16 | - If the server is too busy, keep retrying until success (or timeout) 17 | 18 | The bot is deployed here: [https://telegram.me/dalle_mini_bot](https://telegram.me/dalle_mini_bot) 19 | 20 | ## Changelog 21 | 22 | - v0.2.2 23 | - Support Redis authentication 24 | - Fixed exceptions not being included on "Request failed" log records 25 | - Fixed loguru incompatibility with ApiTelegramException 26 | - Ignore errors when pushing logs to Redis 27 | - v0.2.1 28 | - Graceful shutdown (configurable; wait until pending requests are completed, while not accepting new requests) 29 | - Retry Telegram Bot API requests on 'Too Many Requests' error; usage of requests.Session 30 | - Set bot commands via API on startup, for Telegram hinting 31 | - Limit prompt text length on Generate command (configurable min/max limits via settings) 32 | - Add Redis integration for sending logs to Redis queue 33 | - Improvements in log records 34 | - v0.1.1 35 | - Send message to users while the image is being generated, informing that it may take a while; the message is deleted on success or controlled error 36 | - Add `/about` command 37 | - Timeout chat 'typing-like' action & stop it when bot blocked by user 38 | - Remove chats with count of 0 from Counters 39 | - Setting for deleting Telegram bot Webhook on startup 40 | - v0.0.2 41 | - Add pysocks requirement 42 | - Detect when bot blocked by user on middleware 43 | - v0.0.1 44 | - Initial release 45 | - Generate images from `/generate` command, return as album 46 | - `/generate` command rate limited at chat level (concurrent requests limit) 47 | - `/generate` command sends a 'typing-like' status to the user, while the prompt/s requested are being generated 48 | -------------------------------------------------------------------------------- /__main__.py: -------------------------------------------------------------------------------- 1 | from dalle_telegram_bot.entrypoint import main 2 | 3 | main() 4 | -------------------------------------------------------------------------------- /dalle_telegram_bot/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/David-Lor/dalle-mini-telegram-bot/8bc9ccb68e02101c246bc0fb8ff4c1500b8ab862/dalle_telegram_bot/__init__.py -------------------------------------------------------------------------------- /dalle_telegram_bot/entrypoint.py: -------------------------------------------------------------------------------- 1 | import signal 2 | from threading import Event, Lock 3 | 4 | from .services.bot import Bot 5 | from .services.dalle import Dalle 6 | from .services.redis import Redis 7 | from .settings import Settings 8 | from .logger import logger, setup_logger 9 | 10 | 11 | class BotBackend: 12 | settings: Settings 13 | redis: Redis 14 | dalle: Dalle 15 | bot: Bot 16 | _teardown_event: Event 17 | _teardown_lock: Lock 18 | 19 | def setup(self): 20 | self.settings = Settings() 21 | self._teardown_event = Event() 22 | self._teardown_lock = Lock() 23 | 24 | self.redis = Redis( 25 | settings=self.settings, 26 | ) 27 | setup_logger( 28 | settings=self.settings, 29 | loggers=[self.redis], 30 | ) 31 | logger.debug("Initializing app...") 32 | 33 | self.dalle = Dalle( 34 | settings=self.settings, 35 | ) 36 | self.bot = Bot( 37 | settings=self.settings, 38 | dalle=self.dalle, 39 | ) 40 | logger.debug("App initialized") 41 | 42 | def run(self): 43 | try: 44 | self.start() 45 | self.wait_for_end() 46 | finally: 47 | self.teardown() 48 | 49 | def wait_for_end(self): 50 | self._teardown_event.wait() 51 | 52 | def teardown(self, *_): 53 | with self._teardown_lock: 54 | if self._teardown_event.is_set(): 55 | return 56 | 57 | try: 58 | self.stop() 59 | finally: 60 | self._teardown_event.set() 61 | 62 | def start(self): 63 | logger.debug("Running app...") 64 | self.bot.setup() 65 | self.bot.start() 66 | 67 | def stop(self): 68 | logger.info("Stopping app...") 69 | self.bot.stop() 70 | logger.info("App stopped!") 71 | 72 | 73 | def main(): 74 | app = BotBackend() 75 | app.setup() 76 | 77 | signal.signal(signal.SIGINT, app.teardown) 78 | signal.signal(signal.SIGTERM, app.teardown) 79 | app.run() 80 | -------------------------------------------------------------------------------- /dalle_telegram_bot/logger.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import contextlib 3 | from typing import Collection, Optional 4 | 5 | from loguru import logger 6 | # noinspection PyProtectedMember 7 | from loguru._logger import context as loguru_context 8 | 9 | from .settings import Settings 10 | from .services.logger_abc import AbstractLogger 11 | 12 | __all__ = ("logger", "setup_logger", "get_request_id") 13 | 14 | 15 | LoggerFormat = "{time:YY-MM-DD HH:mm:ss} | " \ 16 | "{level} | " \ 17 | "{function}: {message} | " \ 18 | "{extra} {exception}" 19 | 20 | 21 | def setup_logger(settings: Settings, loggers: Collection[AbstractLogger]): 22 | logger.remove() 23 | logger.add(sys.stderr, level=settings.log_level.upper(), format=LoggerFormat) 24 | 25 | for custom_logger in (loggers or []): 26 | logger.add( 27 | custom_logger.log, 28 | level="TRACE", 29 | enqueue=True, 30 | serialize=True, # record provided as JSON string to the handler 31 | ) 32 | 33 | 34 | def get_request_id() -> Optional[str]: 35 | """Return the current Request ID, if defined, being used by the logger""" 36 | with contextlib.suppress(Exception): 37 | context: dict = loguru_context.get() 38 | return context.get("request_id") 39 | -------------------------------------------------------------------------------- /dalle_telegram_bot/services/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/David-Lor/dalle-mini-telegram-bot/8bc9ccb68e02101c246bc0fb8ff4c1500b8ab862/dalle_telegram_bot/services/__init__.py -------------------------------------------------------------------------------- /dalle_telegram_bot/services/bot/__init__.py: -------------------------------------------------------------------------------- 1 | from .bot import Bot 2 | -------------------------------------------------------------------------------- /dalle_telegram_bot/services/bot/bot.py: -------------------------------------------------------------------------------- 1 | import contextlib 2 | from threading import Thread 3 | from typing import Optional 4 | 5 | import telebot 6 | from telebot.types import Message, InputMediaPhoto, BotCommand 7 | 8 | from . import constants 9 | from .requester import TelegramBotAPIRequester 10 | from .chatactions import ActionManager 11 | from .middlewares import request_middleware, message_request_middleware, RateLimiter 12 | from ..dalle import Dalle, DalleTemporarilyUnavailableException 13 | from ..dalle.models import DalleResponse 14 | from ...settings import Settings 15 | from ...logger import logger 16 | 17 | 18 | class Bot: 19 | def __init__(self, settings: Settings, dalle: Dalle): 20 | self._settings = settings 21 | self._dalle = dalle 22 | self._polling_thread = None 23 | 24 | self._bot = telebot.TeleBot( 25 | token=self._settings.telegram_bot_token, 26 | parse_mode="HTML", 27 | threaded=True, 28 | num_threads=settings.telegram_bot_threads, 29 | ) 30 | self._bot.message_handler(func=lambda message: True)(self._handler_message_entrypoint) 31 | 32 | self._generating_bot_action = ActionManager( 33 | action=self._settings.command_generate_action, 34 | bot=self._bot, 35 | settings=self._settings, 36 | timeout=self._settings.dalle_generation_timeout_seconds, 37 | ) 38 | self._dalle_generate_rate_limiter = RateLimiter( 39 | limit_per_chat=self._settings.command_generate_chat_concurrent_limit, 40 | ) 41 | 42 | self._requester = None 43 | if self._settings.telegram_bot_ratelimit_retry: 44 | self._requester = TelegramBotAPIRequester( 45 | settings=self._settings, 46 | ) 47 | 48 | def setup(self): 49 | """Perform initial setup (delete webhook, set commands)""" 50 | if self._settings.telegram_bot_delete_webhook: 51 | self.delete_webhook() 52 | if self._settings.telegram_bot_set_commands: 53 | self.set_commands() 54 | 55 | def start(self): 56 | """Run the bot in background, by starting a thread running the `run` method.""" 57 | if self._polling_thread: 58 | return 59 | 60 | if self._requester: 61 | self._requester.start() 62 | 63 | self._polling_thread = Thread( 64 | target=self.run, 65 | name="TelegramBotPolling", 66 | daemon=True, 67 | ) 68 | self._polling_thread.start() 69 | 70 | def run(self): 71 | """Run the bot in foreground. Perform initial setup (delete webhook, set commands)""" 72 | logger.info("Running bot with Polling") 73 | self._bot.infinity_polling(logger_level=None) 74 | 75 | def stop(self, graceful_shutdown: Optional[bool] = None): 76 | """Stop the bot execution. 77 | :param graceful_shutdown: if True, wait for pending requests to finalize (but stop accepting new requests). 78 | If false, stop inmediately. If None (default), use the configured setting. 79 | """ 80 | if graceful_shutdown is None: 81 | graceful_shutdown = self._settings.telegram_bot_graceful_shutdown 82 | 83 | if graceful_shutdown: 84 | self._stop_gracefully() 85 | else: 86 | self._stop_force() 87 | 88 | if self._requester: 89 | self._requester.teardown() 90 | 91 | def set_commands(self): 92 | logger.debug("Setting bot commands...") 93 | self._bot.set_my_commands([ 94 | BotCommand(command=k, description=v) 95 | for k, v in constants.COMMANDS_HELP.items() 96 | ]) 97 | logger.info("Bot commands set") 98 | 99 | def delete_webhook(self): 100 | logger.info("Deleting bot webhook...") 101 | self._bot.delete_webhook() 102 | logger.info("Webhook deleted") 103 | 104 | def _stop_force(self): 105 | logger.info("Stopping bot polling (force-stop)...") 106 | self._bot.stop_polling() 107 | logger.info("Bot stopped") 108 | 109 | def _stop_gracefully(self): 110 | logger.info("Stopping bot gracefully (waiting for pending requests to end, not accepting new requests)...") 111 | # stop_bot() waits for remaining requests to complete 112 | self._bot.stop_bot() 113 | logger.info("Bot stopped") 114 | 115 | def _handler_message_entrypoint(self, message: Message): 116 | with request_middleware(chat_id=message.chat.id): 117 | with message_request_middleware(bot=self._bot, message=message): 118 | if self._handler_basic_command(message): 119 | return 120 | if self._handler_command_generate(message): 121 | return 122 | 123 | def _handler_basic_command(self, message: Message) -> bool: 124 | for cmd, reply_text in constants.BASIC_COMMAND_REPLIES.items(): 125 | if message.text.startswith(cmd): 126 | logger.bind(command=cmd).info("Request is Basic command") 127 | 128 | disable_link_preview = cmd in constants.BASIC_COMMAND_DISABLE_LINK_PREVIEWS 129 | self._bot.reply_to( 130 | message=message, 131 | text=reply_text, 132 | disable_web_page_preview=disable_link_preview, 133 | ) 134 | return True 135 | 136 | return False 137 | 138 | def _handler_command_generate(self, message: Message) -> bool: 139 | if not message.text.startswith(constants.COMMAND_GENERATE): 140 | return False 141 | 142 | logger.bind(cmd=constants.COMMAND_GENERATE).info("Request is Generate command") 143 | prompt = self.__command_generate_get_prompt(message) 144 | if not prompt: 145 | return True 146 | 147 | if not self._dalle_generate_rate_limiter.increase(message.chat.id): 148 | logger.bind(chat_id=message.chat.id).info("Generate command Request limit exceeded for this chat") 149 | self._bot.reply_to(message, constants.COMMAND_GENERATE_REPLY_RATELIMIT_EXCEEDED) 150 | return True 151 | 152 | generating_reply_message = self._bot.reply_to(message, constants.COMMAND_GENERATE_REPLY_GENERATING) 153 | self._generating_bot_action.start(message.chat.id) 154 | 155 | response: Optional[DalleResponse] = None 156 | try: 157 | response = self._dalle.generate(prompt) 158 | except DalleTemporarilyUnavailableException: 159 | pass 160 | finally: 161 | self._generating_bot_action.stop(message.chat.id) 162 | self._dalle_generate_rate_limiter.decrease(message.chat.id) 163 | with contextlib.suppress(Exception): 164 | self._bot.delete_message( 165 | chat_id=generating_reply_message.chat.id, 166 | message_id=generating_reply_message.message_id 167 | ) 168 | 169 | if not response: 170 | self._bot.reply_to(message, constants.COMMAND_GENERATE_REPLY_TEMPORARILY_UNAVAILABLE) 171 | return True 172 | 173 | images_telegram = [InputMediaPhoto(image_bytes) for image_bytes in response.images_bytes] 174 | images_telegram[0].caption = prompt 175 | self._bot.send_media_group( 176 | chat_id=message.chat.id, 177 | reply_to_message_id=message.message_id, 178 | media=images_telegram, 179 | ) 180 | return True 181 | 182 | def __command_generate_get_prompt(self, message: Message) -> Optional[str]: 183 | """Get the prompt text from a /generate command and return it. 184 | If the prompt is invalid, replies to the user and returns None.""" 185 | prompt = message.text.replace(constants.COMMAND_GENERATE, "").strip() 186 | prompt_length = len(prompt) 187 | min_length = self._settings.command_generate_prompt_length_min 188 | max_length = self._settings.command_generate_prompt_length_max 189 | 190 | with logger.contextualize(prompt_length=prompt_length): 191 | if prompt_length < min_length: 192 | logger.debug("Generate command prompt too short") 193 | self._bot.reply_to(message, constants.COMMAND_GENERATE_PROMPT_TOO_SHORT.format(characters=min_length)) 194 | return None 195 | 196 | if prompt_length > max_length: 197 | logger.debug("Generate command prompt too long") 198 | self._bot.reply_to(message, constants.COMMAND_GENERATE_PROMPT_TOO_LONG.format(characters=max_length)) 199 | return None 200 | 201 | logger.debug("Generate command prompt is valid") 202 | return prompt 203 | -------------------------------------------------------------------------------- /dalle_telegram_bot/services/bot/chatactions.py: -------------------------------------------------------------------------------- 1 | import time 2 | from threading import Thread, Event, Lock 3 | from collections import Counter 4 | from typing import Dict 5 | 6 | import telebot 7 | 8 | from ...utils import exception_is_bot_blocked_by_user 9 | from ...logger import logger, get_request_id 10 | from ...settings import Settings 11 | 12 | 13 | class ActionManager: 14 | def __init__(self, action: str, timeout: float, bot: telebot.TeleBot, settings: Settings): 15 | self._action = action 16 | self._timeout = timeout 17 | self._bot = bot 18 | self._settings = settings 19 | 20 | self._chatids_events: Dict[int, Event] = dict() 21 | self._chatids_counter = Counter() 22 | self._chatids_counter_lock = Lock() 23 | 24 | def start(self, chat_id: int): 25 | """Register a 'start' Action for a chat. 26 | If no action was currently running for the chat, start it. 27 | In all cases, increase the counter for the chat.""" 28 | do_start = False 29 | with self._chatids_counter_lock: 30 | if self._chatids_counter[chat_id] == 0: 31 | do_start = True 32 | self.increase(chat_id) 33 | 34 | if do_start: 35 | # run the start outside the counter lock 36 | self._start_action_thread(chat_id) 37 | 38 | def stop(self, chat_id: int): 39 | """Register a 'stop' Action for a chat. 40 | Decrease the counter; if just one action was running for the chat, stop it.""" 41 | do_stop = False 42 | with self._chatids_counter_lock: 43 | if self._chatids_counter[chat_id] != 0: 44 | if self.decrease(chat_id) == 0: 45 | do_stop = True 46 | 47 | if do_stop: 48 | # run the stop outside the counter lock 49 | self._stop_action_thread(chat_id) 50 | 51 | def increase(self, chat_id: int) -> int: 52 | """Increase the request counter for a chat, and return the new value. 53 | The counter access should be locked while calling this method.""" 54 | self._chatids_counter[chat_id] += 1 55 | return self._chatids_counter[chat_id] 56 | 57 | def decrease(self, chat_id: int) -> int: 58 | """Decrease the request counter for a chat, and return the new value. 59 | If the new value is 0, remove the referenced that from the Counter. 60 | The counter access should be locked while calling this method.""" 61 | self._chatids_counter[chat_id] -= 1 62 | 63 | current = self._chatids_counter[chat_id] 64 | if current < 0: 65 | current = 0 66 | if current == 0: 67 | del self._chatids_counter[chat_id] 68 | 69 | return current 70 | 71 | def _start_action_thread(self, chat_id: int): 72 | """Start the action thread. The chat_id MUST not be currently running any actions.""" 73 | event = Event() 74 | self._chatids_events[chat_id] = event 75 | request_id = get_request_id() 76 | 77 | Thread( 78 | target=self._action_worker, 79 | kwargs=dict( 80 | request_id=request_id, 81 | chat_id=chat_id, 82 | event=event, 83 | ), 84 | name=f"TelegramBot-ActionWorker-{self._action}-{chat_id}", 85 | daemon=True 86 | ).start() 87 | 88 | def _stop_action_thread(self, chat_id: int): 89 | """Stop the action thread for a chat_id. The chat_id MUST be currently running an action.""" 90 | event = self._chatids_events.pop(chat_id, None) 91 | if event: 92 | event.set() 93 | 94 | def _action_worker(self, request_id: str, chat_id: int, event: Event): 95 | with logger.contextualize(request_id=request_id, chat_id=chat_id, chat_action=self._action): 96 | start = time.time() 97 | while not event.is_set(): 98 | try: 99 | logger.trace("Sending chat action...") 100 | self._bot.send_chat_action( 101 | chat_id=chat_id, 102 | action=self._action, 103 | ) 104 | logger.debug("Chat action sent") 105 | 106 | except Exception as ex: 107 | if exception_is_bot_blocked_by_user(ex): 108 | logger.info("Bot blocked by user, stopping chat action") 109 | self._stop_action_thread(chat_id) 110 | return 111 | logger.opt(exception=ex).warning("Chat action failed delivery") 112 | 113 | elapsed = time.time() - start 114 | if elapsed >= self._timeout: 115 | logger.bind(elapsed_time_seconds=round(elapsed, 3)).warning("Chat action timed out") 116 | self._stop_action_thread(chat_id) 117 | return 118 | 119 | event.wait(4.5) 120 | 121 | logger.debug("Chat action finalized") 122 | -------------------------------------------------------------------------------- /dalle_telegram_bot/services/bot/constants.py: -------------------------------------------------------------------------------- 1 | BASIC_COMMAND_REPLIES = { 2 | "/start": """👋 Hello there!\n\nThis bot returns 9 AI-generated images from any prompt you give, using DALL·E mini.\n\nExample: Send a message like 
/generate a cat eating a hamburger
to generate a set of images.\n\nCheck /help and /about for more info.""", 3 | "/help": "Send a message like 
/generate a cat eating a hamburger
to generate a set of images.", 4 | "/about": """DALL·E mini Telegram bot is a bot that returns 9 AI-generated images from any prompt you give, using DALL·E mini.\n\nHow does it work?\nThe provided text prompt is passed to DALL·E mini, the same way it's done on the official webpage. The returned pictures are forwarded to Telegram via the bot and sent to you.\n\nBuilt with ♥️ using Python and pyTelegramBotAPI.\nSource code: https://github.com/David-Lor/dalle-mini-telegram-bot""", 5 | } 6 | 7 | BASIC_COMMAND_DISABLE_LINK_PREVIEWS = {"/start", "/help", "/about"} 8 | 9 | COMMAND_GENERATE = "/generate" 10 | 11 | COMMAND_GENERATE_REPLY_GENERATING = "The image is being generated. Please wait a few minutes for it..." 12 | 13 | COMMAND_GENERATE_REPLY_RATELIMIT_EXCEEDED = "You have other images being generated. " \ 14 | "Please wait until those are sent to you before asking for more." 15 | 16 | COMMAND_GENERATE_PROMPT_TOO_SHORT = "Your prompt message is too short, try with something longer " \ 17 | "(at least {characters} characters)." 18 | 19 | COMMAND_GENERATE_PROMPT_TOO_LONG = "Your prompt message is too long, try with something shorter " \ 20 | "(at most {characters} characters)." 21 | 22 | COMMAND_GENERATE_REPLY_TEMPORARILY_UNAVAILABLE = "Your image could not be generated. Please try again later." 23 | 24 | UNKNOWN_ERROR_REPLY = "Unknown error. Please try again later." 25 | 26 | COMMANDS_HELP = { 27 | "/generate": "Generate a set of pictures from a given prompt. The prompt must be given after the /generate command", 28 | "/help": "Help about the bot usage", 29 | "/about": "About the bot", 30 | } 31 | -------------------------------------------------------------------------------- /dalle_telegram_bot/services/bot/middlewares.py: -------------------------------------------------------------------------------- 1 | import contextlib 2 | import time 3 | import threading 4 | from threading import Lock 5 | from collections import Counter 6 | 7 | import telebot 8 | from telebot.types import Message 9 | from telebot.apihelper import ApiTelegramException 10 | 11 | from . import constants 12 | from ...logger import logger 13 | from ...utils import get_uuid, exception_is_bot_blocked_by_user 14 | 15 | 16 | @contextlib.contextmanager 17 | def request_middleware(chat_id: int): 18 | request_id = get_uuid() 19 | start = time.time() 20 | 21 | with logger.contextualize(request_id=request_id): 22 | try: 23 | logger.bind( 24 | chat_id=chat_id, 25 | thread_name=threading.current_thread().name 26 | ).info("Request started") 27 | yield 28 | 29 | except Exception as ex: 30 | request_duration = round(time.time() - start, 4) 31 | with logger.contextualize(request_duration=request_duration): 32 | if exception_is_bot_blocked_by_user(ex): 33 | logger.info("Request completed: Bot blocked by the user") 34 | return 35 | 36 | if isinstance(ex, ApiTelegramException): 37 | logger.bind(error_json=ex.result_json).error("Request failed") 38 | return 39 | 40 | logger.exception("Request failed") 41 | 42 | else: 43 | request_duration = round(time.time() - start, 4) 44 | logger.bind(request_duration=request_duration).info("Request completed") 45 | 46 | 47 | @contextlib.contextmanager 48 | def message_request_middleware(bot: telebot.TeleBot, message: Message): 49 | try: 50 | yield 51 | except Exception as ex: 52 | bot.reply_to(message, constants.UNKNOWN_ERROR_REPLY) 53 | raise ex 54 | 55 | 56 | class RateLimiter: 57 | def __init__(self, limit_per_chat: int): 58 | self._limit_per_chat = limit_per_chat 59 | self._counter = Counter() 60 | self._counter_lock = Lock() 61 | 62 | def increase(self, chat_id: int) -> bool: 63 | with self._counter_lock: 64 | current = self._counter[chat_id] 65 | if current >= self._limit_per_chat: 66 | return False 67 | 68 | self._counter[chat_id] = current + 1 69 | return True 70 | 71 | def decrease(self, chat_id: int): 72 | new_value = self._counter[chat_id] - 1 73 | if new_value <= 0: 74 | del self._counter[chat_id] 75 | return 76 | self._counter[chat_id] = new_value 77 | -------------------------------------------------------------------------------- /dalle_telegram_bot/services/bot/requester.py: -------------------------------------------------------------------------------- 1 | import contextlib 2 | import threading 3 | import time 4 | from typing import Dict 5 | 6 | import requests 7 | import wait4it 8 | import telebot.apihelper 9 | 10 | from ...settings import Settings 11 | from ...logger import logger 12 | 13 | 14 | class TelegramBotAPIRequester: 15 | def __init__(self, settings: Settings): 16 | self._settings = settings 17 | self._sessions: Dict[str, requests.Session] = dict() 18 | self._sessions_last_timestamp: Dict[str, float] = dict() 19 | self._sessions_lock = threading.Lock() 20 | self._cleanup_thread = None 21 | 22 | telebot.apihelper.CUSTOM_REQUEST_SENDER = wait4it.wait_for_pass( 23 | exceptions=[TelegramBotAPITooManyRequestsException], # TODO Add other Request errors? 24 | retries=self._settings.telegram_bot_ratelimit_retries_limit, 25 | retries_delay=self._settings.telegram_bot_ratelimit_retry_delay_seconds, 26 | )(self.request) 27 | 28 | def start(self): 29 | """Start the cleanup thread""" 30 | if self._cleanup_thread: 31 | return 32 | 33 | self._cleanup_thread = threading.Thread( 34 | target=self._cleanup_worker, 35 | name="TelegramBotAPIRequester-cleanup", 36 | daemon=True, 37 | ).start() 38 | 39 | def request(self, *args, **kwargs): 40 | if self._settings.telegram_bot_api_sessions_enabled: 41 | session = self.get_session() 42 | r = session.request(*args, **kwargs) 43 | else: 44 | r = requests.request(*args, **kwargs) 45 | 46 | if self._response_is_toomanyrequests(r): 47 | raise TelegramBotAPITooManyRequestsException(r.json().get("description")) 48 | return r 49 | 50 | def get_session(self) -> requests.Session: 51 | thread_name = threading.current_thread().name 52 | with self._sessions_lock: 53 | self._sessions_last_timestamp[thread_name] = time.time() 54 | session = self._sessions.get(thread_name) 55 | 56 | if not session: 57 | session = requests.Session() 58 | self._sessions[thread_name] = session 59 | logger.bind(thread_name=thread_name).trace("New requests.Session created") 60 | 61 | return session 62 | 63 | def teardown(self): 64 | with logger.contextualize(sessions_count=len(self._sessions)): 65 | logger.debug("Closing TelegramBotAPI request sessions...") 66 | for k in list(self._sessions.keys()): 67 | self.stop_session(k) 68 | 69 | logger.info("Closed TelegramBotAPI request sessions") 70 | 71 | def stop_session(self, thread_name: str): 72 | with self._sessions_lock: 73 | session = self._sessions.pop(thread_name, None) 74 | self._sessions_last_timestamp.pop(thread_name, None) 75 | 76 | if not session: 77 | return 78 | 79 | with contextlib.suppress(Exception): 80 | logger.bind(thread_name=thread_name).trace("Stopping requests.Session") 81 | session.close() 82 | 83 | def _cleanup_worker(self): 84 | logger.debug("Start of TelegramBotAPIRequester requests.Sessions cleanup worker") 85 | while True: 86 | with self._sessions_lock: 87 | sessions_timestamps = tuple(self._sessions_last_timestamp.items()) 88 | 89 | now = time.time() 90 | for session_threadname, session_timestamp in sessions_timestamps: 91 | if (now - session_timestamp) >= self._settings.telegram_bot_api_sessions_ttl_seconds: 92 | self.stop_session(session_threadname) 93 | 94 | time.sleep(30) 95 | 96 | @staticmethod 97 | def _response_is_toomanyrequests(response: requests.Response) -> bool: 98 | return response.status_code == 429 99 | 100 | 101 | class TelegramBotAPITooManyRequestsException(Exception): 102 | pass 103 | -------------------------------------------------------------------------------- /dalle_telegram_bot/services/dalle/__init__.py: -------------------------------------------------------------------------------- 1 | from .dalle import * 2 | from .exceptions import * 3 | -------------------------------------------------------------------------------- /dalle_telegram_bot/services/dalle/dalle.py: -------------------------------------------------------------------------------- 1 | import requests 2 | import wait4it 3 | 4 | from .models import DalleResponse 5 | from .exceptions import DalleTemporarilyUnavailableException 6 | from ...settings import Settings 7 | from ...logger import logger 8 | 9 | __all__ = ("Dalle",) 10 | 11 | 12 | class Dalle: 13 | def __init__(self, settings: Settings): 14 | self._settings = settings 15 | 16 | self._generate_until_complete = wait4it.wait_for_pass( 17 | exceptions=(DalleTemporarilyUnavailableException,), 18 | retries=self._settings.dalle_generation_retries_limit, 19 | retries_delay=self._settings.dalle_generation_retry_delay_seconds, 20 | )(lambda prompt: self._simple_request(prompt)) 21 | 22 | def generate(self, prompt: str) -> DalleResponse: 23 | return self._generate_until_complete(prompt) 24 | 25 | def _simple_request(self, prompt: str) -> DalleResponse: 26 | logger.debug("Requesting DALLE...") 27 | body = dict( 28 | prompt=prompt, 29 | ) 30 | response = requests.post( 31 | url=self._settings.dalle_api_url, 32 | json=body, 33 | timeout=self._settings.dalle_api_request_timeout_seconds, 34 | proxies=self._settings.dalle_api_request_socks_proxy_for_requests_lib, 35 | ) 36 | logger.bind(status_code=response.status_code).debug("DALLE response received") 37 | 38 | return self._parse_response( 39 | response=response, 40 | prompt=prompt, 41 | ) 42 | 43 | @staticmethod 44 | def _parse_response(prompt: str, response: requests.Response) -> DalleResponse: 45 | if response.status_code == 503: 46 | raise DalleTemporarilyUnavailableException() 47 | response.raise_for_status() 48 | 49 | return DalleResponse( 50 | **response.json(), 51 | prompt=prompt, 52 | ) 53 | -------------------------------------------------------------------------------- /dalle_telegram_bot/services/dalle/exceptions.py: -------------------------------------------------------------------------------- 1 | __all__ = ("BaseDalleException", "DalleTemporarilyUnavailableException") 2 | 3 | 4 | class BaseDalleException(Exception): 5 | pass 6 | 7 | 8 | class DalleTemporarilyUnavailableException(BaseDalleException): 9 | pass 10 | -------------------------------------------------------------------------------- /dalle_telegram_bot/services/dalle/models.py: -------------------------------------------------------------------------------- 1 | import base64 2 | import pathlib 3 | from typing import List 4 | 5 | import pydantic 6 | 7 | 8 | class DalleResponse(pydantic.BaseModel): 9 | # Fields returned from API response 10 | images: List[str] = pydantic.Field(..., min_items=9, max_items=9) # images as base64 strings 11 | 12 | # Fields we complete 13 | prompt: str 14 | 15 | @property 16 | def images_bytes(self): 17 | images_parsed: List[bytes] = list() 18 | for image_base64 in self.images: 19 | images_parsed.append(base64.b64decode(image_base64)) 20 | return images_parsed 21 | 22 | def save_images(self, directory: str): 23 | directory_path = pathlib.Path(directory) 24 | for i, image_data in enumerate(self.images_bytes, start=1): 25 | image_filename = f"{self.prompt} - {i}.jpg" 26 | image_path = directory_path / image_filename 27 | with open(image_path, "wb") as f: 28 | f.write(image_data) 29 | -------------------------------------------------------------------------------- /dalle_telegram_bot/services/logger_abc.py: -------------------------------------------------------------------------------- 1 | import abc 2 | 3 | 4 | class AbstractLogger(abc.ABC): 5 | @abc.abstractmethod 6 | def log(self, data: str): 7 | """ 8 | :param data: log record as JSON string 9 | """ 10 | pass 11 | -------------------------------------------------------------------------------- /dalle_telegram_bot/services/redis.py: -------------------------------------------------------------------------------- 1 | import redis 2 | 3 | from .logger_abc import AbstractLogger 4 | from ..settings import Settings 5 | 6 | 7 | class Redis(AbstractLogger): 8 | def __init__(self, settings: Settings): 9 | self._settings = settings 10 | self._redis = None 11 | if not self._settings.redis_host: 12 | return 13 | 14 | self._redis = redis.StrictRedis( 15 | host=self._settings.redis_host, 16 | port=self._settings.redis_port, 17 | db=self._settings.redis_db, 18 | **self._get_auth_kwargs(), 19 | ) 20 | 21 | def log(self, data: str): 22 | if not self._redis or not self._settings.redis_logs_queue_name: 23 | return 24 | 25 | try: 26 | self._redis.rpush( 27 | self._settings.redis_logs_queue_name, 28 | data, 29 | ) 30 | except Exception: 31 | # TODO Log errors? 32 | pass 33 | 34 | def _get_auth_kwargs(self): 35 | kwargs = dict() 36 | if self._settings.redis_username: 37 | kwargs["username"] = self._settings.redis_username 38 | if self._settings.redis_password: 39 | kwargs["password"] = self._settings.redis_password 40 | return kwargs 41 | -------------------------------------------------------------------------------- /dalle_telegram_bot/settings.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | import pydantic 4 | 5 | 6 | class Settings(pydantic.BaseSettings): 7 | telegram_bot_token: str 8 | telegram_bot_threads: int = 500 9 | telegram_bot_delete_webhook: bool = False 10 | telegram_bot_graceful_shutdown: bool = False 11 | telegram_bot_set_commands: bool = False 12 | telegram_bot_api_sessions_enabled: bool = True 13 | telegram_bot_api_sessions_ttl_seconds: float = 360 14 | telegram_bot_ratelimit_retry: bool = True 15 | telegram_bot_ratelimit_retry_delay_seconds: float = 5 16 | telegram_bot_ratelimit_retry_timeout_seconds: float = 120 17 | 18 | command_generate_action: str = "typing" 19 | command_generate_chat_concurrent_limit: int = 3 20 | command_generate_prompt_length_min: int = pydantic.Field(default=2, gt=1) 21 | command_generate_prompt_length_max: int = pydantic.Field(default=1000, gt=1) 22 | 23 | dalle_api_url: pydantic.AnyHttpUrl = "https://bf.dallemini.ai/generate" 24 | dalle_api_request_timeout_seconds: float = 3.5 * 60 25 | dalle_api_request_socks_proxy: Optional[pydantic.AnyUrl] = None 26 | dalle_generation_timeout_seconds: float = 6 * 60 27 | dalle_generation_retry_delay_seconds: float = 5 28 | 29 | redis_host: Optional[str] = None 30 | redis_port: int = 6379 31 | redis_db: int = 0 32 | redis_username: Optional[str] = None 33 | redis_password: Optional[str] = None 34 | redis_logs_queue_name: Optional[str] = None 35 | 36 | log_level: str = "INFO" 37 | 38 | @property 39 | def dalle_generation_retries_limit(self) -> int: 40 | return int(self.dalle_generation_timeout_seconds / self.dalle_generation_retry_delay_seconds) 41 | 42 | @property 43 | def telegram_bot_ratelimit_retries_limit(self) -> int: 44 | return int(self.telegram_bot_ratelimit_retry_timeout_seconds / self.telegram_bot_ratelimit_retry_delay_seconds) 45 | 46 | @property 47 | def dalle_api_request_socks_proxy_for_requests_lib(self) -> Optional[dict]: 48 | if not self.dalle_api_request_socks_proxy: 49 | return None 50 | proxy = str(self.dalle_api_request_socks_proxy) 51 | return dict( 52 | http=proxy, 53 | https=proxy, 54 | ) 55 | 56 | class Config: 57 | env_file = ".env" 58 | -------------------------------------------------------------------------------- /dalle_telegram_bot/utils.py: -------------------------------------------------------------------------------- 1 | import shortuuid 2 | from telebot.apihelper import ApiTelegramException 3 | 4 | __all__ = ("get_uuid", "exception_is_bot_blocked_by_user") 5 | 6 | 7 | def get_uuid() -> str: 8 | return shortuuid.uuid() 9 | 10 | 11 | def exception_is_bot_blocked_by_user(ex: Exception) -> bool: 12 | return isinstance(ex, ApiTelegramException) and ex.description == "Forbidden: bot was blocked by the user" 13 | -------------------------------------------------------------------------------- /docs/Telegram-DalleMiniBot-screenshot.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/David-Lor/dalle-mini-telegram-bot/8bc9ccb68e02101c246bc0fb8ff4c1500b8ab862/docs/Telegram-DalleMiniBot-screenshot.png -------------------------------------------------------------------------------- /docs/bot-logo.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/David-Lor/dalle-mini-telegram-bot/8bc9ccb68e02101c246bc0fb8ff4c1500b8ab862/docs/bot-logo.jpg -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | pytelegrambotapi==4.5.1 2 | pydantic==1.9.1 3 | requests==2.28.0 4 | pysocks==1.7.1 5 | shortuuid==1.0.9 6 | wait4it==0.2.1 7 | redis==4.3.3 8 | loguru==0.6.0 9 | python-dotenv==0.20.0 10 | -------------------------------------------------------------------------------- /sample.env: -------------------------------------------------------------------------------- 1 | # TELEGRAM_BOT_TOKEN: bot token returned by BotFather. Keep it safe! 2 | TELEGRAM_BOT_TOKEN= 3 | 4 | # TELEGRAM_BOT_THREADS: number of request handling threads for the bots, equals to amount of concurrent requests that can be handled 5 | TELEGRAM_BOT_THREADS=500 6 | 7 | # TELEGRAM_BOT_DELETE_WEBHOOK: if enabled, delete bot webhook on startup, before starting the bot polling 8 | TELEGRAM_BOT_DELETE_WEBHOOK=0 9 | 10 | # TELEGRAM_BOT_GRACEFUL_SHUTDOWN: if enabled, on app shutdown, wait until remaining requests are completed (but not accept new requests) 11 | TELEGRAM_BOT_GRACEFUL_SHUTDOWN=0 12 | 13 | # TELEGRAM_BOT_RATELIMIT_RETRY: if enabled, retry Telegram Bot API requests if failed because of Too Many Requests error 14 | TELEGRAM_BOT_RATELIMIT_RETRY=1 15 | 16 | # TELEGRAM_BOT_RATELIMIT_RETRY_DELAY_SECONDS: delay between Telegram Bot API retrying requests, because of Too Many Requests error 17 | TELEGRAM_BOT_RATELIMIT_RETRY_DELAY_SECONDS=5 18 | 19 | # TELEGRAM_BOT_RATELIMIT_RETRY_TIMEOUT_SECONDS: timeout for trying to send a Telegram Bot API request, failed because of Too Many Requests error 20 | TELEGRAM_BOT_RATELIMIT_RETRY_TIMEOUT_SECONDS=120 21 | 22 | # TELEGRAM_BOT_API_SESSIONS_ENABLED: if enabled, perform Telegram Bot API requests with requests.Session. Only used if TELEGRAM_BOT_RATELIMIT_RETRY enabled 23 | TELEGRAM_BOT_API_SESSIONS_ENABLED=1 24 | 25 | # TELEGRAM_BOT_API_SESSIONS_TTL_SECONDS: after this time (seconds), requests.Sessions that have not been used will be closed 26 | TELEGRAM_BOT_API_SESSIONS_TTL_SECONDS=360 27 | 28 | # COMMAND_GENERATE_ACTION: chat action to send while generating. One of: https://core.telegram.org/bots/api#sendchataction 29 | COMMAND_GENERATE_ACTION=typing 30 | 31 | # COMMAND_GENERATE_CHAT_CONCURRENT_LIMIT: limit of concurrent work-in-progress requests a single chat can send 32 | COMMAND_GENERATE_CHAT_CONCURRENT_LIMIT=3 33 | 34 | # DALLE_API_URL: complete URL to the DALLE API Generate endpoint 35 | DALLE_API_URL=https://bf.dallemini.ai/generate 36 | 37 | # DALLE_API_REQUEST_TIMEOUT_SECONDS: timeout for individual requests to DALLE API (the time it should take to complete a generation) 38 | DALLE_API_REQUEST_TIMEOUT_SECONDS=210 39 | 40 | # DALLE_API_REQUEST_SOCKS_PROXY: Socks proxy to use for DALLE API requests 41 | #DALLE_API_REQUEST_SOCKS_PROXY=socks5://localhost:8080 42 | 43 | # DALLE_GENERATION_TIMEOUT_SECONDS: timeout for trying to generate an image, including all the retries to the DALLE API 44 | DALLE_GENERATION_TIMEOUT_SECONDS=360 45 | 46 | # DALLE_GENERATION_RETRY_DELAY_SECONDS: delay between DALLE API retrying requests 47 | DALLE_GENERATION_RETRY_DELAY_SECONDS=5 48 | 49 | # REDIS_HOST: host/ip of Redis server; if not set, functionalities using Redis will be disabled 50 | #REDIS_HOST=localhost 51 | 52 | # REDIS_PORT: port of Redis server 53 | REDIS_PORT=6379 54 | 55 | # REDIS_DB: index of Redis DB to use 56 | REDIS_DB=0 57 | 58 | # Redis credentials (username and/or password) 59 | #REDIS_USERNAME= 60 | #REDIS_PASSWORD= 61 | 62 | # REDIS_LOGS_QUEUE_NAME: index name on Redis for the queue where log records will be pushed; if not set, no records will be sent to Redis 63 | REDIS_LOGS_QUEUE_NAME=dallemini-telegrambot/logs 64 | 65 | # LOG_LEVEL: one of: trace, debug, info, warning, error 66 | LOG_LEVEL=INFO 67 | --------------------------------------------------------------------------------