85 lines
2.6 KiB
Python
85 lines
2.6 KiB
Python
import logging
|
|
from telegram.ext import Updater, CommandHandler, MessageHandler, filters, CallbackContext
|
|
|
|
|
|
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
|
|
|
|
# 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)
|
|
|
|
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.')
|
|
|
|
def echo(update, context):
|
|
"""Echo the user message."""
|
|
user_message = update.message.text
|
|
response = query_llm(user_message)
|
|
update.message.reply_text(response)
|
|
|
|
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 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'))
|
|
|
|
def error(update, context):
|
|
"""Log Errors caused by Updates."""
|
|
logger.warning(f'Update "{update}" caused error "{context.error}"')
|
|
|
|
def main():
|
|
"""Start the bot."""
|
|
updater = Updater(TOKEN, use_context=True)
|
|
|
|
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)
|
|
|
|
updater.start_polling()
|
|
updater.idle()
|
|
|
|
if __name__ == '__main__':
|
|
main()
|