153 lines
5.2 KiB
Python
153 lines
5.2 KiB
Python
import discord
|
|
import asyncio
|
|
import atexit
|
|
import os
|
|
import threading
|
|
from dotenv import load_dotenv
|
|
from datetime import datetime, timedelta
|
|
from dashboard import start_dashboard
|
|
from tokenizer import Tokenizer
|
|
from trainer import RubyTrainer
|
|
from reader import BookReader
|
|
import logging
|
|
|
|
# Setup logging
|
|
logging.basicConfig(
|
|
filename="logs/error.log",
|
|
level=logging.ERROR,
|
|
format="%(asctime)s %(levelname)s: %(message)s",
|
|
encoding="utf-8"
|
|
)
|
|
|
|
# Disable Flask, Werkzeug, and other noisy loggers
|
|
for noisy_logger in ["werkzeug", "flask", "flask.app"]:
|
|
logging.getLogger(noisy_logger).setLevel(logging.CRITICAL)
|
|
|
|
# Load environment
|
|
load_dotenv()
|
|
TOKEN = os.getenv("DISCORD_TOKEN")
|
|
|
|
if not TOKEN:
|
|
raise RuntimeError("Bot token not found in .env")
|
|
|
|
# Setup intents
|
|
intents = discord.Intents.default()
|
|
intents.message_content = True
|
|
intents.dm_messages = True
|
|
intents = intents
|
|
|
|
|
|
class Ruby(discord.Client):
|
|
def __init__(self):
|
|
super().__init__(intents=intents)
|
|
self.tokenizer = Tokenizer()
|
|
self.trainer = RubyTrainer(self.tokenizer)
|
|
self.reader = BookReader(trainer=self.trainer,
|
|
book_path="books//wizard_of_oz.txt", # or whatever book you want
|
|
interval=180 # read every 3 minutes (adjust if needed)
|
|
)
|
|
self.last_message_time = datetime.utcnow()
|
|
self.idle_threshold = timedelta(seconds=120)
|
|
self.log_path = os.path.join("logs", "messages.log")
|
|
os.makedirs("logs", exist_ok=True)
|
|
|
|
async def setup_hook(self):
|
|
self.loop.create_task(self.reader.start_reading())
|
|
self.loop.create_task(self.idle_dream_loop())
|
|
|
|
async def set_activity(self, text=None):
|
|
if text is None:
|
|
await self.change_presence(status=discord.Status.online, activity=None)
|
|
else:
|
|
activity = discord.Activity(type=discord.ActivityType.listening, name=text)
|
|
await self.change_presence(status=discord.Status.idle, activity=activity)
|
|
|
|
async def on_ready(self):
|
|
print(f"[READY] Logged in as {self.user} (ID: {self.user.id})")
|
|
await self.set_activity("you...")
|
|
self.trainer.reinforce_core_memory()
|
|
|
|
async def idle_dream_loop(self):
|
|
await self.wait_until_ready()
|
|
while not self.is_closed():
|
|
now = datetime.utcnow()
|
|
if now - self.last_message_time > self.idle_threshold:
|
|
print("[IDLE] Ruby has been idle — entering dream mode.")
|
|
|
|
await self.set_activity("the past...")
|
|
try:
|
|
self.trainer.dream()
|
|
except Exception as e:
|
|
logging.error("Error dreaming: %s", e)
|
|
|
|
await self.set_activity("my thoughts")
|
|
from random import random
|
|
speak = random() < 0.5
|
|
thought = self.trainer.daydream(say_thought=speak)
|
|
|
|
if speak and thought and len(thought.split()) >=4:
|
|
for guild in self.guilds:
|
|
for channel in guild.text_channels:
|
|
if channel.permissions_for(guild.me).send_messages:
|
|
if not thought.endswith("."):
|
|
thought += "."
|
|
await channel.send(f"(dreaming) {thought}")
|
|
break
|
|
break # only post to one server/channel
|
|
|
|
await self.set_activity(None) # reset to normal
|
|
self.last_message_time = datetime.utcnow()
|
|
|
|
await asyncio.sleep(180)
|
|
|
|
async def on_message(self, message: discord.Message):
|
|
if message.author.id == self.user.id:
|
|
return
|
|
|
|
self.log_message(message)
|
|
self.trainer.train_on_tokens_from_text(message.content.strip())
|
|
|
|
reply = self.trainer.generate_reply()
|
|
if reply.strip():
|
|
await message.channel.send(reply)
|
|
else:
|
|
print("[REPLY] Skipped (empty)")
|
|
|
|
def log_message(self, message: discord.Message):
|
|
timestamp = datetime.utcnow().isoformat()
|
|
log_entry = f"{timestamp} | {message.author.name} | {message.content.strip()}\n"
|
|
|
|
with open(self.log_path, "a", encoding="utf-8") as f:
|
|
f.write(log_entry)
|
|
|
|
print(f"[LOGGED] {log_entry.strip()}")
|
|
|
|
def train_on_message(self, message: discord.Message):
|
|
text = message.content.strip()
|
|
self.trainer.train_on_tokens_from_text(text)
|
|
token_tensor = torch.tensor(tokens, dtype=torch.long)
|
|
loss = train_on_tokens(self.model, tokens, self.optimizer, self.criterion, device="cpu")
|
|
print(f"[TRAIN] Tokens: {tokens} | Loss: {loss:.4f}")
|
|
|
|
|
|
# Run Ruby
|
|
client = None
|
|
try:
|
|
client = Ruby()
|
|
|
|
def on_exit():
|
|
if client:
|
|
print("[EXIT] Ruby is gracefully shutting down...")
|
|
client.trainer.dream()
|
|
client.trainer.daydream(rounds=10)
|
|
|
|
atexit.register(on_exit)
|
|
dashboard_thread = threading.Thread(target=start_dashboard, daemon=True)
|
|
dashboard_thread.start()
|
|
client.run(TOKEN)
|
|
finally:
|
|
if client is not None:
|
|
print("[EXIT] Ruby is shutting down — dreaming one last time...")
|
|
client.trainer.dream()
|
|
client.trainer.daydream(rounds=10)
|