diff --git a/telegram_inference_bot.py b/telegram_inference_bot.py index 0a83629..1aefac0 100644 --- a/telegram_inference_bot.py +++ b/telegram_inference_bot.py @@ -6,7 +6,6 @@ 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 openai import OpenAI from dotenv import load_dotenv from tools.base_tool import BaseTool from anthropic import Anthropic @@ -14,24 +13,11 @@ from anthropic import Anthropic # Load environment variables load_dotenv() -openai_client = OpenAI() - anthropic_client = Anthropic( api_key=os.environ.get("ANTHROPIC_API_KEY"), default_headers={"anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15"} ) -GPT_4O = "gpt-4o" -GPT_4O_MINI = "gpt-4o-mini" - -model_max_tokens = { - GPT_4O: 4096, - GPT_4O_MINI: 16384 -} - -use_smart_model = False -use_anthropic = True - # Set up logging to console and file logging.basicConfig(level=logging.WARNING, handlers=[ logging.StreamHandler(), @@ -67,9 +53,18 @@ functions = [] for tool in tools: functions.extend(tool.get_functions()) +def get_keyboard(): + keyboard = [ + ['/status', '/reset'] + ] + return ReplyKeyboardMarkup(keyboard, resize_keyboard=True) + 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.") + 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.", + reply_markup=get_keyboard() + ) async def clear(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: user_id = update.effective_user.id @@ -79,7 +74,7 @@ async def clear(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: 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!") + await update.message.reply_text("Conversation history and image cleared. Let's start fresh!", reply_markup=get_keyboard()) async def update_status_message(context: ContextTypes.DEFAULT_TYPE, chat_id: int, message_id: int, status: str): keyboard = [ @@ -113,17 +108,12 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> response = get_chat_response(messages) tool_calls = [] - if use_anthropic: - 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}) - else: - assistant_message = response.choices[0].message - if hasattr(assistant_message, 'function_call') and assistant_message.function_call is not None: - tool_calls.append(assistant_message.function_call) + 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 @@ -146,30 +136,17 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> messages.append(formatted_result) response = get_chat_response(messages) - assistant_message = "" - if use_anthropic: - 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}) - else: - assistant_message = response.choices[0].message - conversation_history[user_id].append({"role": "assistant", "content": assistant_message}) - if hasattr(assistant_message, 'function_call') and assistant_message.function_call is not None: - tool_calls.append(assistant_message.function_call) - else: - conversation_history[user_id].append({"role": "assistant", "content": assistant_message}) - assistant_reply = assistant_message + 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): - if use_anthropic: - assistant_reply = response.content - else: - assistant_reply = assistant_message + assistant_reply = response.content conversation_history[user_id].append({"role": "assistant", "content": assistant_reply}) if len(conversation_history[user_id]) > 20: @@ -185,28 +162,17 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> 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.") + await update.message.reply_text("Sorry, an error occurred while processing your request.", reply_markup=get_keyboard()) def call_tool(function_call): - function_name = function_call.name if use_anthropic else function_call.name - function_args = json.dumps(function_call.input) if use_anthropic else function_call.arguments + 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) if use_anthropic else get_openai_response(messages) - -def get_openai_response(messages): - model = GPT_4O if use_smart_model else GPT_4O_MINI - response = openai_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 + return get_claude_response(messages) def get_claude_response(messages): anthropic_tools = [ @@ -231,26 +197,8 @@ def get_claude_response(messages): return response -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 switch_providers(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - await clear(update, context) - global use_anthropic - use_anthropic = not use_anthropic - logging.info("Using Anthropic" if use_anthropic else "Using OpenAI") - await update.message.reply_text("Using Anthropic" if use_anthropic else "Using OpenAI") - async def status(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: - if use_anthropic: - await update.message.reply_text("Currently using claude-3-5-sonnet-20240620") - else: - model = GPT_4O if use_smart_model else GPT_4O_MINI - await update.message.reply_text(f"Currently using: {model}") + await update.message.reply_text("Currently using claude-3-5-sonnet-20240620", reply_markup=get_keyboard()) async def abort_processing(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: query = update.callback_query