Add TelegramHelper class

This commit is contained in:
2024-08-19 11:35:29 -05:00
parent 9b16ca0d85
commit 478081ff64
+41 -25
View File
@@ -1,48 +1,34 @@
import os import os
import logging
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackQueryHandler from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackQueryHandler
class TelegramHelper: class TelegramHelper:
def __init__(self, bot): def __init__(self, bot):
self.bot = 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: async def start(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
await self.bot.start() await self.bot.start()
await update.message.reply_text( 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: async def clear(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
user_id = update.effective_user.id user_id = update.effective_user.id
await self.bot.clear(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: async def status(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
status_message = await self.bot.status() status_message = await self.bot.status()
await update.message.reply_text(status_message) await update.message.reply_text(status_message)
async def handle_message(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: async def switch(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
user_id = update.effective_user.id if hasattr(self.bot, 'switch_model'):
user_message = update.message.text status_message = await self.bot.switch_model()
await update.message.reply_text(status_message)
status_message = await update.message.reply_text( else:
"Processing your request...", await update.message.reply_text("Model switching is not supported for this bot.")
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): async def update_status_message(self, context: ContextTypes.DEFAULT_TYPE, chat_id: int, message_id: int, status: str):
keyboard = [ keyboard = [
@@ -56,13 +42,43 @@ class TelegramHelper:
reply_markup=reply_markup 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): 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("start", self.start))
application.add_handler(CommandHandler("clear", self.clear)) 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("status", self.status))
application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_message)) application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_message))
application.add_handler(CallbackQueryHandler(self.abort_processing, pattern='^abort$')) application.add_handler(CallbackQueryHandler(self.abort_processing, pattern='^abort$'))
logging.info("Bot is running...")
application.run_polling() application.run_polling()