Merge pull request #71 from bucolucas/refactor-inference-bots

Refactor inference bots and extract Telegram logic

anthropic_telegram_inference_bot.py

@@ -1,189 +1,31 @@
-import json
 import os
-import importlib
+import json
-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 anthropic import Anthropic
+from base_telegram_inference_bot import BaseTelegramInferenceBot
+from telegram_helper import TelegramHelper
 
-# Load environment variables
+class AnthropicTelegramInferenceBot(BaseTelegramInferenceBot):
-load_dotenv()
+    def __init__(self):
-
+        super().__init__()
-anthropic_client = Anthropic(
+        self.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
+    def get_chat_response(self, messages):
-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 = [
             {
                 "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(
+            response = self.anthropic_client.messages.create(
                 model="claude-3-5-sonnet-20240620",
-            system=system_prompt,
+                system=self.system_prompt,
                 messages=messages,
                 max_tokens=8192,
                 tools=anthropic_tools
@@ -194,39 +36,69 @@ def get_claude_response(messages):
             
         return response
 
-async def status(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+    async def handle_message(self, user_id, user_message):
-    await update.message.reply_text("Currently using claude-3-5-sonnet-20240620")
+        if user_id not in self.conversation_history:
+            self.conversation_history[user_id] = []
 
-async def abort_processing(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+        self.conversation_history[user_id].append({"role": "user", "content": user_message})
-    query = update.callback_query
+        messages = self.conversation_history[user_id]
-    await query.answer()
 
-    user_id = query.from_user.id
+        response = self.get_chat_response(messages)
-    if user_id in processing_status:
+        tool_calls = []
-        processing_status[user_id]["processing"] = False
+        full_message = []
-        await context.bot.edit_message_text(
+        for message_part in response.content:
-            chat_id=query.message.chat_id,
+            full_message.append(message_part)
-            message_id=query.message.message_id,
+            if message_part.type == "tool_use":
-            text="Processing aborted."
+                tool_calls.append(message_part)
-        )
+        messages.append({"role": "assistant", "content": full_message})
-        await clear(update, context)
+
+        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)})
+
+            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:
-        await query.edit_message_text(text="No active processing to abort.")
+            return "No active processing to abort."
 
-def main() -> None:
+def main():
-    # Create the Application and pass it your bot's token
+    bot = AnthropicTelegramInferenceBot()
-    application = Application.builder().token(TELEGRAM_BOT_TOKEN).build()
+    telegram_helper = TelegramHelper(bot)
-
+    telegram_helper.run()
-    # 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$'))
-
-    # Start the Bot
-    logging.info("Bot is running...")
-    application.run_polling()
 
 if __name__ == '__main__':
     main()

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

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

chatgpt_telegram_inference_bot.py

@@ -1,225 +1,98 @@
 import json
 import os
-import importlib
-import inspect
 import logging
-import asyncio
+from base_telegram_inference_bot import BaseTelegramInferenceBot
-from telegram import error as TelegramErrors, Update, __version__ as telegram_version, InlineKeyboardButton, InlineKeyboardMarkup
+from telegram_helper import TelegramHelper
-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 openai import OpenAI
 
-# Load environment variables
+class ChatGPTTelegramInferenceBot(BaseTelegramInferenceBot):
-load_dotenv()
+    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(
+    def get_chat_response(self, messages):
-    api_key=os.environ.get("OPENAI_API_KEY"),
+        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"
+    async def handle_message(self, user_id, user_message):
-GPT_4O_MINI = "gpt-4o-mini"
+        if user_id not in self.conversation_history:
+            self.conversation_history[user_id] = []
 
-model_max_tokens = {
+        self.conversation_history[user_id].append({"role": "user", "content": user_message})
-    GPT_4O: 4096,
+        messages = self.conversation_history[user_id]
-    GPT_4O_MINI: 16384
-}
 
-use_smart_model = False
+        response = self.get_chat_response(messages)
-
-# 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)
         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
+        tool_use_count = 0
-        previous_function_name = ""
+        while len(tool_calls) > 0 and tool_use_count < 50:
-
-        while len(tool_calls) > 0 and toolUseCount < 50 and processing_status[user_id]["processing"]:
             tool_use_results = []
-            while len(tool_calls) > 0:
+            for tool_call in tool_calls:
-                tool_call = tool_calls.pop(0)
+                tool_response = self.call_tool(tool_call)
-                function_name = tool_call.name
+                tool_use_results.append({"role": "function", "name": tool_call.name, "content": json.dumps(tool_response)})
-                
-                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)})
 
             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:
+        if len(self.conversation_history[user_id]) > 20:
-            conversation_history[user_id] = conversation_history[user_id][-20:]
+            self.conversation_history[user_id] = self.conversation_history[user_id][-20:]
 
-        # Remove the status message
+        return messages[-1]["content"]
-        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)}")
 
-    except Exception as e:
+    async def start(self):
-        logging.error(f"An error occurred: {str(e)}")
+        logging.info("Bot started")
-        await update.message.reply_text("Sorry, an error occurred while processing your request.")
 
-def call_tool(function_call):
+    async def clear(self, user_id):
-    function_name = function_call.name
+        super().clear_conversation(user_id)
-    function_args = json.loads(function_call.arguments)
+        logging.info(f"Cleared conversation history for user {user_id}")
-    for tool in tools:
-        if function_name in [f["name"] for f in tool.get_functions()]:
-            return tool.execute(function_name, **function_args)
 
-def get_chat_response(messages):
+    async def status(self):
-    model = GPT_4O if use_smart_model else GPT_4O_MINI
+        return f"Currently using: {self.model}"
-    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 switch(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+    async def abort_processing(self, user_id):
-    global use_smart_model
+        if user_id in self.processing_status:
-    use_smart_model = not use_smart_model
+            self.processing_status[user_id]["processing"] = False
-    model = GPT_4O if use_smart_model else GPT_4O_MINI
+            await self.clear(user_id)
-    logging.info(f"Switched to model: {model}")
+            return "Processing aborted."
-    await update.message.reply_text(f"Switched to model: {model}")
-
-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.")
+            return "No active processing to abort."
 
-def main() -> None:
+    async def switch_model(self):
-    # Create the Application and pass it your bot's token
+        if self.model == "gpt-4o-mini":
-    application = Application.builder().token(TELEGRAM_BOT_TOKEN).build()
+            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}"
 
-    # Add handlers
+def main():
-    application.add_handler(CommandHandler("start", start))
+    bot = ChatGPTTelegramInferenceBot()
-    application.add_handler(CommandHandler("clear", clear))
+    telegram_helper = TelegramHelper(bot)
-    application.add_handler(CommandHandler("switch", switch))
+    telegram_helper.run()
-    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()
 
 if __name__ == '__main__':
     main()

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()