From 478081ff6430e2e2be9afa961b43e806a484923c Mon Sep 17 00:00:00 2001
From: bucolucas <bucolucas@gmail.com>
Date: Mon, 19 Aug 2024 11:35:29 -0500
Subject: [PATCH] Add TelegramHelper class

---
 telegram_helper.py | 66 ++++++++++++++++++++++++++++------------------
 1 file changed, 41 insertions(+), 25 deletions(-)

diff --git a/telegram_helper.py b/telegram_helper.py
index 8595b7a..7c095a6 100644
--- a/telegram_helper.py
+++ b/telegram_helper.py
@@ -1,48 +1,34 @@
 import os
+import logging
 from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
 from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackQueryHandler
 
 class TelegramHelper:
     def __init__(self, bot):
         self.bot = bot
-        self.TELEGRAM_BOT_TOKEN = os.getenv('TELEGRAM_BOT_TOKEN')
+        self.telegram_bot_token = os.getenv('TELEGRAM_BOT_TOKEN')
 
     async def start(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
         await self.bot.start()
         await update.message.reply_text(
-            "Hello! I'm your AI assistant. How can I help you today? You can send me images and then ask questions about them."
+            "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
         await self.bot.clear(user_id)
-        await update.message.reply_text("Conversation history and image cleared. Let's start fresh!")
+        await update.message.reply_text("Conversation history cleared. Let's start fresh!")
 
     async def status(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
         status_message = await self.bot.status()
         await update.message.reply_text(status_message)
 
-    async def handle_message(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
-        user_id = update.effective_user.id
-        user_message = update.message.text
-
-        status_message = await update.message.reply_text(
-            "Processing your request...",
-            reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Abort", callback_data='abort')]])
-        )
-
-        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 update.message.reply_text(response)
-
-    async def abort_processing(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
-        query = update.callback_query
-        await query.answer()
-        
-        user_id = query.from_user.id
-        abort_message = await self.bot.abort_processing(user_id)
-        await query.edit_message_text(text=abort_message)
+    async def switch(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+        if hasattr(self.bot, 'switch_model'):
+            status_message = await self.bot.switch_model()
+            await update.message.reply_text(status_message)
+        else:
+            await update.message.reply_text("Model switching is not supported for this bot.")
 
     async def update_status_message(self, context: ContextTypes.DEFAULT_TYPE, chat_id: int, message_id: int, status: str):
         keyboard = [
@@ -56,13 +42,43 @@ class TelegramHelper:
             reply_markup=reply_markup
         )
 
+    async def handle_message(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+        try:
+            user_id = update.effective_user.id
+            user_message = update.message.text
+
+            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.processing_status[user_id] = {"processing": True, "message_id": status_message.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)
+            del self.bot.processing_status[user_id]
+            await update.message.reply_text(response)
+
+        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.")
+
+    async def abort_processing(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+        query = update.callback_query
+        await query.answer()
+        
+        user_id = query.from_user.id
+        result = await self.bot.abort_processing(user_id)
+        await query.edit_message_text(text=result)
+
     def run(self):
-        application = Application.builder().token(self.TELEGRAM_BOT_TOKEN).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(MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_message))
         application.add_handler(CallbackQueryHandler(self.abort_processing, pattern='^abort$'))
 
+        logging.info("Bot is running...")
         application.run_polling()
\ No newline at end of file