From 48e3948537eafbd6ea9ad57e47dc07001c528730 Mon Sep 17 00:00:00 2001 From: bucolucas Date: Sun, 18 Aug 2024 15:57:45 -0500 Subject: [PATCH] Add status updates and abort functionality --- telegram_inference_bot.py | 56 +++++++++++++++++++++++++++++++++++---- 1 file changed, 51 insertions(+), 5 deletions(-) diff --git a/telegram_inference_bot.py b/telegram_inference_bot.py index 70a2645..cbd0617 100644 --- a/telegram_inference_bot.py +++ b/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: @@ -145,7 +167,10 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> assistant_reply = assistant_message 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...")