diff --git a/src/ChatProcess.py b/src/ChatProcess.py index 594c0e6..42e0235 100644 --- a/src/ChatProcess.py +++ b/src/ChatProcess.py @@ -68,7 +68,9 @@ class Chat: This checks if the bot should actuallly respond to the message or if the message doesn't concern the bot """ returnCriterias = [] - returnCriterias.append(self.guild.sanitizedChannels.get(str(self.message.channel.id), None) != None) + returnCriterias.append( + self.guild.sanitizedChannels.get(str(self.message.channel.id), None) != None + ) returnCriterias.append( self.message.content.find("<@" + str(self.bot.user.id) + ">") != -1 ) @@ -82,5 +84,3 @@ class Chat: self.model = self.settings["model"] self.character = self.settings["character"] self.openai_api_key = self.guild.api_keys.get("openai", None) - - \ No newline at end of file diff --git a/src/cogs/__init__.py b/src/cogs/__init__.py index b84dc2b..70c6c24 100644 --- a/src/cogs/__init__.py +++ b/src/cogs/__init__.py @@ -4,4 +4,4 @@ from src.cogs.help import Help from src.cogs.chat import Chat from src.cogs.manage_chat import ManageChat from src.cogs.moderation import Moderation -from src.cogs.channelSetup import ChannelSetup \ No newline at end of file +from src.cogs.channelSetup import ChannelSetup diff --git a/src/cogs/channelSetup.py b/src/cogs/channelSetup.py index 8c7663e..0c83421 100644 --- a/src/cogs/channelSetup.py +++ b/src/cogs/channelSetup.py @@ -14,41 +14,101 @@ sampleDataFormatExample = { "premium_expiration": 0, } + class ChannelSetup(commands.Cog): def __init__(self, bot: discord.Bot): super().__init__() self.bot = bot - setup = SlashCommandGroup("setup", description="Setup commands for the bot, inlcuding channels, models, and more.") + setup = SlashCommandGroup( + "setup", + description="Setup commands for the bot, inlcuding channels, models, and more.", + ) - setup_channel = setup.create_subgroup(name="channel", description="Setup, add, or remove channels for the bot to use.") + setup_channel = setup.create_subgroup( + name="channel", description="Setup, add, or remove channels for the bot to use." + ) - @setup_channel.command(name="add", description="Add a channel for the bot to use. Can also specify server-wide settings.") - @discord.option(name="channel", description="The channel to setup. If not specified, will use the current channel.", type=discord.TextChannel, required=False) - @discord.option(name="model", description="The model to use for this channel.", type=str, required=False, autocomplete=models.autocomplete) - @discord.option(name="character", description="The character to use for this channel.", type=str, required=False, autocomplete=characters.autocomplete) + @setup_channel.command( + name="add", + description="Add a channel for the bot to use. Can also specify server-wide settings.", + ) + @discord.option( + name="channel", + description="The channel to setup. If not specified, will use the current channel.", + type=discord.TextChannel, + required=False, + ) + @discord.option( + name="model", + description="The model to use for this channel.", + type=str, + required=False, + autocomplete=models.autocomplete, + ) + @discord.option( + name="character", + description="The character to use for this channel.", + type=str, + required=False, + autocomplete=characters.autocomplete, + ) @guild_only() - async def channel(self, ctx: discord.ApplicationContext, channel: discord.TextChannel = None, model: str = models.default, character: str = characters.default): + async def channel( + self, + ctx: discord.ApplicationContext, + channel: discord.TextChannel = None, + model: str = models.default, + character: str = characters.default, + ): if channel is None: channel = ctx.channel guild = Guild(ctx.guild.id) guild.load() if not guild.premium: - if len(guild.channels) >= 1 and guild.channels.get(str(channel.id), None) is None: - await ctx.respond("`Warning: You are not a premium user, and can only have one channel setup. The settings will still be saved, but will not be used.`", ephemeral=True) + if ( + len(guild.channels) >= 1 + and guild.channels.get(str(channel.id), None) is None + ): + await ctx.respond( + "`Warning: You are not a premium user, and can only have one channel setup. The settings will still be saved, but will not be used.`", + ephemeral=True, + ) if model != models.default: - await ctx.respond("`Warning: You are not a premium user, and can only use the default model. The settings will still be saved, but will not be used.`", ephemeral=True) + await ctx.respond( + "`Warning: You are not a premium user, and can only use the default model. The settings will still be saved, but will not be used.`", + ephemeral=True, + ) if character != characters.default: - await ctx.respond("`Warning: You are not a premium user, and can only use the default character. The settings will still be saved, but will not be used.`", ephemeral=True) + await ctx.respond( + "`Warning: You are not a premium user, and can only use the default character. The settings will still be saved, but will not be used.`", + ephemeral=True, + ) if guild.api_keys.get("openai", None) is None: - await ctx.respond("`Error: No openai api key is set. The api key is needed for the openai models, as well as for the content moderation. The openai models will cost you tokens in your openai account. However, if you use one of the llama models, you will not be charged, but the api key is still needed for content moderation, wich is free but requires an api key.`", ephemeral=True) - guild.addChannel(channel, models.matchingDict[model], characters.matchingDict[character]) - await ctx.respond(f"Set channel {channel.mention} with model `{model}` and character `{character}`.") + await ctx.respond( + "`Error: No openai api key is set. The api key is needed for the openai models, as well as for the content moderation. The openai models will cost you tokens in your openai account. However, if you use one of the llama models, you will not be charged, but the api key is still needed for content moderation, wich is free but requires an api key.`", + ephemeral=True, + ) + guild.addChannel( + channel, models.matchingDict[model], characters.matchingDict[character] + ) + await ctx.respond( + f"Set channel {channel.mention} with model `{model}` and character `{character}`." + ) - @setup_channel.command(name="remove", description="Remove a channel from the bot's usage.") - @discord.option(name="channel", description="The channel to remove. If not specified, will use the current channel.", type=discord.TextChannel, required=False) + @setup_channel.command( + name="remove", description="Remove a channel from the bot's usage." + ) + @discord.option( + name="channel", + description="The channel to remove. If not specified, will use the current channel.", + type=discord.TextChannel, + required=False, + ) @guild_only() - async def remove(self, ctx: discord.ApplicationContext, channel: discord.TextChannel = None): + async def remove( + self, ctx: discord.ApplicationContext, channel: discord.TextChannel = None + ): if channel is None: channel = ctx.channel guild = Guild(ctx.guild.id) @@ -67,17 +127,31 @@ class ChannelSetup(commands.Cog): if len(guild.channels) == 0: await ctx.respond("No channels are setup.") return - embed = discord.Embed(title="Channels", description="All channels that are setup.", color=discord.Color.nitro_pink()) + embed = discord.Embed( + title="Channels", + description="All channels that are setup.", + color=discord.Color.nitro_pink(), + ) channels = guild.sanitizedChannels for channel in channels: - discochannel = await self.bot.fetch_channel(int(channel)) + discochannel = await self.bot.fetch_channel(int(channel)) model = models.reverseMatchingDict[channels[channel]["model"]] character = characters.reverseMatchingDict[channels[channel]["character"]] - embed.add_field(name=f"{discochannel.mention}", value=f"Model: `{model}`\nCharacter: `{character}`", inline=False) + embed.add_field( + name=f"{discochannel.mention}", + value=f"Model: `{model}`\nCharacter: `{character}`", + inline=False, + ) await ctx.respond(embed=embed) @setup.command(name="api", description="Set an API key for the bot to use.") - @discord.option(name="api", description="The API to set. Currently only OpenAI is supported.", type=str, required=True, autocomplete=apis.autocomplete) + @discord.option( + name="api", + description="The API to set. Currently only OpenAI is supported.", + type=str, + required=True, + autocomplete=apis.autocomplete, + ) @discord.option(name="key", description="The key to set.", type=str, required=True) @guild_only() async def api(self, ctx: discord.ApplicationContext, api: str, key: str): @@ -92,6 +166,9 @@ class ChannelSetup(commands.Cog): guild = Guild(ctx.guild.id) guild.load() if not guild.premium: - await ctx.respond("You can get your premium subscription at https://www.botator.dev/premium.", ephemeral=True) + await ctx.respond( + "You can get your premium subscription at https://www.botator.dev/premium.", + ephemeral=True, + ) else: - await ctx.respond("This guild is already premium.", ephemeral=True) \ No newline at end of file + await ctx.respond("This guild is already premium.", ephemeral=True) diff --git a/src/cogs/setup.py b/src/cogs/setup.py index 90fa12d..7f2c383 100644 --- a/src/cogs/setup.py +++ b/src/cogs/setup.py @@ -29,6 +29,7 @@ class Setup(discord.Cog): def __init__(self, bot: discord.Bot): super().__init__() self.bot = bot + """ @discord.slash_command(name="setup", description="Setup the bot") @discord.option(name="channel_id", description="The channel id", required=True) @@ -140,6 +141,7 @@ class Setup(discord.Cog): con_data.commit() await ctx.respond("The api key has been added", ephemeral=True) """ + @discord.slash_command( name="delete", description="Delete the information about this server" ) @@ -191,7 +193,7 @@ class Setup(discord.Cog): await ctx.respond("Disabled", ephemeral=True) # create a command calles "add channel" that can only be used in premium servers - + """ @discord.slash_command( name="setup_channel", @@ -261,6 +263,7 @@ class Setup(discord.Cog): return await ctx.respond("You can only add 5 channels", ephemeral=True) """ + # create a command called "remove channel" that can only be used in premium servers @discord.slash_command( name="remove_channel", diff --git a/src/config.py b/src/config.py index 137c295..fef50ce 100644 --- a/src/config.py +++ b/src/config.py @@ -54,7 +54,9 @@ curs_data.execute( """CREATE TABLE IF NOT EXISTS data (guild_id text, channel_id text, api_key text, is_active boolean, max_tokens integer, temperature real, frequency_penalty real, presence_penalty real, uses_count_today integer, prompt_size integer, prompt_prefix text, tts boolean, pretend_to_be text, pretend_enabled boolean)""" ) -con_data.execute('CREATE TABLE IF NOT EXISTS setup_data (guild_id text, guild_settings text)') +con_data.execute( + "CREATE TABLE IF NOT EXISTS setup_data (guild_id text, guild_settings text)" +) # This code creates the model table if it does not exist curs_data.execute( diff --git a/src/guild.py b/src/guild.py index dce5da0..c60db50 100644 --- a/src/guild.py +++ b/src/guild.py @@ -5,6 +5,7 @@ from src.utils.SqlConnector import sql from datetime import datetime from src.utils.variousclasses import models, characters + class Guild: def __init__(self, id: int): self.id = id @@ -20,14 +21,18 @@ class Guild: self.updateDbData() with sql.mainDb as con: curs_data = con.cursor() - curs_data.execute("SELECT * FROM setup_data WHERE guild_id = ?", (self.id,)) + curs_data.execute( + "SELECT * FROM setup_data WHERE guild_id = ?", (self.id,) + ) data = curs_data.fetchone() data = orjson.loads(data[1]) self.premium = data["premium"] self.channels = data["channels"] self.api_keys = data["api_keys"] if self.premium: - self.premium_expiration = datetime.fromisoformat(data.get("premium_expiration", None)) + self.premium_expiration = datetime.fromisoformat( + data.get("premium_expiration", None) + ) self.checkPremiumExpires() else: self.premium_expiration = None @@ -48,7 +53,9 @@ class Guild: "premium": self.premium, "channels": self.channels, "api_keys": self.api_keys, - "premium_expiration": self.premium_expiration.isoformat() if self.premium else None, + "premium_expiration": self.premium_expiration.isoformat() + if self.premium + else None, } else: data = { @@ -61,22 +68,30 @@ class Guild: with sql.mainDb as con: curs_data = con.cursor() if self.isInDb: - curs_data.execute("UPDATE setup_data SET guild_settings = ? WHERE guild_id = ?", (orjson.dumps(data), self.id)) + curs_data.execute( + "UPDATE setup_data SET guild_settings = ? WHERE guild_id = ?", + (orjson.dumps(data), self.id), + ) else: - curs_data.execute("INSERT INTO setup_data (guild_id, guild_settings) VALUES (?, ?)", (self.id, orjson.dumps(data))) + curs_data.execute( + "INSERT INTO setup_data (guild_id, guild_settings) VALUES (?, ?)", + (self.id, orjson.dumps(data)), + ) self.isInDb = True - + def load(self): self.getDbData() - + def addChannel(self, channel: discord.TextChannel, model: str, character: str): - print(f"Adding channel {channel.id} to guild {self.id} with model {model} and character {character}") + print( + f"Adding channel {channel.id} to guild {self.id} with model {model} and character {character}" + ) self.channels[str(channel.id)] = { "model": model, "character": character, } self.updateDbData() - + def delChannel(self, channel: discord.TextChannel): del self.channels[str(channel.id)] self.updateDbData() @@ -99,4 +114,4 @@ class Guild: def addApiKey(self, api: str, key: str): self.api_keys[api] = key - self.updateDbData() \ No newline at end of file + self.updateDbData() diff --git a/src/utils/SqlConnector.py b/src/utils/SqlConnector.py index 68e56a1..d9f2f31 100644 --- a/src/utils/SqlConnector.py +++ b/src/utils/SqlConnector.py @@ -3,24 +3,22 @@ from random import randint class SQLConnection: - - def __init__(self,connection): + def __init__(self, connection): self.connection = connection def __enter__(self): return self.connection - def __exit__(self,exc_type,exc_val,exc_tb): + def __exit__(self, exc_type, exc_val, exc_tb): self.connection.commit() self.connection.close() class _sql: - @property def mainDb(self): - s = connect('./database/data.db') + s = connect("./database/data.db") return SQLConnection(s) -sql: _sql = _sql() \ No newline at end of file +sql: _sql = _sql() diff --git a/src/utils/variousclasses.py b/src/utils/variousclasses.py index 087c240..1647975 100644 --- a/src/utils/variousclasses.py +++ b/src/utils/variousclasses.py @@ -1,5 +1,6 @@ from discord import AutocompleteContext + class models: matchingDict = { "chatGPT (default - free)": "gpt-3.5-turbo", @@ -10,11 +11,13 @@ class models: reverseMatchingDict = {v: k for k, v in matchingDict.items()} default = list(matchingDict.keys())[0] openaimodels = ["gpt-3.5-turbo", "text-davinci-003"] + @classmethod async def autocomplete(cls, ctx: AutocompleteContext) -> list[str]: modls = cls.matchingDict.keys() return [model for model in modls if model.find(ctx.value.lower()) != -1] + class characters: matchingDict = { "Botator (default - free)": "botator", @@ -22,16 +25,21 @@ class characters: } reverseMatchingDict = {v: k for k, v in matchingDict.items()} default = list(matchingDict.keys())[0] + @classmethod async def autocomplete(cls, ctx: AutocompleteContext) -> list[str]: chars = characters = cls.matchingDict.keys() - return [character for character in chars if character.find(ctx.value.lower()) != -1] + return [ + character for character in chars if character.find(ctx.value.lower()) != -1 + ] + class apis: matchingDict = { "OpenAI": "openai", } + @classmethod async def autocomplete(cls, ctx: AutocompleteContext) -> list[str]: apiss = cls.matchingDict.keys() - return [api for api in apiss if api.find(ctx.value.lower()) != -1] \ No newline at end of file + return [api for api in apiss if api.find(ctx.value.lower()) != -1]