276 lines
9.4 KiB
Python
276 lines
9.4 KiB
Python
# bot.py
|
||
|
||
import time
|
||
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 time_logger import log_message_activity
|
||
from autochat import should_auto_reply, generate_auto_reply, update_reply_timer
|
||
|
||
debug_mode = os.getenv("DEBUG_MODE", "false").lower() == "true"
|
||
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}")
|
||
if debug_mode:
|
||
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):
|
||
|
||
if not message.author.bot:
|
||
log_message_activity(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)
|
||
profile = load_user_profile(message.author)
|
||
|
||
# Log summary info about the profile (but don’t inject it)
|
||
logger.info("=" * 60 + " AI Response " + "=" * 60)
|
||
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)
|
||
|
||
elif should_auto_reply():
|
||
reply = await generate_auto_reply(message, bot)
|
||
if reply and len(reply) > 5:
|
||
await message.channel.send(reply)
|
||
update_reply_timer()
|
||
|
||
# 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)
|