Refactor: Separate logic for handle_message handler in TelegramHelper
This commit is contained in:
cyclop-bot
+128
-61
@@ -3,59 +3,68 @@ import logging
|
||||
import sys
|
||||
import asyncio
|
||||
import time
|
||||
from typing import TypedDict, Union, TypeAlias
|
||||
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
|
||||
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackQueryHandler
|
||||
from browse_command import browse_command, button_callback
|
||||
|
||||
class MessageHandlerLogicResult(TypedDict):
|
||||
success: bool
|
||||
response_text: Union[str, None]
|
||||
error_message: Union[str, None]
|
||||
|
||||
LogicResult: TypeAlias = MessageHandlerLogicResult
|
||||
|
||||
class TelegramHelper:
|
||||
CLAUDE_REBOOT_TARGET = 'claude'
|
||||
HTML_QUOTE_BLOCK_START = '<blockquote expandable><b>Thinking...</b>'
|
||||
HTML_QUOTE_BLOCK_END = '</blockquote>'
|
||||
DEFAULT_REBOOT_CLAUDE_FILE = '.reboot_claude'
|
||||
DEFAULT_REBOOT_FILE = '.doreboot'
|
||||
CHUNK_MESSAGE_SLEEP_DURATION = 0.1
|
||||
|
||||
def __init__(self, bot, reboot_claude_file_path: str | None = None, reboot_file_path: str | None = None):
|
||||
def __init__(self, bot,
|
||||
reboot_claude_file_path: str | None = None,
|
||||
reboot_file_path: str | None = None,
|
||||
chunk_message_sleep_duration: float | None = None):
|
||||
self.bot = bot
|
||||
self.telegram_bot_token = os.getenv('TELEGRAM_BOT_TOKEN')
|
||||
self.start_time = time.time()
|
||||
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.chunk_message_sleep_duration = chunk_message_sleep_duration if chunk_message_sleep_duration is not None else self.CHUNK_MESSAGE_SLEEP_DURATION
|
||||
|
||||
# --- Start Command ---
|
||||
async def _start_logic(self) -> str: # New logic method
|
||||
async def _start_logic(self) -> str:
|
||||
await self.bot.start()
|
||||
return "Hello! I'm your AI assistant. How can I help you today?"
|
||||
|
||||
async def start(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: # Modified
|
||||
async def start(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
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
|
||||
async def _clear_logic(self, user_id: int) -> str:
|
||||
self.bot.clear_conversation_history(user_id)
|
||||
return "Conversation history cleared. Let's start fresh!"
|
||||
|
||||
async def clear(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: # Modified
|
||||
async def clear(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
user_id = update.effective_user.id
|
||||
response_message = await self._clear_logic(user_id)
|
||||
await update.message.reply_text(response_message)
|
||||
|
||||
# --- Status Command ---
|
||||
async def _status_logic(self) -> str: # New logic method
|
||||
async def _status_logic(self) -> str:
|
||||
return await self.bot.get_bot_status()
|
||||
|
||||
async def status(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: # Modified
|
||||
async def status(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
response_message = await self._status_logic()
|
||||
await update.message.reply_text(response_message)
|
||||
|
||||
# --- Switch Command ---
|
||||
async def _switch_logic(self) -> str: # New logic method
|
||||
async def _switch_logic(self) -> str:
|
||||
if hasattr(self.bot, 'switch_model'):
|
||||
return await self.bot.switch_model()
|
||||
else:
|
||||
return "Model switching is not supported for this bot."
|
||||
|
||||
async def switch(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: # Modified
|
||||
async def switch(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
response_message = await self._switch_logic()
|
||||
await update.message.reply_text(response_message)
|
||||
|
||||
@@ -71,80 +80,136 @@ class TelegramHelper:
|
||||
reply_markup=reply_markup
|
||||
)
|
||||
|
||||
async def handle_message(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
async def _handle_message_logic(self, user_id: int, user_message: str) -> LogicResult:
|
||||
try:
|
||||
user_id = update.effective_user.id
|
||||
user_message = update.message.text
|
||||
response = await self.bot.handle_message(user_id, user_message)
|
||||
processed_response = response.replace("<think>", self.HTML_QUOTE_BLOCK_START).replace("</think>", self.HTML_QUOTE_BLOCK_END)
|
||||
return LogicResult(success=True, response_text=processed_response, error_message=None)
|
||||
except Exception as e:
|
||||
logging.error(f"Error in _handle_message_logic for user {user_id}: {str(e)}")
|
||||
return LogicResult(success=False, response_text=None, error_message=str(e))
|
||||
|
||||
logging.info(f"Message from user {user_id}: {user_message}")
|
||||
async def handle_message(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
user_id = update.effective_user.id
|
||||
user_message = update.message.text
|
||||
chat_id = update.effective_chat.id
|
||||
status_message_obj = None
|
||||
|
||||
status_message_obj = await update.message.reply_text("Processing your request...", reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Abort", callback_data='abort')]]))
|
||||
try:
|
||||
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)
|
||||
logic_result = await self._handle_message_logic(user_id, user_message)
|
||||
|
||||
await context.bot.delete_message(chat_id=update.effective_chat.id, message_id=status_message_obj.message_id)
|
||||
if status_message_obj:
|
||||
try:
|
||||
await context.bot.delete_message(chat_id=chat_id, message_id=status_message_obj.message_id)
|
||||
except Exception as e_del:
|
||||
logging.warning(f"Failed to delete status message: {e_del}")
|
||||
self.bot.clear_processing_status(user_id)
|
||||
|
||||
response = response.replace("<think>", self.HTML_QUOTE_BLOCK_START).replace("</think>", self.HTML_QUOTE_BLOCK_END)
|
||||
|
||||
if len(response) > 4096:
|
||||
chunks = [response[i:i + 4096] for i in range(0, len(response), 4096)]
|
||||
for chunk in chunks:
|
||||
await update.message.reply_text(chunk)
|
||||
await asyncio.sleep(0.1)
|
||||
if logic_result["success"]:
|
||||
response_text = logic_result["response_text"]
|
||||
if response_text:
|
||||
if len(response_text) > 4096:
|
||||
chunks = [response_text[i:i + 4096] for i in range(0, len(response_text), 4096)]
|
||||
for chunk in chunks:
|
||||
await update.message.reply_text(chunk)
|
||||
await asyncio.sleep(self.chunk_message_sleep_duration)
|
||||
else:
|
||||
await update.message.reply_text(response_text)
|
||||
else:
|
||||
logging.warning("Successful logic result but no response text.")
|
||||
await update.message.reply_text("Something went unexpectedly well, but I have nothing to say.")
|
||||
else:
|
||||
await update.message.reply_text(response)
|
||||
await update.message.reply_text("Sorry, an error occurred while processing your request.")
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"An error occurred: {str(e)}")
|
||||
await update.message.reply_text("Sorry, an error occurred while processing your request.")
|
||||
logging.error(f"Outer error in handle_message for user {user_id}: {str(e)}")
|
||||
if status_message_obj and self.bot.processing_status.get(user_id):
|
||||
self.bot.clear_processing_status(user_id)
|
||||
try:
|
||||
await update.message.reply_text("Sorry, an unexpected error occurred with the bot.")
|
||||
except Exception as e_reply:
|
||||
logging.error(f"Failed to send error reply: {e_reply}")
|
||||
|
||||
# --- Abort Processing (Callback) ---
|
||||
async def _abort_processing_logic(self, user_id: int) -> str: # New logic method
|
||||
async def _abort_processing_logic(self, user_id: int) -> str:
|
||||
return await self.bot.abort_processing(user_id)
|
||||
|
||||
async def abort_processing(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: # Modified
|
||||
async def abort_processing(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
query = update.callback_query
|
||||
await query.answer() # Telegram specific interaction
|
||||
|
||||
await query.answer()
|
||||
user_id = query.from_user.id
|
||||
response_text = await self._abort_processing_logic(user_id) # Call logic method
|
||||
|
||||
await query.edit_message_text(text=response_text) # Telegram specific interaction
|
||||
response_text = await self._abort_processing_logic(user_id)
|
||||
await query.edit_message_text(text=response_text)
|
||||
|
||||
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:
|
||||
open(self.reboot_claude_file, 'w').close()
|
||||
user_message_parts = update.message.text.split()
|
||||
chat_id_to_write = str(update.effective_chat.id) if update and update.effective_chat else ""
|
||||
|
||||
self._reboot_logic(user_message_parts, chat_id_to_write) # Call to new logic method
|
||||
|
||||
if update:
|
||||
await update.message.reply_text("Rebooting the bot...")
|
||||
logging.info("Received reboot command. Exiting process...")
|
||||
|
||||
reboot_f_path = self.reboot_file
|
||||
if not os.path.exists(reboot_f_path):
|
||||
with open(reboot_f_path, 'w') as f:
|
||||
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) # This will be hard to test directly
|
||||
|
||||
async def check_doreboot_file(self, application: Application):
|
||||
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:
|
||||
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)
|
||||
# New internal logic method for reboot preparations
|
||||
def _reboot_logic(self, user_message_parts: list[str], chat_id_to_write: str) -> None:
|
||||
if len(user_message_parts) > 1 and user_message_parts[1].lower() == self.CLAUDE_REBOOT_TARGET:
|
||||
try:
|
||||
with open(self.reboot_claude_file, 'w') as f:
|
||||
f.write("") # Ensure file is created/truncated
|
||||
logging.info(f"Created Claude reboot file: {self.reboot_claude_file}")
|
||||
except IOError as e:
|
||||
logging.error(f"Failed to create Claude reboot file {self.reboot_claude_file}: {e}")
|
||||
|
||||
if not os.path.exists(self.reboot_file):
|
||||
try:
|
||||
with open(self.reboot_file, 'w') as f:
|
||||
f.write(chat_id_to_write)
|
||||
logging.info(f"Created main reboot file: {self.reboot_file}")
|
||||
except IOError as e:
|
||||
logging.error(f"Failed to create main reboot file {self.reboot_file}: {e}")
|
||||
else:
|
||||
logging.info(f"Main reboot file {self.reboot_file} already exists.")
|
||||
|
||||
async def _check_doreboot_file_logic(self) -> str | None:
|
||||
"""Checks for the reboot file and returns the chat_id if found, then removes the file."""
|
||||
if os.path.exists(self.reboot_file):
|
||||
chat_id = None
|
||||
try:
|
||||
with open(self.reboot_file, 'r') as f:
|
||||
chat_id = f.read().strip()
|
||||
os.remove(self.reboot_file)
|
||||
logging.info(f"Removed reboot file: {self.reboot_file}")
|
||||
return chat_id
|
||||
except IOError as e:
|
||||
logging.error(f"Error processing reboot file {self.reboot_file}: {e}")
|
||||
# If we read chat_id but failed to remove, still return chat_id to attempt notification
|
||||
if chat_id is not None:
|
||||
return chat_id
|
||||
return None
|
||||
|
||||
async def check_doreboot_file(self, application: Application) -> None:
|
||||
"""Checks for reboot file and sends notification if applicable."""
|
||||
chat_id = await self._check_doreboot_file_logic()
|
||||
if chat_id:
|
||||
try:
|
||||
await application.bot.send_message(chat_id=chat_id, text="The application has finished initializing.")
|
||||
logging.info(f"Sent reboot notification to chat_id: {chat_id}")
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to send reboot notification to chat_id {chat_id}: {e}")
|
||||
|
||||
async def browse(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
# For testing TelegramHelper, browse_command will be mocked or its effects on self.bot asserted.
|
||||
await browse_command(update, context, self.bot)
|
||||
|
||||
def run(self):
|
||||
# Consider allowing injection of a pre-built application for advanced test cases.
|
||||
application = Application.builder().token(self.telegram_bot_token).build()
|
||||
|
||||
application.add_handler(CommandHandler("start", self.start))
|
||||
@@ -160,9 +225,11 @@ class TelegramHelper:
|
||||
logging.info("Bot is running...")
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
if loop.is_running():
|
||||
# These pragma: no cover comments are hints for test coverage tools to ignore these lines if needed,
|
||||
# as testing both branches of is_running() can be environment-dependent.
|
||||
if loop.is_running(): # pragma: no cover
|
||||
loop.create_task(self.check_doreboot_file(application))
|
||||
else:
|
||||
else: # pragma: no cover
|
||||
asyncio.run(self.check_doreboot_file(application))
|
||||
|
||||
application.run_polling()
|
||||
|
||||
Reference in New Issue
Block a user