Refactored gemini, openai and claude into one file and removed logic from the base class, also made helper class definable from command line
This commit is contained in:
+3
-86
@@ -7,6 +7,7 @@ from typing import TypedDict, Union, TypeAlias, List # Added List for type hint
|
||||
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
|
||||
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackQueryHandler
|
||||
from browse_command import browse_command, button_callback
|
||||
from inference_bot import InferenceBot
|
||||
|
||||
class MessageHandlerLogicResult(TypedDict):
|
||||
success: bool
|
||||
@@ -16,22 +17,15 @@ class MessageHandlerLogicResult(TypedDict):
|
||||
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 : InferenceBot,
|
||||
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
|
||||
|
||||
async def _start_logic(self) -> str:
|
||||
@@ -146,93 +140,16 @@ class TelegramHelper:
|
||||
response_text = await self._abort_processing_logic(user_id)
|
||||
await query.edit_message_text(text=response_text)
|
||||
|
||||
# --- Reboot Command ---
|
||||
def _reboot_logic(self, user_message_parts: List[str], chat_id_to_write: str) -> None:
|
||||
"""Handles the logic for creating reboot files."""
|
||||
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("") # Create/truncate the file
|
||||
logging.info(f"Created/truncated Claude reboot file: {self.reboot_claude_file}")
|
||||
except IOError as e:
|
||||
logging.error(f"Failed to create/truncate Claude reboot file {self.reboot_claude_file}: {e}")
|
||||
|
||||
# Create the main reboot file if it doesn't exist
|
||||
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} with chat_id.")
|
||||
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. Not overwriting chat_id.")
|
||||
|
||||
async def reboot(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
"""Handles the /reboot command, triggers file creation and exits."""
|
||||
user_message_parts = update.message.text.split()
|
||||
chat_id_str = str(update.effective_chat.id) if update and update.effective_chat else ""
|
||||
|
||||
self._reboot_logic(user_message_parts, chat_id_str)
|
||||
|
||||
if update:
|
||||
try:
|
||||
await update.message.reply_text("Rebooting the bot...")
|
||||
except Exception as e_reply:
|
||||
logging.error(f"Failed to send reboot reply: {e_reply}")
|
||||
|
||||
logging.info("Initiating shutdown for reboot...")
|
||||
sys.exit(0) # This part is not directly testable for completion in unit tests
|
||||
|
||||
# --- Check Doreboot File ---
|
||||
async def _check_doreboot_file_logic(self) -> Union[str, None]:
|
||||
"""Checks for the reboot file, reads chat_id, removes file, and returns chat_id."""
|
||||
if os.path.exists(self.reboot_file):
|
||||
chat_id = None
|
||||
try:
|
||||
with open(self.reboot_file, 'r') as f:
|
||||
chat_id = f.read().strip()
|
||||
# Attempt to remove the file after reading
|
||||
try:
|
||||
os.remove(self.reboot_file)
|
||||
logging.info(f"Successfully read and removed reboot file: {self.reboot_file}")
|
||||
except OSError as e_remove:
|
||||
logging.error(f"Failed to remove reboot file {self.reboot_file} after reading: {e_remove}")
|
||||
# Still return chat_id if read was successful, to attempt notification
|
||||
return chat_id
|
||||
except IOError as e_read:
|
||||
logging.error(f"Error reading reboot file {self.reboot_file}: {e_read}")
|
||||
# If reading failed, attempt to remove anyway if it exists, to prevent stale files
|
||||
if os.path.exists(self.reboot_file):
|
||||
try:
|
||||
os.remove(self.reboot_file)
|
||||
logging.warning(f"Removed reboot file {self.reboot_file} after a read error.")
|
||||
except OSError as e_remove_after_fail:
|
||||
logging.error(f"Failed to remove reboot file {self.reboot_file} even after a read error: {e_remove_after_fail}")
|
||||
return None # Reading failed
|
||||
return None # File does not exist
|
||||
|
||||
async def check_doreboot_file(self, application: Application) -> None:
|
||||
"""Checks for reboot file using logic method 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 initialization notification to chat_id: {chat_id}")
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to send reboot initialization notification to chat_id {chat_id}: {e}")
|
||||
|
||||
async def browse(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
await browse_command(update, context, self.bot)
|
||||
|
||||
def run(self):
|
||||
application = Application.builder().token(self.telegram_bot_token).post_init(self.check_doreboot_file).build()
|
||||
application = Application.builder().token(self.telegram_bot_token).build()
|
||||
|
||||
application.add_handler(CommandHandler("start", self.start))
|
||||
application.add_handler(CommandHandler("clear", self.clear))
|
||||
application.add_handler(CommandHandler("switch", self.switch))
|
||||
application.add_handler(CommandHandler("status", self.status))
|
||||
application.add_handler(CommandHandler("reboot", self.reboot))
|
||||
application.add_handler(CommandHandler("browse", self.browse))
|
||||
application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_message))
|
||||
application.add_handler(CallbackQueryHandler(self.abort_processing, pattern='^abort$'))
|
||||
|
||||
Reference in New Issue
Block a user