294 lines
12 KiB
Python
294 lines
12 KiB
Python
import json
|
|
import os
|
|
import importlib
|
|
import inspect
|
|
import logging
|
|
import sys
|
|
import subprocess
|
|
import requests
|
|
from telegram import Update, __version__ as telegram_version, InlineKeyboardButton, InlineKeyboardMarkup, ReplyKeyboardMarkup
|
|
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
|
|
|
|
# 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 bots
|
|
DAEMON_BOT_TOKEN = os.getenv('TELEGRAM_BOT_TOKEN')
|
|
APPRENTICE_BOT_TOKEN = os.getenv('TELEGRAM_APPRENTICE_BOT_TOKEN')
|
|
GITHUB_REPO_OWNER = os.getenv('GITHUB_REPO_OWNER')
|
|
GITHUB_REPO_NAME = os.getenv('GITHUB_REPO_NAME')
|
|
GITHUB_ACCESS_TOKEN = os.getenv('GITHUB_ACCESS_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 = []
|
|
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())
|
|
|
|
def get_keyboard():
|
|
keyboard = [
|
|
['/switch', '/toggle'],
|
|
['/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.",
|
|
reply_markup=get_keyboard()
|
|
)
|
|
|
|
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!", reply_markup=get_keyboard())
|
|
|
|
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
|
|
|
|
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 = {"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]
|
|
|
|
await update.message.reply_text(messages[-1]["content"][0].text, reply_markup=get_keyboard())
|
|
|
|
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.", reply_markup=get_keyboard())
|
|
|
|
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", reply_markup=get_keyboard())
|
|
|
|
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.")
|
|
|
|
async def handover(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
|
if context.bot.token == DAEMON_BOT_TOKEN:
|
|
# Daemon bot initiating handover
|
|
apprentice_chat_id = os.getenv('APPRENTICE_CHAT_ID')
|
|
await context.bot.send_message(chat_id=apprentice_chat_id, text="Handover initiated. Taking control.")
|
|
await update.message.reply_text("Handover initiated. Apprentice bot is now in control.", reply_markup=get_keyboard())
|
|
else:
|
|
await update.message.reply_text("Handover can only be initiated by the daemon bot.", reply_markup=get_keyboard())
|
|
|
|
async def update_apprentice(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
|
if context.bot.token == APPRENTICE_BOT_TOKEN:
|
|
try:
|
|
# Pull latest changes
|
|
subprocess.run(["git", "pull", "origin", "main"], check=True)
|
|
|
|
# Restart the bot
|
|
os.execv(sys.executable, ['python'] + sys.argv)
|
|
except subprocess.CalledProcessError as e:
|
|
logging.error(f"Failed to pull latest changes: {e}")
|
|
await update.message.reply_text("Failed to update. Please check the logs.", reply_markup=get_keyboard())
|
|
except Exception as e:
|
|
logging.error(f"Failed to restart the bot: {e}")
|
|
await update.message.reply_text("Failed to restart. Please check the logs.", reply_markup=get_keyboard())
|
|
else:
|
|
await update.message.reply_text("Update can only be performed by the apprentice bot.", reply_markup=get_keyboard())
|
|
|
|
async def check_for_updates(context: ContextTypes.DEFAULT_TYPE) -> None:
|
|
url = f"https://api.github.com/repos/{GITHUB_REPO_OWNER}/{GITHUB_REPO_NAME}/pulls"
|
|
headers = {"Authorization": f"token {GITHUB_ACCESS_TOKEN}"}
|
|
response = requests.get(url, headers=headers)
|
|
|
|
if response.status_code == 200:
|
|
pull_requests = response.json()
|
|
for pr in pull_requests:
|
|
if pr['state'] == 'closed' and pr['merged']:
|
|
# A pull request was merged, update the apprentice bot
|
|
await context.bot.send_message(chat_id=os.getenv('APPRENTICE_CHAT_ID'), text="A new update is available. Updating now...")
|
|
await update_apprentice(None, context)
|
|
break
|
|
|
|
def main() -> None:
|
|
# Create the Application and pass it your bot's token
|
|
daemon_app = Application.builder().token(DAEMON_BOT_TOKEN).build()
|
|
apprentice_app = Application.builder().token(APPRENTICE_BOT_TOKEN).build()
|
|
|
|
# Add handlers for both bots
|
|
for app in [daemon_app, apprentice_app]:
|
|
app.add_handler(CommandHandler("start", start))
|
|
app.add_handler(CommandHandler("clear", clear))
|
|
app.add_handler(CommandHandler("status", status))
|
|
app.add_handler(CommandHandler("switch", start)) # Placeholder for /switch command
|
|
app.add_handler(CommandHandler("toggle", start)) # Placeholder for /toggle command
|
|
app.add_handler(CommandHandler("reset", clear)) # Use clear function for /reset command
|
|
app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))
|
|
app.add_handler(CallbackQueryHandler(abort_processing, pattern='^abort$'))
|
|
|
|
# Add handover command only to daemon bot
|
|
daemon_app.add_handler(CommandHandler("handover", handover))
|
|
|
|
# Add update command only to apprentice bot
|
|
apprentice_app.add_handler(CommandHandler("update", update_apprentice))
|
|
|
|
# Set up job queue to check for updates every 15 minutes
|
|
job_queue = apprentice_app.job_queue
|
|
job_queue.run_repeating(check_for_updates, interval=900, first=10)
|
|
|
|
# Start both bots
|
|
logging.info("Bots are running...")
|
|
daemon_app.run_polling()
|
|
apprentice_app.run_polling()
|
|
|
|
if __name__ == '__main__':
|
|
main() |