Merge pull request #48 from bucolucas/add-status-updates

Add status updates and abort functionality to Telegram bot

telegram_inference_bot.py

@@ -3,8 +3,9 @@ import os
 import importlib
 import inspect
 import logging
-from telegram import Update, __version__ as telegram_version
-from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes
+import asyncio
+from telegram import Update, __version__ as telegram_version, InlineKeyboardButton, InlineKeyboardMarkup
+from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackQueryHandler
 from openai import OpenAI
 from dotenv import load_dotenv
 from tools.base_tool import BaseTool
@@ -47,6 +48,9 @@ with open("prompts/developer_prompt.txt", "r") as file:
 # Dictionary to store conversation history for each user
 conversation_history = {}
 
+# Dictionary to store processing status for each user
+processing_status = {}
+
 # Load tools
 tools = []
 tools_dir = os.path.join(os.path.dirname(__file__), 'tools')
@@ -77,6 +81,18 @@ async def clear(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
     logging.info(f"Cleared conversation history and image for user {user_id}")
     await update.message.reply_text("Conversation history and image cleared. Let's start fresh!")
 
+async def update_status_message(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
+    )
+
 async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
     try:
         user_id = update.effective_user.id
@@ -89,6 +105,10 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
 
         conversation_history[user_id].append({"role": "user", "content": user_message})
 
+        # Send initial status message
+        status_message = await update.message.reply_text("Processing your request...", reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Abort", callback_data='abort')]]))
+        processing_status[user_id] = {"processing": True, "message_id": status_message.message_id}
+
         messages = conversation_history[user_id]
 
         response = get_chat_response(messages)
@@ -107,16 +127,18 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
 
         toolUseCount = 0
 
-        while len(tool_calls) > 0 and toolUseCount < 50:
+        while len(tool_calls) > 0 and toolUseCount < 50 and processing_status[user_id]["processing"]:
             tool_use_results = []
             while len(tool_calls) > 0:
                 tool_call = tool_calls.pop(0)
                 function_name = tool_call.name
                 
+                # Update status message
+                await update_status_message(context, update.effective_chat.id, status_message.message_id, f"Using tool: {function_name}")
+                
                 tool_response = call_tool(tool_call)
                 tool_use_results.append({"type": "tool_result", "tool_use_id": tool_call.id, "content": json.dumps(tool_response)})
 
-            
             formatted_result = {}
 
             if use_anthropic:
@@ -146,6 +168,9 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
 
             toolUseCount += 1
             
+            # Update status message every 5 seconds
+            if toolUseCount % 5 == 0:
+                await update_status_message(context, update.effective_chat.id, status_message.message_id, f"Tools used: {toolUseCount}")
 
         if (toolUseCount == 0):
             if use_anthropic:
@@ -157,6 +182,10 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) ->
         if len(conversation_history[user_id]) > 20:
             conversation_history[user_id] = conversation_history[user_id][-20:]
 
+        # Remove the status message
+        await context.bot.delete_message(chat_id=update.effective_chat.id, message_id=status_message.message_id)
+        del processing_status[user_id]
+
         if use_anthropic:
             await update.message.reply_text(messages[-1]["content"][0].text)
         else:
@@ -231,6 +260,22 @@ async def status(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
         model = GPT_4O if use_smart_model else GPT_4O_MINI
         await update.message.reply_text(f"Currently using: {model}")
 
+async def abort_processing(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+    query = update.callback_query
+    await query.answer()
+    
+    user_id = query.from_user.id
+    if user_id in processing_status:
+        processing_status[user_id]["processing"] = False
+        await context.bot.edit_message_text(
+            chat_id=query.message.chat_id,
+            message_id=query.message.message_id,
+            text="Processing aborted."
+        )
+        await clear(update, context)
+    else:
+        await query.edit_message_text(text="No active processing to abort.")
+
 def main() -> None:
     # Create the Application and pass it your bot's token
     application = Application.builder().token(TELEGRAM_BOT_TOKEN).build()
@@ -242,6 +287,7 @@ def main() -> None:
     application.add_handler(CommandHandler("toggle", switch_providers))
     application.add_handler(CommandHandler("status", status))
     application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))
+    application.add_handler(CallbackQueryHandler(abort_processing, pattern='^abort$'))
 
     # Start the Bot
     logging.info("Bot is running...")