major ai update

ai/ai_roleplay.py

@@ -0,0 +1,321 @@
+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))