working bot

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