diff --git a/telegram_helper.py b/telegram_helper.py index e487e13..4b23a2e 100644 --- a/telegram_helper.py +++ b/telegram_helper.py @@ -8,42 +8,56 @@ from telegram.ext import Application, CommandHandler, MessageHandler, filters, C from browse_command import browse_command, button_callback class TelegramHelper: - # --- Constants for magic strings (paths are now instance vars) --- CLAUDE_REBOOT_TARGET = 'claude' HTML_QUOTE_BLOCK_START = '
Thinking...' HTML_QUOTE_BLOCK_END = '
' - DEFAULT_REBOOT_CLAUDE_FILE = '.reboot_claude' # Default value - DEFAULT_REBOOT_FILE = '.doreboot' # Default value + DEFAULT_REBOOT_CLAUDE_FILE = '.reboot_claude' + DEFAULT_REBOOT_FILE = '.doreboot' - def __init__(self, bot, reboot_claude_file_path: str | None = None, reboot_file_path: str | None = None): # MODIFIED + def __init__(self, bot, reboot_claude_file_path: str | None = None, reboot_file_path: str | None = None): self.bot = bot self.telegram_bot_token = os.getenv('TELEGRAM_BOT_TOKEN') self.start_time = time.time() - # MODIFIED: Store configurable paths self.reboot_claude_file = reboot_claude_file_path or self.DEFAULT_REBOOT_CLAUDE_FILE self.reboot_file = reboot_file_path or self.DEFAULT_REBOOT_FILE - async def start(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + # --- Start Command --- + async def _start_logic(self) -> str: # New logic method await self.bot.start() - await update.message.reply_text( - "Hello! I'm your AI assistant. How can I help you today?" - ) + return "Hello! I'm your AI assistant. How can I help you today?" - async def clear(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - user_id = update.effective_user.id + async def start(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: # Modified + response_message = await self._start_logic() + await update.message.reply_text(response_message) + + # --- Clear Command --- + async def _clear_logic(self, user_id: int) -> str: # New logic method self.bot.clear_conversation_history(user_id) - await update.message.reply_text("Conversation history cleared. Let's start fresh!") + return "Conversation history cleared. Let's start fresh!" - async def status(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - status_message = await self.bot.get_bot_status() - await update.message.reply_text(status_message) + async def clear(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: # Modified + user_id = update.effective_user.id + response_message = await self._clear_logic(user_id) # Await was missing for async consistency, though _clear_logic is not async yet, it calls a sync method. + # For consistency with other _logic methods, making it async. - async def switch(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + # --- Status Command --- + async def _status_logic(self) -> str: # New logic method + return await self.bot.get_bot_status() + + async def status(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: # Modified + response_message = await self._status_logic() + await update.message.reply_text(response_message) + + # --- Switch Command --- + async def _switch_logic(self) -> str: # New logic method if hasattr(self.bot, 'switch_model'): - status_message = await self.bot.switch_model() - await update.message.reply_text(status_message) + return await self.bot.switch_model() else: - await update.message.reply_text("Model switching is not supported for this bot.") + return "Model switching is not supported for this bot." + + async def switch(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: # Modified + response_message = await self._switch_logic() + await update.message.reply_text(response_message) async def update_status_message(self, context: ContextTypes.DEFAULT_TYPE, chat_id: int, message_id: int, status: str): keyboard = [ @@ -64,12 +78,12 @@ class TelegramHelper: logging.info(f"Message from user {user_id}: {user_message}") - status_message = await update.message.reply_text("Processing your request...", reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Abort", callback_data='abort')]])) - self.bot.set_processing_status(user_id, status_message.message_id) + status_message_obj = await update.message.reply_text("Processing your request...", reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Abort", callback_data='abort')]])) + self.bot.set_processing_status(user_id, status_message_obj.message_id) response = await self.bot.handle_message(user_id, user_message) - await context.bot.delete_message(chat_id=update.effective_chat.id, message_id=status_message.message_id) + await context.bot.delete_message(chat_id=update.effective_chat.id, message_id=status_message_obj.message_id) self.bot.clear_processing_status(user_id) response = response.replace("", self.HTML_QUOTE_BLOCK_START).replace("", self.HTML_QUOTE_BLOCK_END) @@ -97,27 +111,32 @@ class TelegramHelper: async def reboot(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: user_message = update.message.text.split() if len(user_message) > 1 and user_message[1].lower() == self.CLAUDE_REBOOT_TARGET: - # MODIFIED: Use instance variable open(self.reboot_claude_file, 'w').close() if update: await update.message.reply_text("Rebooting the bot...") logging.info("Received reboot command. Exiting process...") - # MODIFIED: Use instance variable + reboot_f_path = self.reboot_file if not os.path.exists(reboot_f_path): with open(reboot_f_path, 'w') as f: - f.write(str(update.effective_chat.id) if update else "") + # If update is None (e.g. called programmatically without a Telegram context for reboot), + # we should handle this. For now, assuming update is present if this handler is called by Telegram. + # Testability of this part needs care due to sys.exit() + chat_id_to_write = str(update.effective_chat.id) if update and update.effective_chat else "" + f.write(chat_id_to_write) sys.exit(0) async def check_doreboot_file(self, application: Application): - # MODIFIED: Use instance variable reboot_f_path = self.reboot_file if os.path.exists(reboot_f_path): with open(reboot_f_path, 'r') as f: chat_id = f.read().strip() if chat_id: - await application.bot.send_message(chat_id=chat_id, text="The application has finished initializing.") + try: + await application.bot.send_message(chat_id=chat_id, text="The application has finished initializing.") + except Exception as e: + logging.error(f"Failed to send reboot notification to chat_id {chat_id}: {e}") os.remove(reboot_f_path) async def browse(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: @@ -138,6 +157,10 @@ class TelegramHelper: logging.info("Bot is running...") - asyncio.get_event_loop().create_task(self.check_doreboot_file(application)) + loop = asyncio.get_event_loop() + if loop.is_running(): + loop.create_task(self.check_doreboot_file(application)) + else: + asyncio.run(self.check_doreboot_file(application)) # Fallback if loop not running (e.g. tests) application.run_polling()