diff --git a/bot/bot.py b/bot/bot.py index 64f132337..0f188bdb5 100644 --- a/bot/bot.py +++ b/bot/bot.py @@ -33,6 +33,8 @@ import base64 +from constants import Commands + # setup db = database.Database() logger = logging.getLogger(__name__) @@ -40,16 +42,16 @@ user_semaphores = {} user_tasks = {} -HELP_MESSAGE = """Commands: -⚪ /retry – Regenerate last bot answer -⚪ /new – Start new dialog -⚪ /mode – Select chat mode -⚪ /settings – Show settings -⚪ /balance – Show balance -⚪ /help – Show help +HELP_MESSAGE = f"""Commands: +⚪ /{Commands.RETRY} – Regenerate last bot answer +⚪ /{Commands.NEW} – Start new dialog +⚪ /{Commands.MODE} – Select chat mode +⚪ /{Commands.SETTINGS} – Show settings +⚪ /{Commands.BALANCE} – Show balance +⚪ /{Commands.HELP} – Show help -🎨 Generate images from text prompts in 👩‍🎨 Artist /mode -👥 Add bot to group chat: /help_group_chat +🎨 Generate images from text prompts in 👩‍🎨 Artist /{Commands.MODE} +👥 Add bot to group chat: /{Commands.HELP_GROUP_CHAT} 🎤 You can send Voice Messages instead of text """ @@ -111,22 +113,22 @@ async def register_user_if_not_exists(update: Update, context: CallbackContext, async def is_bot_mentioned(update: Update, context: CallbackContext): - try: - message = update.message + try: + message = update.message - if message.chat.type == "private": - return True + if message.chat.type == "private": + return True - if message.text is not None and ("@" + context.bot.username) in message.text: - return True + if message.text is not None and ("@" + context.bot.username) in message.text: + return True - if message.reply_to_message is not None: - if message.reply_to_message.from_user.id == context.bot.id: - return True - except: - return True - else: - return False + if message.reply_to_message is not None: + if message.reply_to_message.from_user.id == context.bot.id: + return True + except: + return True + else: + return False async def start_handle(update: Update, context: CallbackContext): @@ -151,14 +153,14 @@ async def help_handle(update: Update, context: CallbackContext): async def help_group_chat_handle(update: Update, context: CallbackContext): - await register_user_if_not_exists(update, context, update.message.from_user) - user_id = update.message.from_user.id - db.set_user_attribute(user_id, "last_interaction", datetime.now()) + await register_user_if_not_exists(update, context, update.message.from_user) + user_id = update.message.from_user.id + db.set_user_attribute(user_id, "last_interaction", datetime.now()) - text = HELP_GROUP_CHAT_MESSAGE.format(bot_username="@" + context.bot.username) + text = HELP_GROUP_CHAT_MESSAGE.format(bot_username="@" + context.bot.username) - await update.message.reply_text(text, parse_mode=ParseMode.HTML) - await update.message.reply_video(config.help_group_chat_video_path) + await update.message.reply_text(text, parse_mode=ParseMode.HTML) + await update.message.reply_video(config.help_group_chat_video_path) async def retry_handle(update: Update, context: CallbackContext): @@ -187,7 +189,7 @@ async def _vision_message_handle_fn( if current_model != "gpt-4-vision-preview": await update.message.reply_text( - "🥲 Images processing is only available for gpt-4-vision-preview model. Please change your settings in /settings", + f"🥲 Images processing is only available for gpt-4-vision-preview model. Please change your settings in /{Commands.SETTINGS}", parse_mode=ParseMode.HTML, ) return @@ -462,7 +464,7 @@ async def fake_gen(): if n_first_dialog_messages_removed == 1: text = "✍️ Note: Your current dialog is too long, so your first message was removed from the context.\n Send /new command to start new dialog" else: - text = f"✍️ Note: Your current dialog is too long, so {n_first_dialog_messages_removed} first messages were removed from the context.\n Send /new command to start new dialog" + text = f"✍️ Note: Your current dialog is too long, so {n_first_dialog_messages_removed} first messages were removed from the context.\n Send /{Commands.NEW} command to start new dialog" await update.message.reply_text(text, parse_mode=ParseMode.HTML) async with user_semaphores[user_id]: @@ -498,7 +500,7 @@ async def is_previous_message_not_answered_yet(update: Update, context: Callback user_id = update.message.from_user.id if user_semaphores[user_id].locked(): text = "⏳ Please wait for a reply to the previous message\n" - text += "Or you can /cancel it" + text += f"Or you can /{Commands.CANCEL} it" await update.message.reply_text(text, reply_to_message_id=update.message.id, parse_mode=ParseMode.HTML) return True else: @@ -641,25 +643,25 @@ async def show_chat_modes_handle(update: Update, context: CallbackContext): async def show_chat_modes_callback_handle(update: Update, context: CallbackContext): - await register_user_if_not_exists(update.callback_query, context, update.callback_query.from_user) - if await is_previous_message_not_answered_yet(update.callback_query, context): return + await register_user_if_not_exists(update.callback_query, context, update.callback_query.from_user) + if await is_previous_message_not_answered_yet(update.callback_query, context): return - user_id = update.callback_query.from_user.id - db.set_user_attribute(user_id, "last_interaction", datetime.now()) + user_id = update.callback_query.from_user.id + db.set_user_attribute(user_id, "last_interaction", datetime.now()) - query = update.callback_query - await query.answer() + query = update.callback_query + await query.answer() - page_index = int(query.data.split("|")[1]) - if page_index < 0: - return + page_index = int(query.data.split("|")[1]) + if page_index < 0: + return - text, reply_markup = get_chat_mode_menu(page_index) - try: - await query.edit_message_text(text, reply_markup=reply_markup, parse_mode=ParseMode.HTML) - except telegram.error.BadRequest as e: - if str(e).startswith("Message is not modified"): - pass + text, reply_markup = get_chat_mode_menu(page_index) + try: + await query.edit_message_text(text, reply_markup=reply_markup, parse_mode=ParseMode.HTML) + except telegram.error.BadRequest as e: + if str(e).startswith("Message is not modified"): + pass async def set_chat_mode_handle(update: Update, context: CallbackContext): @@ -817,12 +819,12 @@ async def error_handle(update: Update, context: CallbackContext) -> None: async def post_init(application: Application): await application.bot.set_my_commands([ - BotCommand("/new", "Start new dialog"), - BotCommand("/mode", "Select chat mode"), - BotCommand("/retry", "Re-generate response for previous query"), - BotCommand("/balance", "Show balance"), - BotCommand("/settings", "Show settings"), - BotCommand("/help", "Show help message"), + BotCommand(f'{Commands.NEW}', "Start new dialog"), + BotCommand(f'{Commands.MODE}', "Select chat mode"), + BotCommand(f'{Commands.RETRY}', "Re-generate response for previous query"), + BotCommand(f'{Commands.BALANCE}', "Show balance"), + BotCommand(f'{Commands.SETTINGS}', "Show settings"), + BotCommand(f'{Commands.HELP}', "Show help message"), ]) def run_bot() -> None: @@ -846,28 +848,28 @@ def run_bot() -> None: group_ids = [x for x in any_ids if x < 0] user_filter = filters.User(username=usernames) | filters.User(user_id=user_ids) | filters.Chat(chat_id=group_ids) - application.add_handler(CommandHandler("start", start_handle, filters=user_filter)) - application.add_handler(CommandHandler("help", help_handle, filters=user_filter)) - application.add_handler(CommandHandler("help_group_chat", help_group_chat_handle, filters=user_filter)) + application.add_handler(CommandHandler(Commands.START, start_handle, filters=user_filter)) + application.add_handler(CommandHandler(Commands.HELP, help_handle, filters=user_filter)) + application.add_handler(CommandHandler(Commands.HELP_GROUP_CHAT, help_group_chat_handle, filters=user_filter)) application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND & user_filter, message_handle)) application.add_handler(MessageHandler(filters.PHOTO & ~filters.COMMAND & user_filter, message_handle)) application.add_handler(MessageHandler(filters.VIDEO & ~filters.COMMAND & user_filter, unsupport_message_handle)) application.add_handler(MessageHandler(filters.Document.ALL & ~filters.COMMAND & user_filter, unsupport_message_handle)) - application.add_handler(CommandHandler("retry", retry_handle, filters=user_filter)) - application.add_handler(CommandHandler("new", new_dialog_handle, filters=user_filter)) - application.add_handler(CommandHandler("cancel", cancel_handle, filters=user_filter)) + application.add_handler(CommandHandler(Commands.RETRY, retry_handle, filters=user_filter)) + application.add_handler(CommandHandler(Commands.NEW, new_dialog_handle, filters=user_filter)) + application.add_handler(CommandHandler(Commands.CANCEL, cancel_handle, filters=user_filter)) application.add_handler(MessageHandler(filters.VOICE & user_filter, voice_message_handle)) - application.add_handler(CommandHandler("mode", show_chat_modes_handle, filters=user_filter)) + application.add_handler(CommandHandler(Commands.MODE, show_chat_modes_handle, filters=user_filter)) application.add_handler(CallbackQueryHandler(show_chat_modes_callback_handle, pattern="^show_chat_modes")) application.add_handler(CallbackQueryHandler(set_chat_mode_handle, pattern="^set_chat_mode")) - application.add_handler(CommandHandler("settings", settings_handle, filters=user_filter)) + application.add_handler(CommandHandler(Commands.SETTINGS, settings_handle, filters=user_filter)) application.add_handler(CallbackQueryHandler(set_settings_handle, pattern="^set_settings")) - application.add_handler(CommandHandler("balance", show_balance_handle, filters=user_filter)) + application.add_handler(CommandHandler(Commands.BALANCE, show_balance_handle, filters=user_filter)) application.add_error_handler(error_handle) @@ -876,4 +878,4 @@ def run_bot() -> None: if __name__ == "__main__": - run_bot() \ No newline at end of file + run_bot() diff --git a/bot/constants.py b/bot/constants.py new file mode 100644 index 000000000..68a2d9dc7 --- /dev/null +++ b/bot/constants.py @@ -0,0 +1,19 @@ +from enum import Enum + + +class Commands(str, Enum): + START = 'start' + HELP = 'help' + HELP_GROUP_CHAT = 'help_group_chat' + RETRY = 'retry' + NEW = 'new' + CANCEL = 'cancel' + MODE = 'mode' + SETTINGS = 'settings' + BALANCE = 'balance' + + def __str__(self): + return self.name.lower() + + def __repr__(self): + return str(self)