AI-Discord-Bot/bot.py

# bot.py

import os
import discord
import yaml
from discord.ext import commands
from dotenv import load_dotenv
from ai import get_ai_response
from personality import apply_personality, set_persona
from discord.ext.commands import (
    cooldown,
    BucketType,
    CooldownMapping,
    CommandOnCooldown
)
import yaml
from scheduler import start_scheduler

with open("settings.yml", "r", encoding="utf-8") as f:
    settings = yaml.safe_load(f)

ROAST_COOLDOWN_SECONDS = settings["cooldowns"]["roast"]
GLOBAL_COOLDOWN_SECONDS = settings["cooldowns"]["global"]
COOLDOWN_MSG_TEMPLATE = settings["messages"]["cooldown"]

load_dotenv()
TOKEN = os.getenv("DISCORD_TOKEN")

intents = discord.Intents.default()
intents.message_content = True

bot = commands.Bot(command_prefix="!", intents=intents)

@bot.event
async def on_command_error(ctx, error):
    if isinstance(error, CommandOnCooldown):
        retry_secs = round(error.retry_after, 1)
        msg = COOLDOWN_MSG_TEMPLATE.replace("{seconds}", str(retry_secs))
        print("🕒 Chill, mortal. You must wait 11.6s before trying again. 😼")
        await ctx.send(msg)
    else:
        raise error

# Global cooldown bucket
global_cooldown = CooldownMapping.from_cooldown(1, GLOBAL_COOLDOWN_SECONDS, BucketType.user)

@bot.check
async def global_command_cooldown(ctx):
    bucket = global_cooldown.get_bucket(ctx.message)
    retry_after = bucket.update_rate_limit()

    if retry_after:
        raise CommandOnCooldown(bucket, retry_after, BucketType.user)
    return True

@bot.command()
async def ping(ctx):
    await ctx.send("🏓 Pong!")

@bot.command()
async def chat(ctx, *, message):
    await ctx.send("🤖 Thinking...")
    reply = get_ai_response(message)
    await ctx.send(reply)

@bot.command()
async def setpersona(ctx, *, description):
    set_persona(description)
    await ctx.send("✅ Persona updated! New style will be used in replies.")


@bot.command(name='roast')
@cooldown(rate=1, per=ROAST_COOLDOWN_SECONDS, type=BucketType.user)
async def roast(ctx):
    # Get the mentioned user (or fallback to the author)
    target = ctx.message.mentions[0].mention if ctx.message.mentions else ctx.author.mention

    # Build the roast prompt
    prompt = f"Roast {target}. Be dramatic, insulting, and sarcastic. Speak in your usual chaotic RGB catgirl personality."

    # Get AI response
    response = get_ai_response(prompt)

    # Send the roast back to the channel
    await ctx.send(f"😼 {response}")

@bot.event
async def on_ready():
    print(f"✅ Logged in as {bot.user.name}")
    bot.loop.create_task(start_scheduler(bot))

bot.run(TOKEN)