Refactor: Separate logic for start, clear, status, switch handlers in TelegramHelper
This commit is contained in:
+51
-28
@@ -8,42 +8,56 @@ from telegram.ext import Application, CommandHandler, MessageHandler, filters, C
|
|||||||
from browse_command import browse_command, button_callback
|
from browse_command import browse_command, button_callback
|
||||||
|
|
||||||
class TelegramHelper:
|
class TelegramHelper:
|
||||||
# --- Constants for magic strings (paths are now instance vars) ---
|
|
||||||
CLAUDE_REBOOT_TARGET = 'claude'
|
CLAUDE_REBOOT_TARGET = 'claude'
|
||||||
HTML_QUOTE_BLOCK_START = '<blockquote expandable><b>Thinking...</b>'
|
HTML_QUOTE_BLOCK_START = '<blockquote expandable><b>Thinking...</b>'
|
||||||
HTML_QUOTE_BLOCK_END = '</blockquote>'
|
HTML_QUOTE_BLOCK_END = '</blockquote>'
|
||||||
DEFAULT_REBOOT_CLAUDE_FILE = '.reboot_claude' # Default value
|
DEFAULT_REBOOT_CLAUDE_FILE = '.reboot_claude'
|
||||||
DEFAULT_REBOOT_FILE = '.doreboot' # Default value
|
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.bot = bot
|
||||||
self.telegram_bot_token = os.getenv('TELEGRAM_BOT_TOKEN')
|
self.telegram_bot_token = os.getenv('TELEGRAM_BOT_TOKEN')
|
||||||
self.start_time = time.time()
|
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_claude_file = reboot_claude_file_path or self.DEFAULT_REBOOT_CLAUDE_FILE
|
||||||
self.reboot_file = reboot_file_path or self.DEFAULT_REBOOT_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 self.bot.start()
|
||||||
await update.message.reply_text(
|
return "Hello! I'm your AI assistant. How can I help you today?"
|
||||||
"Hello! I'm your AI assistant. How can I help you today?"
|
|
||||||
)
|
|
||||||
|
|
||||||
async def clear(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
async def start(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: # Modified
|
||||||
user_id = update.effective_user.id
|
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)
|
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:
|
async def clear(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: # Modified
|
||||||
status_message = await self.bot.get_bot_status()
|
user_id = update.effective_user.id
|
||||||
await update.message.reply_text(status_message)
|
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'):
|
if hasattr(self.bot, 'switch_model'):
|
||||||
status_message = await self.bot.switch_model()
|
return await self.bot.switch_model()
|
||||||
await update.message.reply_text(status_message)
|
|
||||||
else:
|
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):
|
async def update_status_message(self, context: ContextTypes.DEFAULT_TYPE, chat_id: int, message_id: int, status: str):
|
||||||
keyboard = [
|
keyboard = [
|
||||||
@@ -64,12 +78,12 @@ class TelegramHelper:
|
|||||||
|
|
||||||
logging.info(f"Message from user {user_id}: {user_message}")
|
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')]]))
|
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.message_id)
|
self.bot.set_processing_status(user_id, status_message_obj.message_id)
|
||||||
|
|
||||||
response = await self.bot.handle_message(user_id, user_message)
|
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)
|
self.bot.clear_processing_status(user_id)
|
||||||
|
|
||||||
response = response.replace("<think>", self.HTML_QUOTE_BLOCK_START).replace("</think>", self.HTML_QUOTE_BLOCK_END)
|
response = response.replace("<think>", self.HTML_QUOTE_BLOCK_START).replace("</think>", self.HTML_QUOTE_BLOCK_END)
|
||||||
@@ -97,27 +111,32 @@ class TelegramHelper:
|
|||||||
async def reboot(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
async def reboot(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||||
user_message = update.message.text.split()
|
user_message = update.message.text.split()
|
||||||
if len(user_message) > 1 and user_message[1].lower() == self.CLAUDE_REBOOT_TARGET:
|
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()
|
open(self.reboot_claude_file, 'w').close()
|
||||||
|
|
||||||
if update:
|
if update:
|
||||||
await update.message.reply_text("Rebooting the bot...")
|
await update.message.reply_text("Rebooting the bot...")
|
||||||
logging.info("Received reboot command. Exiting process...")
|
logging.info("Received reboot command. Exiting process...")
|
||||||
# MODIFIED: Use instance variable
|
|
||||||
reboot_f_path = self.reboot_file
|
reboot_f_path = self.reboot_file
|
||||||
if not os.path.exists(reboot_f_path):
|
if not os.path.exists(reboot_f_path):
|
||||||
with open(reboot_f_path, 'w') as f:
|
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)
|
sys.exit(0)
|
||||||
|
|
||||||
async def check_doreboot_file(self, application: Application):
|
async def check_doreboot_file(self, application: Application):
|
||||||
# MODIFIED: Use instance variable
|
|
||||||
reboot_f_path = self.reboot_file
|
reboot_f_path = self.reboot_file
|
||||||
if os.path.exists(reboot_f_path):
|
if os.path.exists(reboot_f_path):
|
||||||
with open(reboot_f_path, 'r') as f:
|
with open(reboot_f_path, 'r') as f:
|
||||||
chat_id = f.read().strip()
|
chat_id = f.read().strip()
|
||||||
if chat_id:
|
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)
|
os.remove(reboot_f_path)
|
||||||
|
|
||||||
async def browse(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
async def browse(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||||
@@ -138,6 +157,10 @@ class TelegramHelper:
|
|||||||
|
|
||||||
logging.info("Bot is running...")
|
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()
|
application.run_polling()
|
||||||
|
|||||||
Reference in New Issue
Block a user