diff --git a/agent-py-bot/agent.py b/agent-py-bot/agent.py index c2c8528..fb86aa4 100644 --- a/agent-py-bot/agent.py +++ b/agent-py-bot/agent.py @@ -1,6 +1,9 @@ import logging -from telegram.ext import Updater, CommandHandler, MessageHandler, filters, CallbackContext +import asyncio, nest_asyncio +from telegram import Bot, Update +from telegram.ext import Application, CommandHandler, MessageHandler, filters +# import "gopkg.in/telebot.v3/middleware" import requests import json @@ -10,6 +13,9 @@ from selenium.webdriver.chrome.options import Options from io import BytesIO from PIL import Image +# Apply nest_asyncio +nest_asyncio.apply() + # Set up logging logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO) logger = logging.getLogger(__name__) @@ -23,22 +29,18 @@ TOKEN = '6805059978:AAHNJKuOeazMSJHc3-BXRCsFfEVyFHeFnjw' # LLM API Endpoint LLM_ENDPOINT = "http://192.168.0.11:11434/api/chat" -# Selenium WebDriver setup for screenshots -chrome_options = Options() -chrome_options.add_argument("--headless") -driver = webdriver.Chrome(options=chrome_options) +#! Selenium WebDriver setup for screenshots +#chrome_options = Options() +#chrome_options.add_argument("--headless") +#driver = webdriver.Chrome(options=chrome_options) -def start(update, context): - """Send a message when the command /start is issued.""" - update.message.reply_text('Hi I am ! Send me a message, and I will interact with LLM.') +async def start(update: Update, context): + await context.bot.send_message(chat_id=update.effective_chat.id, text="Hi! I'm your bot.") -def echo(update, context): - """Echo the user message.""" - user_message = update.message.text - response = query_llm(user_message) - update.message.reply_text(response) +async def echo(update: Update, context): + await context.bot.send_message(chat_id=update.effective_chat.id, text=update.message.text) -def query_llm(user_message): +async def query_llm(user_message): """Query the LLM with the user's message.""" data = { "model": "llama2", @@ -51,34 +53,45 @@ def query_llm(user_message): else: return "Error: Unable to reach LLM" -def screenshot(update, context): +async def screenshot(update, context): """Take a screenshot of a webpage.""" url = ' '.join(context.args) - driver.get(url) - screenshot = driver.get_screenshot_as_png() - image_stream = BytesIO(screenshot) - image_stream.seek(0) - image = Image.open(image_stream) - image_stream.close() - image.save('screenshot.png') - update.message.reply_photo(photo=open('screenshot.png', 'rb')) + update.message.reply_text('This will noramlly get a screenshot from: '.url)# url.', but currently under development') -def error(update, context): + + # driver.get(url) + # screenshot = driver.get_screenshot_as_png() + # image_stream = BytesIO(screenshot) + # image_stream.seek(0) + # image = Image.open(image_stream) + # image_stream.close() + # image.save('screenshot.png') + # update.message.reply_photo(photo=open('screenshot.png', 'rb')) + +async def error(update, context): """Log Errors caused by Updates.""" logger.warning(f'Update "{update}" caused error "{context.error}"') -def main(): + +async def main(): """Start the bot.""" - updater = Updater(TOKEN, use_context=True) + # Create an Application instance + application = Application.builder().token(TOKEN).build() - dp = updater.dispatcher - dp.add_handler(CommandHandler("start", start)) - dp.add_handler(MessageHandler(Filters.text, echo)) - dp.add_handler(CommandHandler("screenshot", screenshot, pass_args=True)) - dp.add_error_handler(error) + # Add handlers to the application + application.add_handler(CommandHandler("start", start)) + application.add_handler(MessageHandler(filters.TEXT, echo)) + application.add_handler(CommandHandler("screenshot", screenshot)) # Ensure screenshot function is async + application.add_handler(CommandHandler("ai",query_llm)) + application.add_error_handler(error) # Ensure error function is async + + # Run the bot + await application.run_polling() - updater.start_polling() - updater.idle() if __name__ == '__main__': - main() + loop = asyncio.get_event_loop() + if loop.is_running(): + loop.create_task(main()) + else: + asyncio.run(main())