From b05034cb07c806c25127fc36e141ef62bf304bf8 Mon Sep 17 00:00:00 2001 From: brandons209 Date: Sat, 9 Mar 2019 18:22:17 -0500 Subject: [PATCH] overhauled structure, saving data, and integrated economy cog --- scriptcog/script.py | 300 +++++++++++++++++++++++++++++++++++--------- 1 file changed, 242 insertions(+), 58 deletions(-) diff --git a/scriptcog/script.py b/scriptcog/script.py index 370412b..2847ab9 100644 --- a/scriptcog/script.py +++ b/scriptcog/script.py @@ -12,11 +12,13 @@ import pickle import time import asyncio +from .utils.dataIO import dataIO +from .utils.chat_formatting import warning #loads dictionary from file def _load_dict(path): with open(path, 'rb') as file: - dict = pickle.load(file) + dict = pickle.load(file) return dict #dictionaries for tokenizing puncuation and converting it back @@ -59,15 +61,35 @@ def _sample(prediction, temp=0): probabilities = np.random.multinomial(1, prediction, 1) return np.argmax(probabilities) +def _resolve_role_list(server: discord.Server, roles: list) -> list: + gen = (_role_from_string(server, name) for name in roles) + return list(filter(None, gen)) + +def _role_from_string(server, rolename, roles=None): + if roles is None: + roles = server.roles + + roles = [r for r in roles if r is not None] + role = discord.utils.find(lambda r: r.name.lower() == rolename.lower(), roles) + # if couldnt find by role name, try to find by role id + if role is None: + role = discord.utils.find(lambda r: r.id == rolename, roles) + + return role + +def format_list(*items, join='and', delim=', '): + if len(items) > 1: + return (' %s ' % join).join((delim.join(items[:-1]), items[-1])) + elif items: + return items[0] + else: + return '' """This cog generates scripts based on imported model, I used a keras model. """ class ScriptCog: def __init__(self, bot): self.bot = bot - - os.makedirs("data/scriptcog/", exist_ok=True) - os.makedirs("data/scriptcog/dicts", exist_ok=True) self.model_path = "data/scriptcog/model.h5" self.dict_path = "data/scriptcog/dicts/" @@ -78,19 +100,13 @@ class ScriptCog: except: self.model = None - if os.path.isfile("data/scriptcog/config.txt"): - with open("data/scriptcog/config.txt", "r") as f: - content = f.read() - content = content.split() - self.word_limit = int(content[0]) - self.cooldown_limit = int(content[1]) - self.tv_show = content[2] - else: - with open("data/scriptcog/config.txt", "w") as f: - f.write("100 30 MLP") - self.word_limit = 100 - self.cooldown_limit = 30 - self.tv_show = "MLP" + self.settings_path = "data/scriptcog/settings.json" + self.settings = dataIO.load_json(self.settings_path) + + self.default_word_limit = 300 + self.default_cooldown_limit = 30 + self.default_tv_show = "My Little Pony" + self.default_price = 0 try: self.word_to_int = _load_dict(self.dict_path + 'word_to_int.pkl') @@ -101,41 +117,178 @@ class ScriptCog: self.int_to_word = None self.sequence_length = None - def _write_config(self): - with open("data/scriptcog/config.txt", "w") as f: - f.write("{} {} {}".format(self.word_limit, self.cooldown_limit, self.tv_show)) + @commands.group(pass_context=True, invoke_without_command=True, no_pm=True) + @checks.admin_or_permissions(administrator=True) + async def genscriptset(self, ctx): + if ctx.invoked_subcommand is None: + await self.bot.send_cmd_help(ctx) - @commands.command(pass_context=True, no_pm=True) + @genscriptset.command(pass_context=True, no_pm=True, name="word-limit") @checks.is_owner() - async def setwordlimit(self, ctx, num_words : int = 100): - #if ctx.invoked_subcommand is None: - # await self.bot.say("Usage: setwordlimit limit") - # return - self.word_limit = num_words - self._write_config() - await self.bot.say("Maximum number of words is now {}".format(self.word_limit)) + async def genscriptset_set_word_limit(self, ctx, num_words : int): + """ + Set the word limit for generating scripts. + """ + try: + self.settings[ctx.message.server.id]["WORD_LIMIT"] = num_words + except: + self.settings[ctx.message.server.id] = {} + self.settings[ctx.message.server.id]["WORD_LIMIT"] = num_words + + dataIO.save_json(self.settings_path, self.settings) + await self.bot.say("Maximum number of words is now {}".format(num_words)) + + @genscriptset.command(pass_context=True, no_pm=True, name="tv-show") + @checks.is_owner() + async def genscriptset_set_tv_show(self, ctx, *, show): + """ + Sets the TV show the scripts are generated from. + """ + try: + self.settings[ctx.message.server.id]["TV_SHOW"] = show + except: + self.settings[ctx.message.server.id] = {} + self.settings[ctx.message.server.id]["TV_SHOW"] = show + + dataIO.save_json(self.settings_path, self.settings) + await self.bot.say("TV show is now {}.".format(show)) + + @genscriptset.command(pass_context=True, no_pm=True, name="cooldown") + @checks.is_owner() + async def genscriptset_set_cooldown(self, ctx, cooldown : int): + """ + Sets the cooldown period between generating scripts in seconds. + """ + try: + self.settings[ctx.message.server.id]["COOLDOWN"] = cooldown + except: + self.settings[ctx.message.server.id] = {} + self.settings[ctx.message.server.id]["COOLDOWN"] = cooldown + + dataIO.save_json(self.settings_path, self.settings) + await self.bot.say("Script cooldown is now {}.".format(cooldown)) + + @genscriptset.command(pass_context=True, no_pm=True, name="price") + async def genscriptset_set_price(self, ctx, price : int): + """ + Sets the price for generating scripts. + """ + try: + self.settings[ctx.message.server.id]["PRICE"] = price + except: + self.settings[ctx.message.server.id] = {} + self.settings[ctx.message.server.id]["PRICE"] = price + + dataIO.save_json(self.settings_path, self.settings) + await self.bot.say("Price is now {}.".format(price)) + + @genscriptset.command(pass_context=True, no_pm=True, name="free-roles") + async def genscriptset_free_roles(self, ctx, *, rolelist=None): + """Set roles that do not have to pay to generate scripts. + + COMMA SEPARATED LIST (e.g. Admin,Staff,Mod), Can also use role IDs as well. + + To get current list, run command with no roles. + + Add role_list_clear as the role to clear the server\'s free role list. + """ + server = ctx.message.server + current_roles = _resolve_role_list(server, self.settings[server.id].get("FREE_ROLE_LIST", [])) + + if rolelist is None: + if current_roles: + names_list = format_list(*(r.name for r in current_roles)) + await self.bot.say("Current list of roles that do not have to pay: {}".format(names_list)) + else: + await self.bot.say("No roles defined.") + return + elif "role_list_clear" in rolelist.lower(): + await self.bot.say("Free role list cleared.") + self.settings[server.id]["FREE_ROLE_LIST"] = [] + dataIO.save_json(self.settings_path, self.settings) + return + + found_roles = set() + notfound_names = set() + + for lookup in rolelist.split(","): + lookup = lookup.strip() + role = _role_from_string(server, lookup) + + if role: + found_roles.add(role) + else: + notfound_names.add(lookup) + + if notfound_names: + fmt_list = format_list(*("`{}`".format(x) for x in notfound_names)) + await self.bot.say(warning("These roles were not found: {}\n\nPlease try again.".format(fmt_list))) + elif server.default_role in found_roles: + await self.bot.say(warning("The everyone role cannot be added.\n\nPlease try again.")) + elif found_roles == set(current_roles): + await self.bot.say("No changes to make.") + else: + if server.id not in self.settings: + self.settings[server.id] = {} + else: + extra = "" + + self.settings[server.id]["FREE_ROLE_LIST"] = [r.id for r in found_roles] + dataIO.save_json(self.settings_path, self.settings) + + fmt_list = format_list(*(r.name for r in found_roles)) + await self.bot.say("These roles will not have to pay for scripts: {}.{}".format(fmt_list, extra)) @commands.command(pass_context=True, no_pm=True) async def genscriptinfo(self, ctx): - await self.bot.say("Word Limit: {}, Cooldown Time: {}, Show: {}".format(self.word_limit, self.cooldown_limit, self.tv_show)) + server_id = ctx.message.server.id + word_limit = self.settings.get(server_id, {}).get("WORD_LIMIT", self.default_word_limit) + cooldown = self.settings.get(server_id, {}).get("COOLDOWN", self.default_cooldown_limit) + tv_show = self.settings.get(server_id, {}).get("TV_SHOW", self.default_tv_show) + price = self.settings.get(server_id, {}).get("PRICE", self.default_price) + await self.bot.say("Word Limit: {}, Cooldown Time: {}, Show: {}, Price: {}".format(word_limit, cooldown, tv_show, price)) @commands.command(pass_context=True, no_pm=True) async def genscripthelp(self, ctx): - await self.bot.say("--------------------\nGenerate original TV scripts for {} using Neural Networks!\nUsage: `genscript number_of_words_to_generate word_variance starting_text`\nUse starting texts such as:\n`\"pinkie pie::\"`\n`fluttershy::`\n`\"twilight sparkle::\"`\nor other names of characters in the show. Otherwise, you can use any words said in the show.\n\n**If your starting text has spaces, enclose it with quotes: \"\"**\nWord variance helps gives the script better results. A variance of 0 will mean that with the same starting text, it will always have the same output. Variance up to 1.0 will give more variety to words, however going closer to 1 can introduce more grammar and spelling mistakes.\n-------------------".format(self.tv_show)) + await self.bot.say("--------------------\nGenerate original TV scripts for {} using Neural Networks!\nUsage: `genscript `\nUse starting texts such as:\n`pinkie pie::`\n`fluttershy::`\n`twilight sparkle::`\nor other names of characters in the show. Otherwise, you can use any words said in the show.\n\nWord variance helps gives the script better results. A variance of 0 will mean that with the same starting text, it will always have the same output. Variance up to 1.0 will give more variety to words, however going closer to 1 can introduce more grammar and spelling mistakes.\n-------------------".format(self.settings.get(ctx.message.server.id, {}).get("TV_SHOW", self.default_tv_show))) @commands.command(pass_context=True, no_pm=True) - @checks.is_owner() - async def settvshow(self, ctx, show : str = "MLP"): - self.tv_show = show - self._write_config() - await self.bot.say("TV show is now {}.".format(self.tv_show)) + async def genscript(self, ctx, num_words_to_generate : int, variance : float, *, seed): + """ + Generate a script using the power of Neural Networks! + Please use the genscripthelp command to get a complete explaination of how to use the command. + """ + server_id = ctx.message.server.id + word_limit = self.settings.get(server_id, {}).get("WORD_LIMIT", self.default_word_limit) + cooldown = self.settings.get(server_id, {}).get("COOLDOWN", self.default_cooldown_limit) + price = self.settings.get(server_id, {}).get("PRICE", self.default_price) + user = ctx.message.author - @commands.command(pass_context=True, no_pm=True) - @checks.is_owner() - async def setcooldown(self, ctx, cooldown : int = 30): - self.cooldown_limit = cooldown - self._write_config() - await self.bot.say("Script cooldown is now {}.".format(self.cooldown_limit)) + if num_words_to_generate > word_limit: + await self.bot.say("Please keep script sizes to {} words or less.".format(word_limit)) + return + elif time.time() - self.cooldown < cooldown: + await self.bot.say("Sorry, I am cooling down, please wait {:.0f} seconds.".format(cooldown - (time.time() - self.cooldown))) + return + + if not self.check_free_role(user): + return_val = self.charge_user(user, price) + if return_val == 1: #if this worked, dont need to check when getting econ cog + balance = self.bot.get_cog('Economy').bank.get_balance(user) + await self.bot.say("Charged: {}, Balance: {}".format(price, balance)) + else: + await self.charge_user_check(return_val, price, user) + return + + self.cooldown = time.time() + + if variance > 1.0: + variance = 1.0 + elif variance < 0: + variance = 0 + + await self.bot.say("Generating script, please wait...") + await self.get_model_output(num_words_to_generate, variance, seed) async def get_model_output(self, num_words, temp, seed): input_text = seed @@ -162,27 +315,58 @@ class ScriptCog: await self.bot.say(result) await self.bot.say("------------------------") - @commands.command(pass_context=True, no_pm=True) - async def genscript(self, ctx, num_words_to_generate : int = 100, variance : float = 0.5, seed : str = "pinkie pie::"): - #if ctx.invoked_subcommand is None: - # await self.bot.say("Usage: genscript num_words randomness(between 0 and 1) seed_text") - # return - if num_words_to_generate > self.word_limit: - await self.bot.say("Please keep script sizes to {} words or less.".format(self.word_limit)) - return - elif time.time() - self.cooldown < self.cooldown_limit: - await self.bot.say("Sorry, I am cooling down, please wait {:.0f} seconds.".format(self.cooldown_limit - (time.time() - self.cooldown))) + def charge_user(self, user, amount): + """ + Takes a user and a amount and charges the user. Returns 1 on success, 0 if economy cog cannot be loaded, -1 if the user does not have a bank account, and -2 if the user doesn't have enough credits. + """ + econ_cog = self.bot.get_cog('Economy') + if not econ_cog: + return 0 + if not econ_cog.bank.account_exists(user): + return -1 + if not econ_cog.bank.can_spend(user, amount): + return -2 + econ_cog.bank.withdraw_credits(user, amount) + return 1 + + async def charge_user_check(self, return_val, amount, user): + """ + takes in return vals from charge_user and gives apporiate response + """ + econ_cog = self.bot.get_cog('Economy') + if not econ_cog: + await self.bot.say("Error loading economy cog!") return + if return_val == 0: + await self.bot.say("Economy cog not found! Please check to make sure the economy cog is loaded.") + elif return_val == -1: + await self.bot.say("You appear to not have a bank account. Use [p]bank register to open an account.") + elif return_val == -2: + await self.bot.say("You do not have enough credits. The sound command costs {} and you have {} credits in your bank account.".format(amount, econ_cog.bank.get_balance(user))) - self.cooldown = time.time() + def check_free_role(self, user): + """ + Checks if the user has a role that grants them free sounds. Returns 1 if the user does have one of these roles, zero otherwise. + """ + server = user.server + free_roles = self.settings.get(server.id, {}).get("FREE_ROLE_LIST", []) + if not free_roles: + return 0 + user_roles = [r.id for r in user.roles] - if variance > 1.0: - variance = 1.0 - elif variance < 0: - variance = 0 + for role in free_roles: + if role in user_roles: + return 1 - await self.bot.say("Generating script, please wait...") - await self.get_model_output(num_words_to_generate, variance, seed) + return 0 + +def check_folders(): + os.makedirs("data/scriptcog", exist_ok=True) + os.makedirs("data/scriptcog/dicts", exist_ok=True) + f = "data/scriptcog/settings.json" + if not dataIO.is_valid_json(f): + dataIO.save_json(f, {}) def setup(bot): + check_folders() bot.add_cog(ScriptCog(bot))