From aabc938434f61115a63213217b8b3c077bb240f3 Mon Sep 17 00:00:00 2001 From: brandons209 Date: Wed, 22 Jul 2020 02:09:05 -0400 Subject: [PATCH] add markov cog --- markov/__init__.py | 5 ++ markov/info.json | 22 +++++++ markov/markov.py | 150 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 177 insertions(+) create mode 100644 markov/__init__.py create mode 100644 markov/info.json create mode 100644 markov/markov.py diff --git a/markov/__init__.py b/markov/__init__.py new file mode 100644 index 0000000..b29f415 --- /dev/null +++ b/markov/__init__.py @@ -0,0 +1,5 @@ +from .markov import Markov + + +def setup(bot): + bot.add_cog(Markov(bot)) diff --git a/markov/info.json b/markov/info.json new file mode 100644 index 0000000..2593acd --- /dev/null +++ b/markov/info.json @@ -0,0 +1,22 @@ +{ + "author": [ + "Brandons209" + ], + "bot_version": [ + 3, + 0, + 0 + ], + "description": "Uses markov chains to generate unique text in your guild! Text generated is based on what users say per channel. The markov model can be reset if needed as well per channel.", + "hidden": false, + "install_msg": "Thank you for using this cog! Make sure to set bot prefixes so the bot won't train on commands from users.", + "requirements": [], + "short": "Markov chain text generation", + "tags": [ + "brandons209", + "redbot", + "markov", + "generaton", + "markov chain" + ] +} diff --git a/markov/markov.py b/markov/markov.py new file mode 100644 index 0000000..c102ddb --- /dev/null +++ b/markov/markov.py @@ -0,0 +1,150 @@ +from redbot.core.utils.chat_formatting import * +from redbot.core import Config, checks, commands +import discord +import random + + +class Markov(commands.Cog): + """ Generate text based on what your members say per channel""" + + def __init__(self, bot): + self.bot = bot + self.config = Config.get_conf(self, identifier=5989735216541313, force_registration=True) + + default_guild = {"model": {}, "prefixes": [], "max_len": 200} + self.config.register_guild(**default_guild) + + @commands.group() + @checks.admin_or_permissions(administrator=True) + @commands.guild_only() + async def markovset(self, ctx): + """ Manage Markov Settings """ + pass + + @markovset.command(name="clear") + async def markovset_clear(self, ctx, *, channel: discord.TextChannel): + """ Clear data for a specific channel """ + async with self.config.guild(ctx.guild).model() as model: + del model[str(channel.id)] + + await ctx.tick() + + @markovset.command(name="prefix") + async def markovset_prefixes(self, ctx, *, prefixes: str = None): + """ Set prefixes for bots in your server + This is so markov won't log bot commands. + """ + if not prefixes: + current = await self.config.guild(ctx.guild).prefixes() + curr = [f"`{p}`" for p in current] + if not current: + await ctx.send("No prefixes set, setting this bot's prefix.") + await self.config.guild(ctx.guild).prefixes.set([ctx.clean_prefix]) + return + + await ctx.send("Current Prefixes: " + humanize_list(curr)) + return + + prefixes = [p for p in prefixes.split(" ")] + await self.config.guild(ctx.guild).prefixes.set(prefixes) + prefixes = [f"`{p}`" for p in prefixes] + await ctx.send("Prefixes set to: " + humanize_list(prefixes)) + + @markovset.command(name="len") + async def markovset_length(self, ctx, length: int = None): + """ + Set max characters of generated text. + + Max size limited by discord is 3000. + """ + if not length: + curr = await self.config.guild(ctx.guild).max_len() + await ctx.send(f"Current max length of generated text is `{curr}` characters.") + return + + await self.config.guild(ctx.guild).max_len.set(length) + await ctx.tick() + + @commands.command(name="markov") + @commands.guild_only() + @commands.cooldown(rate=1, per=10, type=commands.BucketType.user) + @checks.bot_has_permissions(embed_links=True) + async def markov(self, ctx, *, starting_text: str = None): + """ Generate text using markov chains! + + Text generated is based on what users say in the current channel + """ + model = await self.config.guild(ctx.guild).model() + try: + model = model[str(ctx.channel.id)] + except KeyError: + await ctx.send(error("This channel has no data, try talking in it for a bit first!")) + return + + last_word = starting_text.split(" ")[-1] if starting_text else None + + if not starting_text: + markov_text = [random.choice(list(model.keys()))] + elif not model[last_word]: + markov_text = [last_word, random.choice(list(model.keys()))] + else: + markov_text = [last_word, random.choice(model[last_word])] + + max_len = await self.config.guild(ctx.guild).max_len() + + tries = 0 + max_tries = 20 + while len(markov_text) < max_len and tries < max_tries: + if "?" in markov_text[-1]: + break + if "\r" in markov_text[-1]: + break + if "." in markov_text[-1]: + break + if "!" in markov_text[-1]: + break + + # make sure word is in the model and there is data for the word + if markov_text[-1] in model and model[markov_text[-1]]: + markov_text.append(random.choice(model[markov_text[-1]])) + else: + markov_text.append(random.choice(list(model.keys()))) + tries += 1 + + markov_text = " ".join(markov_text) + member = ctx.author + embed = discord.Embed(title="Generated Text", description=markov_text, colour=member.colour) + + if member.avatar: + avatar = member.avatar_url_as(static_format="png") + embed.set_thumbnail(url=avatar) + + embed.set_footer(text=f"Generated by {member.display_name}") + await ctx.send(embed=embed) + + # Listener + @commands.Cog.listener() + async def on_message(self, message): + # updates model + content = message.content + if not content or message.author == message.guild.me: + return + + # check if this is a bot message + prefixes = await self.config.guild(message.guild).prefixes() + for prefix in prefixes: + if prefix == content[: len(prefix)]: + return + + async with self.config.guild(message.guild).model() as model: + content = content.split(" ") + try: + model[str(message.channel.id)] + except: + model[str(message.channel.id)] = {} + + for i in range(len(content) - 1): + if content[i] not in model[str(message.channel.id)]: + model[str(message.channel.id)][content[i]] = list() + + model[str(message.channel.id)][content[i]].append(content[i + 1])