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 anthropic import Anthropic from database import db # Import the database handler # Load environment variables load_dotenv() 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 processing status for each user processing_status = {} # Load tools tools = [] tools_dir = os.path.join(os.path.dirname(__file__), 'tools') for filename in os.listdir(tools_dir): if filename.endswith('.py') and filename != '__init__.py' and filename != 'base_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 db.clear_conversation(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}") conversation_history = db.get_conversation(user_id) conversation_history.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 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 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 = {} 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 messages.append({"role": "assistant", "content": assistant_reply}) if len(messages) > 20: messages = messages[-20:] # Save the updated conversation history db.save_conversation(user_id, messages) # 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 ] 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 async def status(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: await update.message.reply_text("Currently using claude-3-5-sonnet-20240620") 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("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()