GPT3Discord/cogs/moderations_service_cog.py

import asyncio

import discord
from sqlitedict import SqliteDict

from services.environment_service import EnvService
from services.moderations_service import Moderation, ThresholdSet

MOD_DB = None
try:
    print("Attempting to retrieve the General and Moderations DB")
    MOD_DB = SqliteDict("main_db.sqlite", tablename="moderations", autocommit=True)
except Exception as e:
    print("Failed to retrieve the General and Moderations DB")
    raise e


class ModerationsService(discord.Cog, name="ModerationsService"):
    """Cog containing moderation tools and features"""

    def __init__(
        self,
        bot,
        usage_service,
        model,
    ):
        super().__init__()
        self.bot = bot
        self.usage_service = usage_service
        self.model = model

        # Moderation service data
        self.moderation_queues = {}
        self.moderation_alerts_channel = EnvService.get_moderations_alert_channel()
        self.moderation_enabled_guilds = []
        self.moderation_tasks = {}
        self.moderations_launched = []

        # Defaults
        self.default_warn_set = ThresholdSet(0.01, 0.05, 0.05, 0.91, 0.1, 0.45, 0.1)
        self.default_delete_set = ThresholdSet(0.26, 0.26, 0.1, 0.95, 0.03, 0.85, 0.4)

    @discord.Cog.listener()
    async def on_ready(self):
        """Check moderation service for each guild"""
        for guild in self.bot.guilds:
            self.get_or_set_warn_set(guild.id)
            self.get_or_set_delete_set(guild.id)
            await self.check_and_launch_moderations(guild.id)
        print("The moderation service is ready.")

    def check_guild_moderated(self, guild_id):
        """Given guild id, return bool of moderation status"""
        return guild_id in MOD_DB and MOD_DB[guild_id]["moderated"]

    def get_moderated_alert_channel(self, guild_id):
        """Given guild id, return alert channel"""
        return MOD_DB[guild_id]["alert_channel"]

    def set_moderated_alert_channel(self, guild_id, channel_id):
        """Given guild id and channel id, set channel to recieve alerts"""
        MOD_DB[guild_id] = {"moderated": True, "alert_channel": channel_id}
        MOD_DB.commit()

    def get_or_set_warn_set(self, guild_id):
        """Get warn_set set for the guild, if not set them from default values"""
        guild_id = str(guild_id)
        key = guild_id + "_warn_set"
        if key not in MOD_DB:
            MOD_DB[key] = zip(
                self.default_warn_set.keys, self.default_warn_set.thresholds
            )
            MOD_DB.commit()
        return dict(MOD_DB[key])

    def get_or_set_delete_set(self, guild_id):
        """Get delete_set set for the guild, if not set them from default values"""
        guild_id = str(guild_id)
        key = guild_id + "_delete_set"
        if key not in MOD_DB:
            MOD_DB[key] = zip(
                self.default_delete_set.keys, self.default_delete_set.thresholds
            )
            MOD_DB.commit()
        return dict(MOD_DB[key])

    def set_warn_set(self, guild_id, threshold_set):
        """Set threshold for warning a message"""
        guild_id = str(guild_id)
        key = guild_id + "_warn_set"
        MOD_DB[key] = zip(threshold_set.keys, threshold_set.thresholds)
        MOD_DB.commit()

    def set_delete_set(self, guild_id, threshold_set):
        """Set threshold for deleting a message"""
        guild_id = str(guild_id)
        key = guild_id + "_delete_set"
        MOD_DB[key] = zip(threshold_set.keys, threshold_set.thresholds)
        MOD_DB.commit()

    def set_guild_moderated(self, guild_id, status=True):
        """Set the guild to moderated or not"""
        if guild_id not in MOD_DB:
            MOD_DB[guild_id] = {"moderated": status, "alert_channel": 0}
            MOD_DB.commit()
            return
        MOD_DB[guild_id] = {
            "moderated": status,
            "alert_channel": self.get_moderated_alert_channel(guild_id),
        }
        MOD_DB.commit()

    async def check_and_launch_moderations(self, guild_id, alert_channel_override=None):
        """Create the moderation service"""
        if self.check_guild_moderated(guild_id):
            Moderation.moderation_queues[guild_id] = asyncio.Queue()

            moderations_channel = await self.bot.fetch_channel(
                self.get_moderated_alert_channel(guild_id)
                if not alert_channel_override
                else alert_channel_override
            )
            warn_set_nums = self.get_or_set_warn_set(guild_id).values()
            delete_set_nums = self.get_or_set_delete_set(guild_id).values()
            warn_set = ThresholdSet(*warn_set_nums)
            delete_set = ThresholdSet(*delete_set_nums)

            Moderation.moderation_tasks[guild_id] = asyncio.ensure_future(
                Moderation.process_moderation_queue(
                    Moderation.moderation_queues[guild_id],
                    0.25,
                    0.25,
                    moderations_channel,
                    warn_set,
                    delete_set,
                )
            )
            print("Launched the moderations service for guild " + str(guild_id))
            Moderation.moderations_launched.append(guild_id)
            return moderations_channel

        return None

    async def moderations_command(
        self, ctx: discord.ApplicationContext, status: str, alert_channel_id: str
    ):
        """command handler for toggling moderation and setting an alert channel"""
        await ctx.defer()

        try:
            if alert_channel_id:
                int(alert_channel_id)
        except ValueError:
            # the alert_channel_id was passed in as a channel NAME instead of an ID, fetch the ID.
            alert_channel = discord.utils.get(ctx.guild.channels, name=alert_channel_id)
            alert_channel_id = alert_channel.id

        if status == "on":
            # Check if the current guild is already in the database and if so, if the moderations is on
            if self.check_guild_moderated(ctx.guild_id):
                await ctx.respond("Moderations is already enabled for this guild")
                return

            # Create the moderations service.
            await self.start_moderations_service(
                guild_id=ctx.guild_id, alert_channel_id=alert_channel_id
            )
            await ctx.respond("Moderations is now enabled for this guild")

        elif status == "off":
            # Cancel the moderations service.
            await self.stop_moderations_service(ctx.guild_id)
            await ctx.respond(
                "Moderations is now disabled for this guild", ephemeral=True
            )

    async def stop_moderations_service(self, guild_id):
        """Remove guild moderation status and stop the service"""
        self.set_guild_moderated(guild_id, False)
        Moderation.moderation_tasks[guild_id].cancel()
        Moderation.moderation_tasks[guild_id] = None
        Moderation.moderation_queues[guild_id] = None
        Moderation.moderations_launched.remove(guild_id)

    async def start_moderations_service(self, guild_id, alert_channel_id=None):
        """Set guild moderation and start the service"""
        self.set_guild_moderated(guild_id)
        moderations_channel = await self.check_and_launch_moderations(
            guild_id,
            Moderation.moderation_alerts_channel
            if not alert_channel_id
            else alert_channel_id,
        )
        self.set_moderated_alert_channel(guild_id, moderations_channel.id)

    async def restart_moderations_service(self, ctx):
        """restarts the moderation of the guild it's run in"""
        if not self.check_guild_moderated(ctx.guild_id):
            await ctx.respond(
                "Moderations are not enabled, can't restart",
                ephemeral=True,
                delete_after=30,
            )
            return

        await ctx.respond(
            "The moderations service is being restarted...",
            ephemeral=True,
            delete_after=30,
        )
        await self.stop_moderations_service(ctx.guild_id)
        await ctx.send_followup(
            "The moderations service was stopped..", ephemeral=True, delete_after=30
        )
        await self.start_moderations_service(
            ctx.guild_id, self.get_moderated_alert_channel(ctx.guild_id)
        )
        await ctx.send_followup(
            "The moderations service was restarted successfully.",
            ephemeral=True,
            delete_after=30,
        )

    async def build_moderation_settings_embed(self, category, mod_set):

        embed = discord.Embed(
            title="Moderation Settings",
            description="The moderation settings for this guild for the type: "
            + category,
            color=discord.Color.yellow() if type == "warn" else discord.Color.red(),
        )

        # Add each key_value pair in the mod_set to the embed, make them fairly small
        for key, value in mod_set.items():
            embed.add_field(name=key, value=value, inline=False)

        return embed

    async def config_command(
        self,
        ctx: discord.ApplicationContext,
        config_type: str,
        hate,
        hate_threatening,
        self_harm,
        sexual,
        sexual_minors,
        violence,
        violence_graphic,
    ):
        """command handler for assigning threshold values for warn or delete"""
        all_args = [
            hate,
            hate_threatening,
            self_harm,
            sexual,
            sexual_minors,
            violence,
            violence_graphic,
        ]
        await ctx.defer(ephemeral=True)

        # Case for printing the current config
        if not any(all_args) and config_type != "reset":
            await ctx.respond(
                ephemeral=True,
                embed=await self.build_moderation_settings_embed(
                    config_type,
                    self.get_or_set_warn_set(ctx.guild_id)
                    if config_type == "warn"
                    else self.get_or_set_delete_set(ctx.guild_id),
                ),
            )
            return

        if config_type == "warn":
            # Check if no args were
            warn_set = self.get_or_set_warn_set(ctx.guild_id)

            new_warn_set = ThresholdSet(
                hate if hate else warn_set["hate"],
                hate_threatening if hate_threatening else warn_set["hate/threatening"],
                self_harm if self_harm else warn_set["self-harm"],
                sexual if sexual else warn_set["sexual"],
                sexual_minors if sexual_minors else warn_set["sexual/minors"],
                violence if violence else warn_set["violence"],
                violence_graphic if violence_graphic else warn_set["violence/graphic"],
            )
            self.set_warn_set(ctx.guild_id, new_warn_set)
            await self.restart_moderations_service(ctx)

        elif config_type == "delete":
            delete_set = self.get_or_set_delete_set(ctx.guild_id)

            new_delete_set = ThresholdSet(
                hate if hate else delete_set["hate"],
                hate_threatening
                if hate_threatening
                else delete_set["hate/threatening"],
                self_harm if self_harm else delete_set["self-harm"],
                sexual if sexual else delete_set["sexual"],
                sexual_minors if sexual_minors else delete_set["sexual/minors"],
                violence if violence else delete_set["violence"],
                violence_graphic
                if violence_graphic
                else delete_set["violence/graphic"],
            )
            self.set_delete_set(ctx.guild_id, new_delete_set)
            await self.restart_moderations_service(ctx)

        elif config_type == "reset":
            self.set_delete_set(ctx.guild_id, self.default_delete_set)
            self.set_warn_set(ctx.guild_id, self.default_warn_set)
            await self.restart_moderations_service(ctx)

    async def moderations_test_command(
        self, ctx: discord.ApplicationContext, prompt: str
    ):
        """command handler for checking moderation values of a given input"""
        await ctx.defer()
        response = await self.model.send_moderations_request(prompt)
        await ctx.respond(response["results"][0]["category_scores"])
        await ctx.send_followup(response["results"][0]["flagged"])