118 lines
3.9 KiB
Python
118 lines
3.9 KiB
Python
import logging
|
|
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
|
|
import base64
|
|
from selenium import webdriver
|
|
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__)
|
|
|
|
# Telegram Bot Token
|
|
# t.me/kevin_ai_robot
|
|
TOKEN = '6805059978:AAHNJKuOeazMSJHc3-BXRCsFfEVyFHeFnjw'
|
|
# t.me/artitherobot 6749075936:AAHUHiPTDEIu6JH7S2fQdibwsu6JVG3FNG0
|
|
|
|
|
|
# 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)
|
|
|
|
async def start(update: Update, context):
|
|
await context.bot.send_message(chat_id=update.effective_chat.id, text="Hi! I'm your AI bot. Ask me aything with /ask")
|
|
|
|
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):
|
|
"""Query the LLM with the user's message."""
|
|
data = {
|
|
"model": "llama2",
|
|
"messages": [{"role": "user", "content": user_message}]
|
|
}
|
|
response = requests.post(LLM_ENDPOINT, json=data)
|
|
if response.status_code == 200:
|
|
response_data = response.json()
|
|
return response_data.get('message', {}).get('content', 'No response')
|
|
else:
|
|
return "Error: Unable to reach LLM"
|
|
|
|
def ask(update, context):
|
|
user_message = ' '.join(context.args)
|
|
llm_response = query_llm(user_message)
|
|
update.message.reply_text(llm_response)
|
|
|
|
async def screenshot(update, context):
|
|
"""Take a screenshot of a webpage."""
|
|
url = ' '.join(context.args)
|
|
update.message.reply_text('This will noramlly get a screenshot from: '.url)# url.', but currently under development')
|
|
|
|
|
|
# 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_handler(update: Update): #context: CallbackContext
|
|
"""Handle errors occurred during the execution of commands."""
|
|
# Log the error before we do anything else
|
|
# logger.error(f"An error occurred: {context.error}")
|
|
logger.error(f"An error occurred:")
|
|
|
|
# Send a message to the user
|
|
# error_message = "Sorry, an unexpected error occurred. Please try again later."
|
|
# if update.effective_message:
|
|
# await context.bot.send_message(chat_id=update.effective_chat.id, text=error_message)
|
|
|
|
|
|
async def main():
|
|
"""Start the bot."""
|
|
# Create an Application instance
|
|
application = Application.builder().token(TOKEN).build()
|
|
|
|
# Add handlers to the application
|
|
# Command handlers should be registered before the generic message handler
|
|
application.add_handler(CommandHandler("start", start))
|
|
application.add_handler(CommandHandler("screenshot", screenshot)) # Ensure screenshot function is async
|
|
application.add_handler(CommandHandler("ai", ask))
|
|
application.add_handler(CommandHandler("ask", ask))
|
|
application.add_handler(CommandHandler("llm", ask))
|
|
|
|
# This handler should be last as it's the most generic
|
|
application.add_handler(MessageHandler(filters.TEXT, echo))
|
|
|
|
# Register the error handler
|
|
application.add_error_handler(error_handler)
|
|
|
|
# Run the bot
|
|
await application.run_polling()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
loop = asyncio.get_event_loop()
|
|
if loop.is_running():
|
|
loop.create_task(main())
|
|
else:
|
|
asyncio.run(main())
|