322 lines
12 KiB
Python
322 lines
12 KiB
Python
import os
|
|
import re
|
|
import time
|
|
import asyncio
|
|
import discord
|
|
from discord.ext import commands
|
|
from collections import defaultdict, deque
|
|
from openai import AsyncOpenAI
|
|
|
|
#THIS IS DEEPSEEK-V4-FLASH, BUT CAN SWITCH TO V4-PRO
|
|
|
|
#only last 12 messages on discord will be fed into the ai's memory
|
|
MEMORY_SIZE = 12
|
|
|
|
#the very much needed ai shit
|
|
|
|
#loads persona file
|
|
def load_persona(name, tier='flash'):
|
|
suffix = '_prompt_pro.txt' if tier == 'pro' else '_prompt.txt'
|
|
path = os.path.join(os.path.dirname(__file__), 'personalities', f'{name.lower()}{suffix}')
|
|
if not os.path.exists(path):
|
|
# fallback to base prompt if pro version doesn't exist yet
|
|
path = os.path.join(os.path.dirname(__file__), 'personalities', f'{name.lower()}_prompt.txt')
|
|
if not os.path.exists(path):
|
|
return None
|
|
with open(path, 'r', encoding='utf-8') as f:
|
|
return f.read().strip()
|
|
|
|
#this command lists all personas
|
|
def list_personas():
|
|
names = []
|
|
persona_dir = os.path.join(os.path.dirname(__file__), 'personalities')
|
|
if os.path.exists(persona_dir):
|
|
for fname in os.listdir(persona_dir):
|
|
# only show base prompts, not _pro variants
|
|
if fname.endswith('_prompt.txt') and '_prompt_pro' not in fname:
|
|
names.append(fname.replace('_prompt.txt', ''))
|
|
return sorted(names)
|
|
|
|
#this splits messages into chunks if they are too long (optimization)
|
|
def split_msg(text, limit=2000):
|
|
chunks = []
|
|
while len(text) > limit:
|
|
split_at = text.rfind('\n', 0, limit)
|
|
if split_at == -1:
|
|
split_at = limit
|
|
chunks.append(text[:split_at])
|
|
text = text[split_at:].lstrip('\n')
|
|
if text:
|
|
chunks.append(text)
|
|
return chunks
|
|
|
|
class AIRoleplay(commands.Cog):
|
|
def __init__(self, bot):
|
|
self.bot = bot
|
|
self.client = AsyncOpenAI(
|
|
api_key=os.getenv('DEEPSEEK_API'),
|
|
base_url='https://api.deepseek.com'
|
|
)
|
|
|
|
self.allowed_users = set()
|
|
for key in ('OWNER_ID', 'SIDE_OWNER_ID'):
|
|
val = os.getenv(key)
|
|
if val:
|
|
try:
|
|
self.allowed_users.add(int(val))
|
|
except ValueError:
|
|
pass
|
|
|
|
self.memory = defaultdict(lambda: deque(maxlen=MEMORY_SIZE))
|
|
self.active_personas = {}
|
|
self.prompts = {}
|
|
|
|
self.input_tokens = 0
|
|
self.output_tokens = 0
|
|
self.session_start = time.time()
|
|
self.webhooks = {}
|
|
self.active_model = 'deepseek-v4-flash'
|
|
self.active_channels = set()
|
|
|
|
#only accepts me and hizuki
|
|
def is_allowed(self, user):
|
|
return user.id in self.allowed_users
|
|
|
|
#this loads the persona file
|
|
def get_prompt(self, channel_id):
|
|
tier = 'pro' if 'pro' in self.active_model else 'flash'
|
|
cache_key = (channel_id, tier)
|
|
if cache_key not in self.prompts:
|
|
name = self.active_personas.get(channel_id, 'yuuka')
|
|
prompt = load_persona(name, tier)
|
|
if not prompt:
|
|
prompt = load_persona('yuuka', tier) or ''
|
|
self.prompts[cache_key] = prompt
|
|
return self.prompts[cache_key]
|
|
|
|
#this gets the channel webhook, creates one if it doesnt exist
|
|
async def get_webhook(self, channel):
|
|
if channel.id in self.webhooks:
|
|
return self.webhooks[channel.id]
|
|
try:
|
|
whs = await channel.webhooks()
|
|
wh = next((w for w in whs if w.name == 'YuukaAI'), None)
|
|
if not wh:
|
|
wh = await channel.create_webhook(name='YuukaAI')
|
|
self.webhooks[channel.id] = wh
|
|
return wh
|
|
except discord.Forbidden:
|
|
return None
|
|
|
|
#this sends the message as the persona, with vision support (if images are present)
|
|
async def send_as_char(self, message, text, persona):
|
|
avatar_path = os.path.join(os.path.dirname(__file__), 'personalities', f'{persona.lower()}_avatar.png')
|
|
avatar_bytes = None
|
|
if os.path.exists(avatar_path):
|
|
with open(avatar_path, 'rb') as f:
|
|
avatar_bytes = f.read()
|
|
|
|
wh = await self.get_webhook(message.channel)
|
|
chunks = split_msg(text)
|
|
|
|
if wh:
|
|
display_name = persona.capitalize()
|
|
for i, chunk in enumerate(chunks):
|
|
kwargs = {'content': chunk, 'username': display_name, 'wait': True}
|
|
if i == 0 and avatar_bytes:
|
|
kwargs['avatar'] = avatar_bytes
|
|
await wh.send(**kwargs)
|
|
else:
|
|
first = True
|
|
for chunk in chunks:
|
|
if first:
|
|
await message.reply(chunk)
|
|
first = False
|
|
else:
|
|
await message.channel.send(chunk)
|
|
|
|
#this is the main listener, it handles all messages
|
|
@commands.Cog.listener()
|
|
async def on_message(self, message):
|
|
if message.author.bot or not message.guild:
|
|
return
|
|
|
|
mentioned = self.bot.user in message.mentions
|
|
in_active = message.channel.id in self.active_channels
|
|
|
|
# reject non-allowed users who actively ping the bot
|
|
if mentioned and not self.is_allowed(message.author):
|
|
await message.reply("You dont have access to AI features")
|
|
return
|
|
|
|
# determine if we should respond at all
|
|
if not mentioned and not in_active:
|
|
return
|
|
|
|
# in active mode: skip commands and ignore non-allowed users silently
|
|
if in_active and not mentioned:
|
|
if not self.is_allowed(message.author):
|
|
return
|
|
if message.content.startswith('!'):
|
|
return
|
|
|
|
#this removes the mention from the message
|
|
content = re.sub(r'<@!?' + str(self.bot.user.id) + r'>', '', message.content).strip()
|
|
if not content:
|
|
return
|
|
|
|
cid = message.channel.id
|
|
mem = self.memory[cid]
|
|
persona = self.active_personas.get(cid, 'yuuka')
|
|
prompt = self.get_prompt(cid)
|
|
|
|
# VISION SUPPORT
|
|
user_content = [{"type": "text", "text": f'{message.author.display_name}: {content}'}]
|
|
has_images = False
|
|
|
|
for att in message.attachments:
|
|
if att.content_type and att.content_type.startswith('image/'):
|
|
has_images = True
|
|
user_content.append({
|
|
"type": "image_url",
|
|
"image_url": {"url": att.url}
|
|
})
|
|
|
|
# DeepSeek and older models prefer pure strings if there are no images
|
|
final_content = user_content if has_images else user_content[0]["text"]
|
|
mem.append({'role': 'user', 'content': final_content})
|
|
|
|
api_msgs = [{'role': 'system', 'content': prompt}] + list(mem)
|
|
model_to_use = self.active_model
|
|
|
|
#VISION SUPPORT END
|
|
|
|
#this sends the message to the ai and handles the response
|
|
async with message.channel.typing():
|
|
try:
|
|
resp = await self.client.chat.completions.create(
|
|
model=model_to_use,
|
|
messages=api_msgs,
|
|
max_tokens=1024
|
|
)
|
|
except Exception as e:
|
|
await message.reply(f'api error: {e}')
|
|
mem.pop()
|
|
return
|
|
|
|
reply = resp.choices[0].message.content.strip()
|
|
|
|
if resp.usage:
|
|
self.input_tokens += resp.usage.prompt_tokens
|
|
self.output_tokens += resp.usage.completion_tokens
|
|
|
|
mem.append({'role': 'assistant', 'content': reply})
|
|
|
|
# simple typing delay calculation for better rp
|
|
words = len(reply.split())
|
|
delay = max(1.5, min(words / 60 * 2, 5.0))
|
|
await asyncio.sleep(delay)
|
|
|
|
await self.send_as_char(message, reply, persona)
|
|
|
|
#this command lists all the ai commands, not an actual ai command##
|
|
@commands.command()
|
|
async def ai(self, ctx):
|
|
ai_commands = [f'!{cmd.name}' for cmd in self.get_commands()]
|
|
await ctx.reply('\n'.join(ai_commands))
|
|
######################################################################
|
|
|
|
#this sets the persona, only usable by me and hizu
|
|
@commands.command()
|
|
async def setpersona(self, ctx, name: str):
|
|
if not self.is_allowed(ctx.author):
|
|
return
|
|
|
|
prompt = load_persona(name)
|
|
if not prompt:
|
|
avail = ', '.join(list_personas()) or 'none'
|
|
return await ctx.reply(f'couldnt find {name}. available: {avail}')
|
|
|
|
self.active_personas[ctx.channel.id] = name.lower()
|
|
self.prompts.pop(ctx.channel.id, None)
|
|
self.memory[ctx.channel.id].clear()
|
|
await ctx.reply(f'switched to {name} and wiped memory')
|
|
|
|
#this clears the memory of the ai, only usable by me and hizu
|
|
@commands.command()
|
|
async def clearchat(self, ctx):
|
|
if not self.is_allowed(ctx.author):
|
|
return
|
|
self.memory[ctx.channel.id].clear()
|
|
await ctx.reply('memory cleared')
|
|
|
|
#this enables active chat, only usable by me and hizu
|
|
@commands.command()
|
|
async def activechat(self, ctx):
|
|
if not self.is_allowed(ctx.author):
|
|
return
|
|
self.active_channels.add(ctx.channel.id)
|
|
await ctx.reply('active chat enabled — no need to ping, just talk')
|
|
|
|
#this disables active chat, only usable by me and hizu
|
|
@commands.command()
|
|
async def stopchat(self, ctx):
|
|
if not self.is_allowed(ctx.author):
|
|
return
|
|
self.active_channels.discard(ctx.channel.id)
|
|
await ctx.reply('active chat disabled')
|
|
|
|
#this lists all the personas, only usable by me and hizu
|
|
@commands.command()
|
|
async def personas(self, ctx):
|
|
if not self.is_allowed(ctx.author):
|
|
return
|
|
names = list_personas()
|
|
current = self.active_personas.get(ctx.channel.id, 'yuuka')
|
|
lines = [f'- {n} {"(active)" if n == current else ""}' for n in names]
|
|
await ctx.reply('**personas:**\n' + '\n'.join(lines) if lines else 'none found')
|
|
|
|
#this shows the current token usage, only usable by me and hizu
|
|
@commands.command()
|
|
async def aicost(self, ctx):
|
|
if not self.is_allowed(ctx.author):
|
|
return
|
|
|
|
elapsed = time.time() - self.session_start
|
|
hours, r = divmod(int(elapsed), 3600)
|
|
mins, _ = divmod(r, 60)
|
|
|
|
cost_in = self.input_tokens * (0.14 / 1_000_000)
|
|
cost_out = self.output_tokens * (0.28 / 1_000_000)
|
|
total = cost_in + cost_out
|
|
|
|
embed = discord.Embed(title='ai usage', color=0x00ff00)
|
|
embed.add_field(name='uptime', value=f'{hours}h {mins}m', inline=False)
|
|
embed.add_field(name='tokens', value=f'in: {self.input_tokens:,}\nout: {self.output_tokens:,}', inline=True)
|
|
embed.add_field(name='cost', value=f'~${total:.5f}', inline=True)
|
|
await ctx.reply(embed=embed)
|
|
|
|
#THIS IS WHERE TURBO COMMANDS LIVE IN
|
|
@commands.command()
|
|
async def turbo(self, ctx):
|
|
if not self.is_allowed(ctx.author):
|
|
return
|
|
if self.active_model == 'deepseek-v4-flash':
|
|
self.active_model = 'deepseek-v4-pro'
|
|
self.prompts.clear()
|
|
await ctx.reply('switched to **deepseek-v4-pro**')
|
|
else:
|
|
self.active_model = 'deepseek-v4-flash'
|
|
self.prompts.clear()
|
|
await ctx.reply('switched to **deepseek-v4-flash**')
|
|
|
|
#this shows the current model, only usable by me and hizu
|
|
@commands.command()
|
|
async def current(self, ctx):
|
|
if not self.is_allowed(ctx.author):
|
|
return
|
|
await ctx.reply(f'currently using: **{self.active_model}**')
|
|
|
|
async def setup(bot):
|
|
await bot.add_cog(AIRoleplay(bot))
|