Add telegram_helper.py with TelegramHelper class

2024-08-19 11:31:06 -05:00
parent f7d70a927f
commit 0ece8b354f
1 changed files with 68 additions and 0 deletions
@@ -0,0 +1,68 @@
+import os
+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')
+
+    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."
+        )
+
+    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!")
+
+    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 update_status_message(self, context: ContextTypes.DEFAULT_TYPE, chat_id: int, message_id: int, status: str):
+        keyboard = [
+            [InlineKeyboardButton("Abort", callback_data='abort')]
+        ]
+        reply_markup = InlineKeyboardMarkup(keyboard)
+        await context.bot.edit_message_text(
+            chat_id=chat_id,
+            message_id=message_id,
+            text=f"Current status: {status}",
+            reply_markup=reply_markup
+        )
+
+    def run(self):
+        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("status", self.status))
+        application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_message))
+        application.add_handler(CallbackQueryHandler(self.abort_processing, pattern='^abort$'))
+
+        application.run_polling()