AI-Discord-Bot/src/bot.py

# bot.py

import os
import discord
import yaml
import random
from dotenv import load_dotenv
from textwrap import wrap
from discord.ext import commands
from discord.ext.commands import (
    cooldown,
    BucketType,
    CooldownMapping,
    CommandOnCooldown
)

# Local imports
from scheduler import start_scheduler
from profilepic import set_avatar_from_bytes
from context import fetch_recent_context, format_context
from user_profiles import (
    load_user_profile,
    update_last_seen,
    increment_interactions,
    format_profile_for_block,
    set_pronouns,
    set_custom_prompt
)
from personality import apply_personality, set_persona
from logger import setup_logger
from ai import (
    unload_model,
    load_model,
    get_current_model,
    get_ai_response,
    TAGS_ENDPOINT
)

from user_profiles import format_profile_for_block as format_user_profile_block

# Setup logger and environment
logger = setup_logger("bot")
dotenv_path = os.path.join(os.path.dirname(__file__), '..', '.env')
load_dotenv(dotenv_path)

# Load model settings
MODEL_NAME = os.getenv("MODEL_NAME", "llama3:latest")
logger.info(f"🔍 Loaded MODEL_NAME from .env: {MODEL_NAME}")
logger.info(f"🧹 Attempting to clear VRAM before loading {MODEL_NAME}...")
unload_model(MODEL_NAME)

if load_model(MODEL_NAME):
    logger.info(f"🚀 Model `{MODEL_NAME}` preloaded on startup.")
else:
    logger.warning(f"⚠️ Failed to preload model `{MODEL_NAME}`.")
logger.info(f"✅ Final model in use: {MODEL_NAME}")

# Load YAML settings
base_dir = os.path.dirname(__file__)
settings_path = os.path.join(base_dir, "settings.yml")
with open(settings_path, "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"]

# Configure Discord bot
TOKEN = os.getenv("DISCORD_TOKEN")
if not TOKEN:
    logger.error("❌ DISCORD_TOKEN not set in .env file.")
    raise SystemExit("DISCORD_TOKEN not set.")

intents = discord.Intents.default()
intents.message_content = True
bot = commands.Bot(command_prefix="!", intents=intents)

# Handle cooldown errors globally
@bot.event
async def on_command_error(ctx, error):
    if isinstance(error, CommandOnCooldown):
        retry_secs = round(error.retry_after, 1)
        template = random.choice(COOLDOWN_MSG_TEMPLATE) if isinstance(COOLDOWN_MSG_TEMPLATE, list) else COOLDOWN_MSG_TEMPLATE
        msg = template.replace("{seconds}", str(retry_secs))
        logger.info(f"Command {ctx.command} on cooldown. Retry after {retry_secs} seconds.")
        await ctx.send(msg)
    else:
        raise error

# Global cooldown
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

# Handle direct bot mentions
@bot.event
async def on_message(message):
    # Ignore messages from the bot itself
    if message.author == bot.user:
        return

    # Only respond if the bot is mentioned
    if bot.user.mentioned_in(message):
        # Strip the mention from the message to extract the actual prompt
        prompt = message.content.replace(f"<@{bot.user.id}>", "").strip()
        if not prompt:
            return  # Nothing to respond to

        # Update the user's interaction history
        user_id = str(message.author.id)
        update_last_seen(user_id)
        #increment_interactions(user_id) <---------------------------------------------------Test
        profile = load_user_profile(message.author)

        # Log summary info about the profile (but don’t inject it)
        logger.info(f"🧠 Profile loaded for {profile['display_name']} (interactions: {profile['interactions']})")

        # Fetch recent messages for conversation context
        context_msgs = await fetch_recent_context(message.channel)
        formatted_context = format_context(context_msgs)

        # Log number of context messages, not the entire block (to reduce clutter)
        logger.info(f"📚 Retrieved {len(context_msgs)} messages for context")

        # Let ai.py handle all prompt construction (persona + context + profile)
        async with message.channel.typing():
            reply = get_ai_response(prompt, context=formatted_context, user_profile=profile)
            await message.channel.send(reply)

    # Ensure bot commands still work (!setprompt, !roast, etc.)
    await bot.process_commands(message)


# Bot startup event
@bot.event
async def on_ready():
    print(f"✅ Logged in as {bot.user.name}")
    logger.info(f"Logged in as {bot.user.name}")
    for guild in bot.guilds:
        me = guild.me
        if me.nick != "Delta":
            try:
                await me.edit(nick="Delta")
                logger.info(f"🔄 Renamed self to Delta in {guild.name}")
            except Exception as e:
                logger.warning(f"⚠️ Failed to rename in {guild.name}: {e}")
    bot.loop.create_task(start_scheduler(bot))

# Commands
@bot.command(name="setprompt")
async def set_prompt_cmd(ctx, *, prompt):
    set_custom_prompt(ctx.author.id, prompt)
    await ctx.send("✅ Custom prompt saved.")

@bot.command(name="setpronouns")
async def set_pronouns_cmd(ctx, *, pronouns):
    success = set_pronouns(ctx.author, pronouns)
    if success:
        await ctx.send(f"✅ Got it, {ctx.author.display_name}! Your pronouns have been updated.")
    else:
        await ctx.send("⚠️ Failed to update pronouns. Try interacting with Delta first to generate your profile.")

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

@bot.command()
async def chat(ctx, *, prompt):
    await ctx.send("🤖 Thinking...")
    reply = get_ai_response(prompt)
    for chunk in wrap(reply, 2000):
        await ctx.send(chunk)

@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):
    target = ctx.message.mentions[0].mention if ctx.message.mentions else ctx.author.mention
    prompt = f"Roast {target}. Be dramatic, insulting, and sarcastic. Speak in your usual chaotic RGB catgirl personality."
    response = get_ai_response(prompt)
    await ctx.send(f"😼 {response}")

@bot.command(name="clearmodel")
async def clear_model(ctx):
    model = get_current_model()
    success = unload_model(model)
    msg = f"✅ Unloaded model: `{model}`" if success else f"❌ Failed to unload model: `{model}`"
    await ctx.send(msg)

@bot.command(name="model")
async def current_model(ctx):
    model = get_current_model()
    await ctx.send(f"📦 Current model: `{model}`")

@bot.command(name="setmodel")
async def set_model(ctx, *, model_name):
    current_model = get_current_model()
    if model_name == current_model:
        return await ctx.send(f"⚠️ `{model_name}` is already active.")

    await ctx.send(f"🔄 Switching from `{current_model}` to `{model_name}`…")

    if unload_model(current_model):
        await ctx.send(f"🧽 Unloaded `{current_model}` from VRAM.")
    else:
        await ctx.send(f"⚠️ Couldn’t unload `{current_model}`.")

    if not load_model(model_name):
        return await ctx.send(f"❌ Failed to pull `{model_name}`.")

    os.environ["MODEL_NAME"] = model_name
    env_path = os.path.join(os.path.dirname(__file__), '..', '.env')
    lines = []
    with open(env_path, 'r', encoding='utf-8') as f:
        for line in f:
            lines.append(f"MODEL_NAME={model_name}\n" if line.startswith("MODEL_NAME=") else line)
    with open(env_path, 'w', encoding='utf-8') as f:
        f.writelines(lines)

    await ctx.send(f"✅ Model switched to `{model_name}` and `.env` updated.")

@bot.command(name="models")
async def list_models(ctx):
    import requests
    try:
        resp = requests.get(TAGS_ENDPOINT)
        models = [m["name"] for m in resp.json().get("models", [])]
        if models:
            await ctx.send("🧠 Available models:\n" + "\n".join(f"- `{m}`" for m in models))
        else:
            await ctx.send("❌ No models found.")
    except Exception as e:
        await ctx.send(f"❌ Failed to fetch models: {e}")

@bot.command(name="setavatar")
@commands.is_owner()
async def set_avatar(ctx):
    if not ctx.message.attachments:
        return await ctx.send("❌ Please attach an image (PNG) to use as the new avatar.")

    image = ctx.message.attachments[0]
    image_bytes = await image.read()
    token = os.getenv("DISCORD_TOKEN")
    if not token:
        return await ctx.send("❌ Bot token not found in environment.")

    success = set_avatar_from_bytes(image_bytes, token)
    await ctx.send("✅ Avatar updated successfully!" if success else "❌ Failed to update avatar.")

# Run bot
bot.run(TOKEN)