diff --git a/anthropic_telegram_inference_bot.py b/anthropic_telegram_inference_bot.py index bb8c5c2..16d5344 100644 --- a/anthropic_telegram_inference_bot.py +++ b/anthropic_telegram_inference_bot.py @@ -1,232 +1,104 @@ -import json import os -import importlib -import inspect +import json import logging -import asyncio -from telegram import error as TelegramErrors, Update, __version__ as telegram_version, InlineKeyboardButton, InlineKeyboardMarkup -from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackQueryHandler -from dotenv import load_dotenv -from tools.base_tool import BaseTool -from tools.metrics_tool import MetricsTool from anthropic import Anthropic +from base_telegram_inference_bot import BaseTelegramInferenceBot +from telegram_helper import TelegramHelper -# Load environment variables -load_dotenv() +class AnthropicTelegramInferenceBot(BaseTelegramInferenceBot): + def __init__(self): + super().__init__() + self.anthropic_client = Anthropic( + api_key=os.environ.get("ANTHROPIC_API_KEY"), + default_headers={"anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15"} + ) -anthropic_client = Anthropic( - api_key=os.environ.get("ANTHROPIC_API_KEY"), - default_headers={"anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15"} -) - -# Set up logging to console and file -logging.basicConfig(level=logging.WARNING, handlers=[ - logging.StreamHandler(), - logging.FileHandler('logs/output.log', mode='a') -]) - -# Set up Telegram bot -TELEGRAM_BOT_TOKEN = os.getenv('TELEGRAM_BOT_TOKEN') - -# Load system prompt -with open("prompts/developer_prompt.txt", "r") as file: - system_prompt = file.read().strip() - -# Dictionary to store conversation history for each user -conversation_history = {} - -# Dictionary to store processing status for each user -processing_status = {} - -# Load tools -tools = [MetricsTool()] # Add MetricsTool instance -tools_dir = os.path.join(os.path.dirname(__file__), 'tools') -for filename in os.listdir(tools_dir): - if filename.endswith('.py') and filename not in ['__init__.py', 'base_tool.py', 'metrics_tool.py']: - module_name = f'tools.{filename[:-3]}' - module = importlib.import_module(module_name) - for name, obj in inspect.getmembers(module): - if inspect.isclass(obj) and issubclass(obj, BaseTool) and obj != BaseTool: - tools.append(obj()) - -# Collect all function definitions -functions = [] -for tool in tools: - functions.extend(tool.get_functions()) - -async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - logging.info("Bot started") - 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(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - user_id = update.effective_user.id - if user_id in conversation_history: - del conversation_history[user_id] - for tool in tools: - tool.clear() - - 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 - user_message = update.message.text - - logging.info(f"Message from user {user_id}: {user_message}") - - if user_id not in conversation_history: - conversation_history[user_id] = [] - - 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) - tool_calls = [] - fullMessage = [] - for message_part in response.content: - fullMessage.append(message_part) - if message_part.type == "tool_use": - tool_calls.append(message_part) - messages.append({"role": "assistant", "content": fullMessage}) - - toolUseCount = 0 - - previous_function_name = "" - - 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 - - if previous_function_name != function_name: - await update_status_message(context, update.effective_chat.id, status_message.message_id, f"Using tool: {function_name}") - previous_function_name = 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 = {} - - formatted_result = {"role": "user", "content":tool_use_results} - - messages.append(formatted_result) - - response = get_chat_response(messages) - fullMessage = [] - for message_part in response.content: - fullMessage.append(message_part) - if message_part.type == "tool_use": - tool_calls.append(message_part) - messages.append({"role": "assistant", "content": fullMessage}) - - toolUseCount += 1 - - if (toolUseCount == 0): - assistant_reply = response.content - conversation_history[user_id].append({"role": "assistant", "content": assistant_reply}) - - 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] - try: - await update.message.reply_text(messages[-1]["content"][0].text) - except TelegramErrors.BadRequest as e: - logging.error(f"An error occurred when trying to send a message in telegram: {str(e)}") - - 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.") - -def call_tool(function_call): - function_name = function_call.name - function_args = json.dumps(function_call.input) - for tool in tools: - if function_name in [f["name"] for f in tool.get_functions()]: - return tool.execute(function_name, **json.loads(function_args)) - -def get_chat_response(messages): - return get_claude_response(messages) - -def get_claude_response(messages): - anthropic_tools = [ + def get_chat_response(self, messages): + anthropic_tools = [ { "name": function['name'], "description": function['description'], "input_schema": function['parameters'] if function['parameters'] not in [None, {}] else {"type": "object", "properties": {"param1": {"type": "string", "description": "Unnecessary"}}, "required": []} } - for function in functions + for function in self.functions ] - try: - response = anthropic_client.messages.create( - model="claude-3-5-sonnet-20240620", - system=system_prompt, - messages=messages, - max_tokens=8192, - tools=anthropic_tools - ) - except Exception as e: - logging.error(f"An error occurred: {str(e)}") - return None - - return response + try: + response = self.anthropic_client.messages.create( + model="claude-3-5-sonnet-20240620", + system=self.system_prompt, + messages=messages, + max_tokens=8192, + tools=anthropic_tools + ) + except Exception as e: + logging.error(f"An error occurred: {str(e)}") + return None + + return response -async def status(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - await update.message.reply_text("Currently using claude-3-5-sonnet-20240620") + async def handle_message(self, user_id, user_message): + if user_id not in self.conversation_history: + self.conversation_history[user_id] = [] -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.") + self.conversation_history[user_id].append({"role": "user", "content": user_message}) + messages = self.conversation_history[user_id] -def main() -> None: - # Create the Application and pass it your bot's token - application = Application.builder().token(TELEGRAM_BOT_TOKEN).build() + response = self.get_chat_response(messages) + tool_calls = [] + full_message = [] + for message_part in response.content: + full_message.append(message_part) + if message_part.type == "tool_use": + tool_calls.append(message_part) + messages.append({"role": "assistant", "content": full_message}) - # Add handlers - application.add_handler(CommandHandler("start", start)) - application.add_handler(CommandHandler("clear", clear)) - application.add_handler(CommandHandler("status", status)) - application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message)) - application.add_handler(CallbackQueryHandler(abort_processing, pattern='^abort$')) + tool_use_count = 0 + while len(tool_calls) > 0 and tool_use_count < 50: + tool_use_results = [] + for tool_call in tool_calls: + tool_response = self.call_tool(tool_call) + tool_use_results.append({"type": "tool_result", "tool_use_id": tool_call.id, "content": json.dumps(tool_response)}) - # Start the Bot - logging.info("Bot is running...") - application.run_polling() + messages.append({"role": "user", "content": tool_use_results}) + + response = self.get_chat_response(messages) + full_message = [] + tool_calls = [] + for message_part in response.content: + full_message.append(message_part) + if message_part.type == "tool_use": + tool_calls.append(message_part) + messages.append({"role": "assistant", "content": full_message}) + + tool_use_count += 1 + + if len(self.conversation_history[user_id]) > 20: + self.conversation_history[user_id] = self.conversation_history[user_id][-20:] + + return messages[-1]["content"][0].text + + async def start(self): + logging.info("Bot started") + + async def clear(self, user_id): + super().clear_conversation(user_id) + logging.info(f"Cleared conversation history and image for user {user_id}") + + async def status(self): + return "Currently using claude-3-5-sonnet-20240620" + + async def abort_processing(self, user_id): + if user_id in self.processing_status: + self.processing_status[user_id]["processing"] = False + await self.clear(user_id) + return "Processing aborted." + else: + return "No active processing to abort." + +def main(): + bot = AnthropicTelegramInferenceBot() + telegram_helper = TelegramHelper(bot) + telegram_helper.run() if __name__ == '__main__': main() \ No newline at end of file diff --git a/base_inference_bot.py b/base_inference_bot.py new file mode 100644 index 0000000..68a4609 --- /dev/null +++ b/base_inference_bot.py @@ -0,0 +1,80 @@ +import json +import os +import importlib +import inspect +import logging +from abc import ABC, abstractmethod +from dotenv import load_dotenv +from tools.base_tool import BaseTool +from tools.metrics_tool import MetricsTool + +class BaseInferenceBot(ABC): + def __init__(self): + load_dotenv() + self.setup_logging() + self.load_system_prompt() + self.load_tools() + self.conversation_history = {} + self.processing_status = {} + + def setup_logging(self): + logging.basicConfig(level=logging.WARNING, handlers=[ + logging.StreamHandler(), + logging.FileHandler('logs/output.log', mode='a') + ]) + + def load_system_prompt(self): + with open("prompts/developer_prompt.txt", "r") as file: + self.system_prompt = file.read().strip() + + def load_tools(self): + self.tools = [MetricsTool()] + tools_dir = os.path.join(os.path.dirname(__file__), 'tools') + for filename in os.listdir(tools_dir): + if filename.endswith('.py') and filename not in ['__init__.py', 'base_tool.py', 'metrics_tool.py']: + module_name = f'tools.{filename[:-3]}' + module = importlib.import_module(module_name) + for name, obj in inspect.getmembers(module): + if inspect.isclass(obj) and issubclass(obj, BaseTool) and obj != BaseTool: + self.tools.append(obj()) + + self.functions = [] + for tool in self.tools: + self.functions.extend(tool.get_functions()) + + def clear_conversation(self, user_id): + if user_id in self.conversation_history: + del self.conversation_history[user_id] + for tool in self.tools: + tool.clear() + + def call_tool(self, function_call): + function_name = function_call.name + function_args = json.loads(function_call.arguments) + for tool in self.tools: + if function_name in [f["name"] for f in tool.get_functions()]: + return tool.execute(function_name, **function_args) + + @abstractmethod + def get_chat_response(self, messages): + pass + + @abstractmethod + async def handle_message(self, user_id, user_message): + pass + + @abstractmethod + async def start(self): + pass + + @abstractmethod + async def clear(self, user_id): + pass + + @abstractmethod + async def status(self): + pass + + @abstractmethod + async def abort_processing(self, user_id): + pass \ No newline at end of file diff --git a/base_telegram_inference_bot.py b/base_telegram_inference_bot.py new file mode 100644 index 0000000..2d26516 --- /dev/null +++ b/base_telegram_inference_bot.py @@ -0,0 +1,56 @@ +import os +import json +import logging +from abc import ABC, abstractmethod + +class BaseTelegramInferenceBot(ABC): + def __init__(self): + self.conversation_history = {} + self.processing_status = {} + self.system_prompt = self.load_system_prompt() + self.functions = self.load_functions() + + @staticmethod + def load_system_prompt(): + with open("prompts/developer_prompt.txt", "r") as file: + return file.read().strip() + + @staticmethod + def load_functions(): + # Implement function loading logic here + return [] + + @abstractmethod + def get_chat_response(self, messages): + pass + + @abstractmethod + async def handle_message(self, user_id, user_message): + pass + + def clear_conversation(self, user_id): + if user_id in self.conversation_history: + del self.conversation_history[user_id] + + def call_tool(self, function_call): + function_name = function_call.name + function_args = json.loads(function_call.arguments) + for tool in self.tools: + if function_name in [f["name"] for f in tool.get_functions()]: + return tool.execute(function_name, **function_args) + + @abstractmethod + async def start(self): + pass + + @abstractmethod + async def clear(self, user_id): + pass + + @abstractmethod + async def status(self): + pass + + @abstractmethod + async def abort_processing(self, user_id): + pass \ No newline at end of file diff --git a/chatgpt_telegram_inference_bot.py b/chatgpt_telegram_inference_bot.py index f7c6052..4bf6f91 100644 --- a/chatgpt_telegram_inference_bot.py +++ b/chatgpt_telegram_inference_bot.py @@ -1,225 +1,98 @@ import json import os -import importlib -import inspect import logging -import asyncio -from telegram import error as TelegramErrors, Update, __version__ as telegram_version, InlineKeyboardButton, InlineKeyboardMarkup -from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackQueryHandler -from dotenv import load_dotenv -from tools.base_tool import BaseTool -from tools.metrics_tool import MetricsTool +from base_telegram_inference_bot import BaseTelegramInferenceBot +from telegram_helper import TelegramHelper from openai import OpenAI -# Load environment variables -load_dotenv() +class ChatGPTTelegramInferenceBot(BaseTelegramInferenceBot): + def __init__(self): + super().__init__() + self.client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY")) + self.model = "gpt-4o-mini" + self.max_tokens = 16384 -client = OpenAI( - api_key=os.environ.get("OPENAI_API_KEY"), -) + def get_chat_response(self, messages): + response = self.client.chat.completions.create( + model=self.model, + messages=[{"role": "system", "content": self.system_prompt}] + messages, + functions=self.functions, + function_call="auto", + max_tokens=self.max_tokens + ) + return response -GPT_4O = "gpt-4o" -GPT_4O_MINI = "gpt-4o-mini" + async def handle_message(self, user_id, user_message): + if user_id not in self.conversation_history: + self.conversation_history[user_id] = [] -model_max_tokens = { - GPT_4O: 4096, - GPT_4O_MINI: 16384 -} + self.conversation_history[user_id].append({"role": "user", "content": user_message}) + messages = self.conversation_history[user_id] -use_smart_model = False - -# Set up logging to console and file -logging.basicConfig(level=logging.WARNING, handlers=[ - logging.StreamHandler(), - logging.FileHandler('logs/output.log', mode='a') -]) - -# Set up Telegram bot -TELEGRAM_BOT_TOKEN = os.getenv('TELEGRAM_BOT_TOKEN') - -# Load system prompt -with open("prompts/developer_prompt.txt", "r") as file: - system_prompt = file.read().strip() - -# Dictionary to store conversation history for each user -conversation_history = {} - -# Dictionary to store processing status for each user -processing_status = {} - -# Load tools -tools = [MetricsTool()] # Add MetricsTool instance -tools_dir = os.path.join(os.path.dirname(__file__), 'tools') -for filename in os.listdir(tools_dir): - if filename.endswith('.py') and filename not in ['__init__.py', 'base_tool.py', 'metrics_tool.py']: - module_name = f'tools.{filename[:-3]}' - module = importlib.import_module(module_name) - for name, obj in inspect.getmembers(module): - if inspect.isclass(obj) and issubclass(obj, BaseTool) and obj != BaseTool: - tools.append(obj()) - -# Collect all function definitions -functions = [] -for tool in tools: - functions.extend(tool.get_functions()) - -async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - logging.info("Bot started") - 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(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - user_id = update.effective_user.id - if user_id in conversation_history: - del conversation_history[user_id] - for tool in tools: - tool.clear() - - 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 - user_message = update.message.text - - logging.info(f"Message from user {user_id}: {user_message}") - - if user_id not in conversation_history: - conversation_history[user_id] = [] - - 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) + response = self.get_chat_response(messages) assistant_message = response.choices[0].message tool_calls = [] if hasattr(assistant_message, 'function_call') and assistant_message.function_call is not None: tool_calls.append(assistant_message.function_call) - toolUseCount = 0 - previous_function_name = "" - - while len(tool_calls) > 0 and toolUseCount < 50 and processing_status[user_id]["processing"]: + tool_use_count = 0 + while len(tool_calls) > 0 and tool_use_count < 50: tool_use_results = [] - while len(tool_calls) > 0: - tool_call = tool_calls.pop(0) - function_name = tool_call.name - - if function_name != previous_function_name: - # Update status message - await update_status_message(context, update.effective_chat.id, status_message.message_id, f"Using tool: {function_name}") - previous_function_name = function_name - - tool_response = call_tool(tool_call) - tool_use_results.append({"role": "function", "name": function_name, "content": json.dumps(tool_response)}) + for tool_call in tool_calls: + tool_response = self.call_tool(tool_call) + tool_use_results.append({"role": "function", "name": tool_call.name, "content": json.dumps(tool_response)}) messages.extend(tool_use_results) - response = get_chat_response(messages) + response = self.get_chat_response(messages) assistant_message = response.choices[0].message messages.append({"role": "assistant", "content": assistant_message.content}) + tool_calls = [] if hasattr(assistant_message, 'function_call') and assistant_message.function_call is not None: tool_calls.append(assistant_message.function_call) - toolUseCount += 1 + tool_use_count += 1 - if toolUseCount == 0: + if tool_use_count == 0: messages.append({"role": "assistant", "content": assistant_message.content}) - if len(conversation_history[user_id]) > 20: - conversation_history[user_id] = conversation_history[user_id][-20:] + if len(self.conversation_history[user_id]) > 20: + self.conversation_history[user_id] = self.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] - try: - await update.message.reply_text(messages[-1]["content"]) - except TelegramErrors.BadRequest as e: - logging.error(f"An error occurred when trying to send a message in telegram: {str(e)}") + return messages[-1]["content"] - 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 start(self): + logging.info("Bot started") -def call_tool(function_call): - function_name = function_call.name - function_args = json.loads(function_call.arguments) - for tool in tools: - if function_name in [f["name"] for f in tool.get_functions()]: - return tool.execute(function_name, **function_args) + async def clear(self, user_id): + super().clear_conversation(user_id) + logging.info(f"Cleared conversation history for user {user_id}") -def get_chat_response(messages): - model = GPT_4O if use_smart_model else GPT_4O_MINI - response = client.chat.completions.create( - model=model, - messages = [{"role": "system", "content": system_prompt}] + messages, - functions=functions, - function_call="auto", - max_tokens=model_max_tokens[model] - ) - return response + async def status(self): + return f"Currently using: {self.model}" -async def switch(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - global use_smart_model - use_smart_model = not use_smart_model - model = GPT_4O if use_smart_model else GPT_4O_MINI - logging.info(f"Switched to model: {model}") - await update.message.reply_text(f"Switched to model: {model}") + async def abort_processing(self, user_id): + if user_id in self.processing_status: + self.processing_status[user_id]["processing"] = False + await self.clear(user_id) + return "Processing aborted." + else: + return "No active processing to abort." -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 switch_model(self): + if self.model == "gpt-4o-mini": + self.model = "gpt-4o" + self.max_tokens = 4096 + else: + self.model = "gpt-4o-mini" + self.max_tokens = 16384 + logging.info(f"Switched to model: {self.model}") + return f"Switched to model: {self.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() - - # Add handlers - application.add_handler(CommandHandler("start", start)) - application.add_handler(CommandHandler("clear", clear)) - application.add_handler(CommandHandler("switch", switch)) - 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...") - application.run_polling() +def main(): + bot = ChatGPTTelegramInferenceBot() + telegram_helper = TelegramHelper(bot) + telegram_helper.run() if __name__ == '__main__': main() \ No newline at end of file diff --git a/telegram_helper.py b/telegram_helper.py new file mode 100644 index 0000000..7c095a6 --- /dev/null +++ b/telegram_helper.py @@ -0,0 +1,84 @@ +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') + + 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?" + ) + + 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 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 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 = [ + [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(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.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