diff --git a/README.md b/README.md index 53e2e6f..ccf7bfd 100644 --- a/README.md +++ b/README.md @@ -121,11 +121,11 @@ These commands are grouped, so each group has a prefix but you can easily tab co ### Automatic AI Moderation -`/system moderations status:on` - Turn on automatic chat moderations. +`/mod set status:on` - Turn on automatic chat moderations. -`/system moderations status:off` - Turn off automatic chat moderations +`/mod set status:off` - Turn off automatic chat moderations -`/system moderations status:off alert_channel_id:` - Turn on moderations and set the alert channel to the channel ID you specify in the command. +`/mod set status:off alert_channel_id:` - Turn on moderations and set the alert channel to the channel ID you specify in the command. - The bot needs Administrative permissions for this, and you need to set `MODERATIONS_ALERT_CHANNEL` to the channel ID of a desired channel in your .env file if you want to receive alerts about moderated messages. - This uses the OpenAI Moderations endpoint to check for messages, requests are only sent to the moderations endpoint at a MINIMUM request gap of 0.5 seconds, to ensure you don't get blocked and to ensure reliability. diff --git a/cogs/commands.py b/cogs/commands.py index b9a02c6..6815f27 100644 --- a/cogs/commands.py +++ b/cogs/commands.py @@ -20,6 +20,7 @@ class Commands(discord.Cog, name="Commands"): converser_cog, image_draw_cog, image_service_cog, + moderations_cog, ): super().__init__() self.bot = bot @@ -30,6 +31,7 @@ class Commands(discord.Cog, name="Commands"): self.converser_cog = converser_cog self.image_draw_cog = image_draw_cog self.image_service_cog = image_service_cog + self.moderations_cog = moderations_cog # Create slash command groups dalle = discord.SlashCommandGroup( @@ -50,6 +52,12 @@ class Commands(discord.Cog, name="Commands"): guild_ids=ALLOWED_GUILDS, checks=[Check.check_admin_roles()], ) + mod = discord.SlashCommandGroup( + name="mod", + description="AI-Moderation commands for the bot", + guild_ids=ALLOWED_GUILDS, + checks=[Check.check_admin_roles()], + ) """ System commands @@ -140,9 +148,9 @@ class Commands(discord.Cog, name="Commands"): (system) Moderation commands """ - @add_to_group("system") + @add_to_group("mod") @discord.slash_command( - name="moderations-test", + name="test", description="Used to test a prompt and see what threshold values are returned by the moderations endpoint", guild_ids=ALLOWED_GUILDS, ) @@ -153,13 +161,13 @@ class Commands(discord.Cog, name="Commands"): ) @discord.guild_only() async def moderations_test(self, ctx: discord.ApplicationContext, prompt: str): - await self.converser_cog.moderations_test_command(ctx, prompt) + await self.moderations_cog.moderations_test_command(ctx, prompt) - @add_to_group("system") + @add_to_group("mod") @discord.slash_command( - name="moderations", - description="The AI moderations service", + name="set", + description="Turn the moderations service on and off", guild_ids=ALLOWED_GUILDS, ) @discord.option( @@ -176,7 +184,7 @@ class Commands(discord.Cog, name="Commands"): async def moderations( self, ctx: discord.ApplicationContext, status: str, alert_channel_id: str ): - await self.converser_cog.moderations_command(ctx, status, alert_channel_id) + await self.moderations_cog.moderations_command(ctx, status, alert_channel_id) """ GPT commands diff --git a/cogs/moderations_service_cog.py b/cogs/moderations_service_cog.py new file mode 100644 index 0000000..f44d48f --- /dev/null +++ b/cogs/moderations_service_cog.py @@ -0,0 +1,126 @@ +import asyncio + +import discord +from sqlitedict import SqliteDict + +from services.environment_service import EnvService +from services.moderations_service import Moderation + +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"): + 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 = [] + @discord.Cog.listener() + async def on_ready(self): + # Check moderation service for each guild + for guild in self.bot.guilds: + await self.check_and_launch_moderations(guild.id) + + def check_guild_moderated(self, guild_id): + return guild_id in MOD_DB and MOD_DB[guild_id]["moderated"] + + def get_moderated_alert_channel(self, guild_id): + return MOD_DB[guild_id]["alert_channel"] + + def set_moderated_alert_channel(self, guild_id, channel_id): + MOD_DB[guild_id] = {"moderated": True, "alert_channel": channel_id} + MOD_DB.commit() + + def set_guild_moderated(self, guild_id, status=True): + 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 moderations service. + print("Checking and attempting to launch moderations 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 + ) + + Moderation.moderation_tasks[guild_id] = asyncio.ensure_future( + Moderation.process_moderation_queue( + Moderation.moderation_queues[guild_id], 1, 1, moderations_channel + ) + ) + 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 + ): + await ctx.defer() + + status = status.lower().strip() + if status not in ["on", "off"]: + await ctx.respond("Invalid status, please use on or off") + return + + 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. + self.set_guild_moderated(ctx.guild_id) + moderations_channel = await self.check_and_launch_moderations( + ctx.guild_id, + Moderation.moderation_alerts_channel + if not alert_channel_id + else alert_channel_id, + ) + self.set_moderated_alert_channel(ctx.guild_id, moderations_channel.id) + + await ctx.respond("Moderations service enabled") + + elif status == "off": + # Cancel the moderations service. + self.set_guild_moderated(ctx.guild_id, False) + Moderation.moderation_tasks[ctx.guild_id].cancel() + Moderation.moderation_tasks[ctx.guild_id] = None + Moderation.moderation_queues[ctx.guild_id] = None + Moderation.moderations_launched.remove(ctx.guild_id) + await ctx.respond("Moderations service disabled") + + async def moderations_test_command(self, ctx: discord.ApplicationContext, prompt: str): + 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"]) + diff --git a/cogs/text_service_cog.py b/cogs/text_service_cog.py index bcba0be..a32beae 100644 --- a/cogs/text_service_cog.py +++ b/cogs/text_service_cog.py @@ -112,12 +112,6 @@ class GPT3ComCon(discord.Cog, name="GPT3ComCon"): self.conversation_threads = {} self.summarize = self.model.summarize_conversations - # 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 = [] # Pinecone data self.pinecone_service = pinecone_service @@ -207,10 +201,6 @@ class GPT3ComCon(discord.Cog, name="GPT3ComCon"): ) print("The debug channel was acquired") - # Check moderation service for each guild - for guild in self.bot.guilds: - await self.check_and_launch_moderations(guild.id) - await self.bot.sync_commands( commands=None, method="individual", @@ -512,40 +502,19 @@ class GPT3ComCon(discord.Cog, name="GPT3ComCon"): # Moderation if not isinstance(after.channel, discord.DMChannel): if ( - after.guild.id in self.moderation_queues - and self.moderation_queues[after.guild.id] is not None + after.guild.id in Moderation.moderation_queues + and Moderation.moderation_queues[after.guild.id] is not None ): # Create a timestamp that is 0.5 seconds from now timestamp = ( datetime.datetime.now() + datetime.timedelta(seconds=0.5) ).timestamp() - await self.moderation_queues[after.guild.id].put( + await Moderation.moderation_queues[after.guild.id].put( Moderation(after, timestamp) - ) + ) # TODO Don't proceed if message was deleted! await TextService.process_conversation_edit(self, after, original_message) - async def check_and_launch_moderations(self, guild_id, alert_channel_override=None): - # Create the moderations service. - print("Checking and attempting to launch moderations service...") - if self.check_guild_moderated(guild_id): - self.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 - ) - - self.moderation_tasks[guild_id] = asyncio.ensure_future( - Moderation.process_moderation_queue( - self.moderation_queues[guild_id], 1, 1, moderations_channel - ) - ) - print("Launched the moderations service for guild " + str(guild_id)) - self.moderations_launched.append(guild_id) - return moderations_channel - return None @discord.Cog.listener() async def on_message(self, message): @@ -554,18 +523,18 @@ class GPT3ComCon(discord.Cog, name="GPT3ComCon"): content = message.content.strip() - # Moderations service + # Moderations service is done here. if ( - message.guild.id in self.moderation_queues - and self.moderation_queues[message.guild.id] is not None + message.guild.id in Moderation.moderation_queues + and Moderation.moderation_queues[message.guild.id] is not None ): # Create a timestamp that is 0.5 seconds from now timestamp = ( datetime.datetime.now() + datetime.timedelta(seconds=0.5) ).timestamp() - await self.moderation_queues[message.guild.id].put( + await Moderation.moderation_queues[message.guild.id].put( Moderation(message, timestamp) - ) + ) # TODO Don't proceed to conversation processing if the message is deleted by moderations. # Process the message if the user is in a conversation @@ -646,7 +615,7 @@ class GPT3ComCon(discord.Cog, name="GPT3ComCon"): inline=False, ) embed.add_field( - name="/system moderations", + name="/mod", value="The automatic moderations service", inline=False, ) @@ -939,49 +908,6 @@ class GPT3ComCon(discord.Cog, name="GPT3ComCon"): if thread.id in self.awaiting_thread_responses: self.awaiting_thread_responses.remove(thread.id) - async def moderations_test_command(self, ctx: discord.ApplicationContext, prompt: str): - 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"]) - - async def moderations_command( - self, ctx: discord.ApplicationContext, status: str, alert_channel_id: str - ): - await ctx.defer() - - status = status.lower().strip() - if status not in ["on", "off"]: - await ctx.respond("Invalid status, please use on or off") - return - - 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. - self.set_guild_moderated(ctx.guild_id) - moderations_channel = await self.check_and_launch_moderations( - ctx.guild_id, - self.moderation_alerts_channel - if not alert_channel_id - else alert_channel_id, - ) - self.set_moderated_alert_channel(ctx.guild_id, moderations_channel.id) - - await ctx.respond("Moderations service enabled") - - elif status == "off": - # Cancel the moderations service. - self.set_guild_moderated(ctx.guild_id, False) - self.moderation_tasks[ctx.guild_id].cancel() - self.moderation_tasks[ctx.guild_id] = None - self.moderation_queues[ctx.guild_id] = None - self.moderations_launched.remove(ctx.guild_id) - await ctx.respond("Moderations service disabled") - async def end_command(self, ctx: discord.ApplicationContext): await ctx.defer(ephemeral=True) @@ -1039,23 +965,3 @@ class GPT3ComCon(discord.Cog, name="GPT3ComCon"): # Otherwise, process the settings change await self.process_settings(ctx, parameter, value) - def check_guild_moderated(self, guild_id): - return guild_id in MOD_DB and MOD_DB[guild_id]["moderated"] - - def get_moderated_alert_channel(self, guild_id): - return MOD_DB[guild_id]["alert_channel"] - - def set_moderated_alert_channel(self, guild_id, channel_id): - MOD_DB[guild_id] = {"moderated": True, "alert_channel": channel_id} - MOD_DB.commit() - - def set_guild_moderated(self, guild_id, status=True): - 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() diff --git a/gpt3discord.py b/gpt3discord.py index 08aafda..bbe9232 100644 --- a/gpt3discord.py +++ b/gpt3discord.py @@ -8,6 +8,7 @@ import pinecone from pycord.multicog import apply_multicog import os +from cogs.moderations_service_cog import ModerationsService from services.pinecone_service import PineconeService if sys.platform == "win32": @@ -100,6 +101,9 @@ async def main(): if not data_path.exists(): raise OSError(f"Data path: {data_path} does not exist ... create it?") + # Load the cog for the moderations service + bot.add_cog(ModerationsService(bot, usage_service, model)) + # Load the main GPT3 Bot service bot.add_cog( GPT3ComCon( @@ -147,10 +151,12 @@ async def main(): deletion_queue, bot.get_cog("GPT3ComCon"), bot.get_cog("DrawDallEService"), - bot.get_cog("ImgPromptOptimizer") + bot.get_cog("ImgPromptOptimizer"), + bot.get_cog("ModerationsService"), ) ) + apply_multicog(bot) await bot.start(os.getenv("DISCORD_TOKEN")) diff --git a/services/moderations_service.py b/services/moderations_service.py index cb69490..8420652 100644 --- a/services/moderations_service.py +++ b/services/moderations_service.py @@ -7,6 +7,7 @@ from pathlib import Path import discord from models.openai_model import Model +from services.environment_service import EnvService from services.usage_service import UsageService usage_service = UsageService(Path(os.environ.get("DATA_DIR", os.getcwd()))) @@ -51,6 +52,13 @@ class ThresholdSet: class Moderation: + # Moderation service data + moderation_queues = {} + moderation_alerts_channel = EnvService.get_moderations_alert_channel() + moderation_enabled_guilds = [] + moderation_tasks = {} + moderations_launched = [] + def __init__(self, message, timestamp): self.message = message self.timestamp = timestamp