280 lines
9.0 KiB
Python
280 lines
9.0 KiB
Python
import os
|
|
import openai
|
|
import simplematrixbotlib as botlib
|
|
import validators
|
|
import wget
|
|
|
|
from bofa import BofaData
|
|
from org import OrgData
|
|
from prometeo import get_bank_information
|
|
from settings import (
|
|
BANK_ACCOUNT_NUMBERS,
|
|
BROU_PASSWORD,
|
|
BROU_USERNAME,
|
|
ITAU_PASSWORD,
|
|
ITAU_USERNAME,
|
|
MATRIX_PASSWORD,
|
|
MATRIX_URL,
|
|
MATRIX_USER,
|
|
MATRIX_USERNAME,
|
|
MATRIX_USERNAMES,
|
|
OPEN_AI_API_KEY
|
|
)
|
|
|
|
openai.api_key = OPEN_AI_API_KEY
|
|
|
|
creds = botlib.Creds(MATRIX_URL, MATRIX_USER, MATRIX_PASSWORD)
|
|
bot = botlib.Bot(creds)
|
|
|
|
PREFIX = "!"
|
|
|
|
MATRIX_USERNAMES = MATRIX_USERNAMES.split(",")
|
|
|
|
starting_prompt = {"role": "system", "content": "You are a very helpful chatbot"}
|
|
|
|
CONVERSATION = {}
|
|
for username in MATRIX_USERNAMES:
|
|
CONVERSATION[username] = [starting_prompt]
|
|
|
|
|
|
@bot.listener.on_message_event
|
|
async def todo(room, message):
|
|
"""
|
|
Function that adds new TODOs
|
|
Usage:
|
|
user: !todo title - objective - extra (optional)
|
|
bot: TODO added!
|
|
"""
|
|
match = botlib.MessageMatch(room, message, bot, PREFIX)
|
|
if match.is_not_from_this_bot() and match.prefix():
|
|
user = message.sender
|
|
keyword = str(message).split()[1].replace("!", "").lower()
|
|
|
|
if user == MATRIX_USERNAME and keyword in ["todo", "repeat", "next", "waiting", "someday", "proj"]:
|
|
message = " ".join(message.body.split(" ")[1:])
|
|
room_id = room.room_id
|
|
splitted_message = message.split("-")
|
|
|
|
try:
|
|
todo_title = splitted_message[0].strip()
|
|
todo_objective = splitted_message[1].strip()
|
|
except IndexError:
|
|
return await bot.api.send_text_message(
|
|
room_id,
|
|
"An objective is needed. The correct format is 'Title - Objective - Extra (optional)'",
|
|
)
|
|
|
|
try:
|
|
todo_extra = splitted_message[2].strip()
|
|
except IndexError:
|
|
todo_extra = ""
|
|
|
|
print(f"Room: {room_id}, User: {user}, Message: {message}")
|
|
OrgData().add_new_todo(keyword, todo_title, todo_objective, todo_extra)
|
|
await bot.api.send_text_message(room_id, "TODO added!")
|
|
|
|
|
|
@bot.listener.on_message_event
|
|
async def list_todos(room, message):
|
|
"""
|
|
Function that lists today's plan
|
|
Usage:
|
|
user: !list [free,work]
|
|
bot: [prints a list with today's todos]
|
|
"""
|
|
match = botlib.MessageMatch(room, message, bot, PREFIX)
|
|
if match.is_not_from_this_bot() and match.prefix() and match.command("list"):
|
|
user = message.sender
|
|
|
|
if user == MATRIX_USERNAME:
|
|
message = " ".join(message.body.split(" ")[1:]).strip() or "free"
|
|
room_id = room.room_id
|
|
|
|
if message not in ["free", "work"]:
|
|
return await bot.api.send_text_message(
|
|
room_id,
|
|
f'"{message}" not accepted. Accepted options are "free" and "work"',
|
|
)
|
|
|
|
print(f"Room: {room_id}, User: {user}, Message: {message}")
|
|
plan = OrgData().list_plan(message)
|
|
await bot.api.send_text_message(room_id, plan)
|
|
|
|
|
|
@bot.listener.on_message_event
|
|
async def list_bofa_status(room, message):
|
|
"""
|
|
Function that lists bofa status
|
|
Usage:
|
|
user: !bofa
|
|
bot: [prints bofa current status]
|
|
"""
|
|
match = botlib.MessageMatch(room, message, bot, PREFIX)
|
|
if match.is_not_from_this_bot() and match.prefix() and match.command("bofa"):
|
|
user = message.sender
|
|
|
|
if user == MATRIX_USERNAME:
|
|
room_id = room.room_id
|
|
|
|
print(f"Room: {room_id}, User: {user}, Message: bofa")
|
|
|
|
bofa_data = BofaData().get()
|
|
|
|
return_data = ""
|
|
for person, amount in bofa_data.items():
|
|
if amount != "0 USD":
|
|
return_data += f"{person}: {amount}\n"
|
|
|
|
await bot.api.send_text_message(room_id, return_data)
|
|
|
|
|
|
def generate_return_data_from_account(bank: str, accounts: list) -> str:
|
|
bank_message = f"{bank.upper()}\nNot available. Try again later!\n"
|
|
|
|
for account in accounts:
|
|
if account.get("number", "") in BANK_ACCOUNT_NUMBERS:
|
|
account_number = account.get("number", "")
|
|
balance = account.get("balance", "")
|
|
|
|
bank_message = f"{bank.upper()}\nAccount Number: {account_number}\nBalance: {balance} USD\n"
|
|
return bank_message
|
|
|
|
|
|
@bot.listener.on_message_event
|
|
async def list_bank_information(room, message):
|
|
"""
|
|
Function that lists banks information
|
|
Usage:
|
|
user: !banks
|
|
bot: [prints current status of banks]
|
|
"""
|
|
match = botlib.MessageMatch(room, message, bot, PREFIX)
|
|
if match.is_not_from_this_bot() and match.prefix() and match.command("banks"):
|
|
user = message.sender
|
|
|
|
if user == MATRIX_USERNAME:
|
|
room_id = room.room_id
|
|
|
|
print(f"Room: {room_id}, User: {user}, Message: banks")
|
|
await bot.api.send_text_message(room_id, "Looking for bank data, just a sec...")
|
|
|
|
brou_accounts = get_bank_information("brou", BROU_USERNAME, BROU_PASSWORD)
|
|
itau_accounts = get_bank_information("itau", ITAU_USERNAME, ITAU_PASSWORD)
|
|
|
|
return_data = ""
|
|
return_data += generate_return_data_from_account("brou", brou_accounts)
|
|
return_data += "\n"
|
|
return_data += generate_return_data_from_account("itau", itau_accounts)
|
|
|
|
await bot.api.send_text_message(room_id, return_data)
|
|
|
|
|
|
@bot.listener.on_message_event
|
|
async def save_link(room, message):
|
|
"""
|
|
Function that lists banks information
|
|
Usage:
|
|
user: !banks
|
|
bot: [prints current status of banks]
|
|
"""
|
|
match = botlib.MessageMatch(room, message, bot)
|
|
message_content = message.body
|
|
if match.is_not_from_this_bot() and validators.url(message_content):
|
|
user = message.sender
|
|
|
|
if user == MATRIX_USERNAME:
|
|
room_id = room.room_id
|
|
|
|
print(f"Room: {room_id}, User: {user}, Message: {message_content}")
|
|
|
|
OrgData().add_new_link(f"- {message_content}\n")
|
|
await bot.api.send_text_message(room_id, "Link added!")
|
|
|
|
|
|
@bot.listener.on_message_event
|
|
async def chatgpt(room, message):
|
|
"""
|
|
Function that starts a conversation with chatgpt
|
|
Usage:
|
|
user: !chatgpt Hello!
|
|
bot: [prints chatgpt response]
|
|
"""
|
|
match = botlib.MessageMatch(room, message, bot, PREFIX)
|
|
message_content = message.body
|
|
if match.is_not_from_this_bot() and not match.prefix() and not validators.url(message_content):
|
|
user = message.sender
|
|
|
|
if user in MATRIX_USERNAMES:
|
|
personal_conversation = CONVERSATION[user]
|
|
room_id = room.room_id
|
|
|
|
print(f"Room: {room_id}, User: {user}, Message: chatgpt")
|
|
|
|
def format_message(message):
|
|
return {"role": "user", "content": message}
|
|
|
|
personal_conversation.append(format_message(message_content))
|
|
|
|
try:
|
|
completion = openai.ChatCompletion.create(model="gpt-4", messages=personal_conversation)
|
|
response = completion.choices[0].message.content
|
|
personal_conversation.append(completion.choices[0].message)
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|
|
response = "There was a problem with your prompt"
|
|
|
|
await bot.api.send_text_message(room_id, response)
|
|
|
|
|
|
@bot.listener.on_message_event
|
|
async def reset_chatgpt(room, message):
|
|
"""
|
|
Function that resets a conversation with chatgpt
|
|
Usage:
|
|
user: !reset
|
|
bot: Conversation reset!
|
|
"""
|
|
match = botlib.MessageMatch(room, message, bot, PREFIX)
|
|
if match.is_not_from_this_bot() and match.prefix() and match.command("reset"):
|
|
user = message.sender
|
|
|
|
if user in MATRIX_USERNAMES:
|
|
CONVERSATION[user] = [starting_prompt]
|
|
room_id = room.room_id
|
|
|
|
await bot.api.send_text_message(room_id, "Conversation reset!")
|
|
|
|
|
|
@bot.listener.on_message_event
|
|
async def dall_e(room, message):
|
|
"""
|
|
Function that generates a Dall-E image
|
|
Usage:
|
|
user: !dalle A sunny caribbean beach
|
|
bot: returns an image
|
|
"""
|
|
match = botlib.MessageMatch(room, message, bot, PREFIX)
|
|
if match.is_not_from_this_bot() and match.prefix() and match.command("dalle"):
|
|
user = message.sender
|
|
|
|
if user in MATRIX_USERNAMES:
|
|
room_id = room.room_id
|
|
message = " ".join(message.body.split(" ")[1:]).strip()
|
|
|
|
print(f"Room: {room_id}, User: {user}, Message: dalle")
|
|
await bot.api.send_text_message(room_id, "Generating image...")
|
|
|
|
try:
|
|
image = openai.Image.create(prompt=message)
|
|
image_url = image["data"][0]["url"]
|
|
image_filename = wget.download(image_url)
|
|
|
|
await bot.api.send_image_message(room_id, image_filename)
|
|
os.remove(image_filename)
|
|
except Exception as e:
|
|
print(f"Error sending image: {e}")
|
|
await bot.api.send_text_message(room_id, f"Error sending image: {e}")
|
|
|
|
|
|
bot.run()
|