KaniumCogs/reginaldCog/reginald.py

import discord
import openai
import random
import asyncio
import traceback
from redbot.core import Config, commands
from openai import OpenAIError

class ReginaldCog(commands.Cog):
    def __init__(self, bot):
        self.bot = bot
        self.config = Config.get_conf(self, identifier=71717171171717)
        self.memory_locks = {}  # ✅ Prevents race conditions per channel

        # ✅ Properly Registered Configuration Keys
        default_global = {"openai_model": "gpt-4o-mini"}
        default_guild = {
            "openai_api_key": None,
            "short_term_memory": {},  # Tracks last 100 messages per channel
            "mid_term_memory": {},  # Stores condensed summaries
            "long_term_profiles": {},  # Stores persistent knowledge
            "admin_role": None,
            "allowed_role": None
        }
        self.config.register_global(**default_global)
        self.config.register_guild(**default_guild)
    async def is_admin(self, ctx):
        """✅ Checks if the user is an admin (or has an assigned admin role)."""
        admin_role_id = await self.config.guild(ctx.guild).admin_role()
        if admin_role_id:
            return any(role.id == admin_role_id for role in ctx.author.roles)
        return ctx.author.guild_permissions.administrator

    async def is_allowed(self, ctx):
        """✅ Checks if the user is allowed to use Reginald based on role settings."""
        allowed_role_id = await self.config.guild(ctx.guild).allowed_role()
        return any(role.id == allowed_role_id for role in ctx.author.roles) if allowed_role_id else False

    @commands.command(name="reginald", help="Ask Reginald a question in shared channels")
    @commands.cooldown(1, 10, commands.BucketType.user)
    async def reginald(self, ctx, *, prompt=None):
        """Handles multi-user memory tracking in shared channels, recognizing mentions properly."""

        if not await self.is_admin(ctx) and not await self.is_allowed(ctx):
            await ctx.send("You do not have the required role to use this command.")
            return

        if prompt is None:
            await ctx.send(random.choice(["Yes?", "How may I assist?", "You rang?"]))
            return

        api_key = await self.config.guild(ctx.guild).openai_api_key()
        if not api_key:
            await ctx.send("OpenAI API key not set. Use `!setreginaldcogapi`.")
            return

        channel_id = str(ctx.channel.id)
        user_id = str(ctx.author.id)
        user_name = ctx.author.display_name  # Uses Discord nickname if available

        # ✅ Convert mentions into readable names
        for mention in ctx.message.mentions:
            prompt = prompt.replace(f"<@{mention.id}>", mention.display_name)

        # ✅ Ensure only one update per channel at a time
        if channel_id not in self.memory_locks:
            self.memory_locks[channel_id] = asyncio.Lock()

        async with self.memory_locks[channel_id]:  # ✅ Prevent race conditions
            async with self.config.guild(ctx.guild).short_term_memory() as short_memory, \
                       self.config.guild(ctx.guild).mid_term_memory() as mid_memory, \
                       self.config.guild(ctx.guild).long_term_profiles() as long_memory:

                # ✅ Retrieve memory
                memory = short_memory.get(channel_id, [])
                user_profile = long_memory.get(user_id, {})

                # ✅ Format messages properly
                formatted_messages = [{"role": "system", "content": (
                    "You are Reginald, the esteemed butler of The Kanium Estate. "
                    "This estate is home to Lords, Ladies, and distinguished guests, each with unique personalities and demands. "
                    "Your duty is to uphold decorum while providing assistance with wit and intelligence. "
                    "You should recognize individual names and use them sparingly. Prefer natural conversation flow—do not force names where unnecessary."
                )}]

                # ✅ Add long-term knowledge if available
                if user_profile:
                    knowledge_summary = f"Previous knowledge about {user_name}: {user_profile.get('summary', 'No detailed memory yet.')}"
                    formatted_messages.append({"role": "system", "content": knowledge_summary})

                # ✅ Add recent conversation history
                formatted_messages += [{"role": "user", "content": f"{entry['user']}: {entry['content']}"} for entry in memory]
                formatted_messages.append({"role": "user", "content": f"{user_name}: {prompt}"})

                # ✅ Generate response
                response_text = await self.generate_response(api_key, formatted_messages)

                # ✅ Store new messages in memory
                memory.append({"user": user_name, "content": prompt})  # Store user message
                memory.append({"user": "Reginald", "content": response_text})  # Store response

                # ✅ Keep memory within limit
                if len(memory) > 100:
                    summary = await self.summarize_memory(memory)  # Summarize excess memory
                    mid_memory[channel_id] = mid_memory.get(channel_id, "") + "\n" + summary  # Store in Mid-Term Memory
                    memory = memory[-100:]  # Prune old memory

                short_memory[channel_id] = memory  # ✅ Atomic update inside async context

        await ctx.send(response_text[:2000])  # Discord character limit safeguard

    async def summarize_memory(self, messages):
        """✅ Generates a summary of past conversations for mid-term storage."""
        summary_prompt = (
        "Analyze and summarize the following conversation in a way that retains key details, nuances, and unique insights. "
        "Your goal is to create a structured yet fluid summary that captures important points without oversimplifying. "
        "Maintain resolution on individual opinions, preferences, debates, and shared knowledge. "
        "If multiple topics are discussed, summarize each distinctly rather than blending them together."
        )

        summary_text = "\n".join(f"{msg['user']}: {msg['content']}" for msg in messages)

        try:
            client = openai.AsyncClient(api_key=await self.config.openai_model())
            response = await client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[{"role": "system", "content": summary_prompt}, {"role": "user", "content": summary_text}],
                max_tokens=256
            )
            return response.choices[0].message.content.strip()
        except OpenAIError:
            return "Summary unavailable due to an error."


    async def generate_response(self, api_key, messages):
        """✅ Generates a response using OpenAI's new async API client (OpenAI v1.0+)."""
        model = await self.config.openai_model()
        try:
            client = openai.AsyncClient(api_key=api_key)  # ✅ Correct API usage
            response = await client.chat.completions.create(
                model=model,
                messages=messages,
                max_tokens=1024,
                temperature=0.7,
                presence_penalty=0.5,
                frequency_penalty=0.5
            )

            if not response.choices:
                return "I fear I have no words to offer at this time."

            return response.choices[0].message.content.strip()

        except OpenAIError as e:
            error_message = f"OpenAI Error: {e}"
            reginald_responses = [
                f"Regrettably, I must inform you that I have encountered a bureaucratic obstruction:\n\n```{error_message}```",
                f"It would seem that a most unfortunate technical hiccup has befallen my faculties:\n\n```{error_message}```",
                f"Ah, it appears I have received an urgent memorandum stating:\n\n```{error_message}```",
                f"I regret to inform you that my usual eloquence is presently obstructed by an unforeseen complication:\n\n```{error_message}```"
            ]
            return random.choice(reginald_responses)

    @commands.command(name="reginald_allowrole", help="Allow a role to use the Reginald command")
    @commands.has_permissions(administrator=True)
    async def allow_role(self, ctx, role: discord.Role):
        """✅ Grants permission to a role to use Reginald."""
        await self.config.guild(ctx.guild).allowed_role.set(role.id)
        await ctx.send(f"The role `{role.name}` (ID: `{role.id}`) is now allowed to use the Reginald command.")

    @commands.command(name="reginald_disallowrole", help="Remove a role's ability to use the Reginald command")
    @commands.has_permissions(administrator=True)
    async def disallow_role(self, ctx):
        """✅ Removes a role's permission to use Reginald."""
        await self.config.guild(ctx.guild).allowed_role.clear()
        await ctx.send("The role's permission to use the Reginald command has been revoked.")
        
    @commands.guild_only()
    @commands.has_permissions(manage_guild=True)
    @commands.command(help="Set the OpenAI API key")
    async def setreginaldcogapi(self, ctx, api_key):
        """Allows an admin to set the OpenAI API key for Reginald."""
        await self.config.guild(ctx.guild).openai_api_key.set(api_key)
        await ctx.send("OpenAI API key set successfully.")

async def setup(bot):
    """✅ Correct async cog setup for Redbot"""
    await bot.add_cog(ReginaldCog(bot))