From a38251f971464dbc848c895cc59bbfb4a1665c39 Mon Sep 17 00:00:00 2001 From: LBlend Date: Tue, 18 Apr 2023 21:19:15 +0900 Subject: [PATCH] add wordcloud cog --- requirements.txt | 3 + src/cogs/word_cloud.py | 475 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 478 insertions(+) create mode 100644 src/cogs/word_cloud.py diff --git a/requirements.txt b/requirements.txt index e768269..a92fdfa 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,8 +2,11 @@ aiocron==1.8.* asyncio==3.4.* discord.py==2.1.* mcrcon==0.7.* +nltk==3.6.* +numpy==1.22.* pillow==9.3.* psutil==5.9.* psycopg2==2.9.* pyyaml==6.0.* requests==2.28.* +wordcloud==1.8.* diff --git a/src/cogs/word_cloud.py b/src/cogs/word_cloud.py new file mode 100644 index 0000000..a571f9d --- /dev/null +++ b/src/cogs/word_cloud.py @@ -0,0 +1,475 @@ +import discord +from discord import app_commands +from discord.ext import commands, tasks + +from collections import defaultdict +import functools +import itertools +from io import BytesIO, StringIO +import json +import nltk +from nltk.corpus import stopwords +import numpy as np +from PIL import Image +from psycopg2.extras import execute_values +import re +from wordcloud import WordCloud as WCloud # Avoid naming conflicts with cog class name +# from wordcloud import ImageColorGenerator + +from cogs.utils import embed_templates + + +class WordCloud(commands.Cog): + """Generate a wordcloud based on the most frequent words posted""" + + def __init__(self, bot: commands.Bot): + """ + Parameters + ---------- + bot (commands.Bot): The bot instance + """ + + self.bot = bot + self.cursor = self.bot.db_connection.cursor() + self.init_db() + + # Default dict where all users are empty defaultdicts + # The user's defaultdict has a default value of 0 + self.word_freq_cache = defaultdict(lambda: defaultdict(int)) + + self.consenting_users = [] + self.populate_consenting_users() + + nltk.download('stopwords') + + def init_db(self): + """ + Create the necessary tables for the Ordsky cog to work + """ + + self.cursor.execute( + """ + CREATE TABLE IF NOT EXISTS wordcloud_metadata ( + discord_user_id BIGINT PRIMARY KEY, + tracked_since_message_channel_id BIGINT NOT NULL, + tracked_since_message_id BIGINT NOT NULL UNIQUE + ) + """ + ) + self.cursor.execute( + """ + CREATE TABLE IF NOT EXISTS wordcloud_words ( + discord_user_id BIGINT NOT NULL, + word TEXT NOT NULL, + frequency INTEGER NOT NULL, + PRIMARY KEY (discord_user_id, word) + ) + """ + ) + + def populate_consenting_users(self): + """ + Populates the cached list of consenting users + """ + + self.cursor.execute( + """ + SELECT discord_user_id + FROM wordcloud_metadata + """ + ) + self.consenting_users = [row[0] for row in self.cursor.fetchall()] + + async def batch_update_word_freqs(self): + """ + Inserts cached word frequencies into the database. + We do this in order to prevent excess database writes + """ + + if not self.word_freq_cache: + return + + # Flatten the defaultdict into a list of tuples + # This is a bit of a hack but it works + word_freqs = list(itertools.chain.from_iterable( + [(user_id, word, freq) for word, freq in user_word_freqs.items()] + for user_id, user_word_freqs in self.word_freq_cache.items() + )) + + # Insert cache into database + try: + execute_values( + self.cursor, + """ + INSERT INTO wordcloud_words (discord_user_id, word, frequency) + VALUES (%s, %s, %s) + ON CONFLICT (discord_user_id, word) + DO UPDATE SET frequency = wordcloud_words.frequency + EXCLUDED.frequency + """, word_freqs + ) + except: + self.bot.db_connection.rollback() + self.bot.logger.exception("Failed to insert wordcloud cache into database") + + self.word_freq_cache.clear() + + @tasks.loop(minutes=20) + async def batch_update_word_freqs_loop(self): + """ + Clears and inserts word frequency cache into the database every 20 minutes + """ + + self.bot.logger.info("Dumping word cloud cache to database...") + await self.batch_update_word_freqs() + + async def word_freq_listener(self, message: discord.Message): + """ + Listens for all messages of consenting users + + Parameters + ---------- + message (discord.Message): Message to count + """ + + if message.author.id not in self.consenting_users: + return + + # Ignore bot commands + # This is a very naive approach but it works for our use case + if not message.clean_content[:2].isalpha(): + return + + # Divide into words + # Also very naive but luckily we only speak norwegian and english :) + words = message.clean_content.split(" ") + + for word in words: + # Filter urls + # do I really need to comment on this being naive again? + if re.match(r'https?://', word): + continue + + # Remove punctuation + word = word.lower().strip("/+-=~|$%@#*_.,;:!?()[]{}<>\"'`\n") + if not word: + continue + + # Enter into cache + # This may take a performance hit but who tf cares. We're using python anyway + self.word_freq_cache[message.author.id][word] += 1 + + @staticmethod + def generate_wordcloud( + text: str, + max_words: int = 4000, + allow_bigrams: bool = False + ) -> BytesIO: + """ + Generates a wordcloud + + Parameters + ---------- + text (str): Text to generate wordcloud from + max_words (int): Maximum number of words to include in the wordcloud. Defaults to 4000 + allow_bigrams (bool): Whether to allow bigrams in the wordcloud. Defaults to False + + Returns + ---------- + BytesIO: BytesIO object containing the wordcloud image + """ + + filter_words = set(stopwords.words('norwegian') + stopwords.words('english')) + mask = np.array(Image.open('./src/assets/word_cloud/mask.png')) + + wc = WCloud( + max_words=max_words, + mask=mask, + repeat=False, + stopwords=filter_words, + min_word_length=3, + collocations=allow_bigrams + ) + wc.process_text(text) + wc.generate(text) + + # Color the wordcloud based on the mask + # wc.recolor(color_func=ImageColorGenerator(mask)) + + img = wc.to_image() + b = BytesIO() + img.save(b, 'png') + b.seek(0) + return b + + wordcloud_group = app_commands.Group( + name='ordsky', + description='Generer en ordsky basert på dine mest frekvente sagte ord' + ) + wordcloud_generate_group = app_commands.Group( + parent=wordcloud_group, + name='generer', + description='Generer en ordsky basert på dine mest frekvente sagte ord' + ) + + @app_commands.checks.bot_has_permissions(embed_links=True) + @wordcloud_group.command(name='samtykke', description='Gi samtykke til å samle meldingsdataen din for ordskyer') + async def consent(self, interaction: discord.Interaction): + """ + Gi samtykke til å samle meldingsdataen din + + Parameters + ---------- + interaction (discord.Interaction): Slash command context object + """ + + if interaction.user.id in self.consenting_users: + return await interaction.response.send_message( + embed=embed_templates.error_warning(interaction, text='Du har allerede samtykket'), + ephemeral=False + ) + + try: + self.cursor.execute( + """ + INSERT INTO wordcloud_metadata (discord_user_id, tracked_since_message_channel_id, tracked_since_message_id) + VALUES (%s, %s, %s) + """, (interaction.user.id, interaction.channel_id, interaction.id) + ) + except: + self.bot.db_connection.rollback() + self.bot.logger.exception("Failed to insert wordcloud metadata into database") + return await interaction.response.send_message( + embed=embed_templates.error_warning(interaction, text='Klarte ikke å skrive til database'), + ephemeral=False + ) + + self.consenting_users.append(interaction.user.id) + + embed = embed_templates.success(interaction, 'Samtykke registrert!') + await interaction.response.send_message(embed=embed, ephemeral=False) + + @app_commands.checks.bot_has_permissions(embed_links=True) + @wordcloud_group.command(name='slett', description='Fjern samtykke og slett alle dine ordskydata') + async def consent_remove(self, interaction: discord.Interaction): + """ + Fjern samtykke og slett meldingsdata + + Parameters + ---------- + interaction (discord.Interaction): Slash command context object + """ + + try: + self.consenting_users.remove(interaction.user.id) + except ValueError: + return await interaction.response.send_message( + embed=embed_templates.error_warning(interaction, text='Fant ingen data om deg'), + ephemeral=False + ) + + self.word_freq_cache.pop(f'{interaction.user.id}', None) + + try: + self.cursor.execute( + """ + DELETE FROM wordcloud_metadata WHERE discord_user_id = %s; + DELETE FROM wordcloud_words WHERE discord_user_id = %s; + """, (interaction.user.id, interaction.user.id) + ) + except: + self.bot.db_connection.rollback() + self.bot.logger.exception("Failed to delete wordcloud metadata from database") + return await interaction.response.send_message( + embed=embed_templates.error_warning(interaction, text='Klarte ikke å slette fra database'), + ephemeral=False + ) + + embed = embed_templates.success(interaction, 'Meldingsdata er slettet!') + await interaction.response.send_message(embed=embed) + + @app_commands.checks.bot_has_permissions(embed_links=True, attach_files=True) + @app_commands.checks.cooldown(1, 10) + @wordcloud_group.command(name='data', description='Få tilsendt dine ordskydata i JSON-format') + async def data(self, interaction: discord.Interaction): + """ + Få tilsendt dine data + + Parameters + ---------- + interaction (discord.Interaction): Slash command context object + """ + + self.cursor.execute( + """ + SELECT word, frequency + FROM wordcloud_words + WHERE discord_user_id = %s + ORDER BY frequency DESC + """, (interaction.user.id,) + ) + result = self.cursor.fetchall() + + if not result: + return await interaction.response.send_message( + embed=embed_templates.error_warning(interaction, text='Fant ingen data om deg'), + ephemeral=False + ) + + freq_list = {f'{word}': freq for word, freq in result} + + embed = discord.Embed(description='Her er dataen jeg har lagret om deg') + buffer = StringIO() + buffer.write(json.dumps(freq_list, indent=4)) + buffer.seek(0) + file = discord.File(buffer, filename=f'uiog_word_freqs_{interaction.user.id}.json') + await interaction.response.send_message(embed=embed, file=file, ephemeral=True) + + @app_commands.checks.bot_has_permissions(embed_links=True, attach_files=True) + @app_commands.checks.cooldown(1, 10) + @wordcloud_generate_group.command( + name='alle', + description='Generer en ordsky basert på dine mest frekvente sagte ord siden du ga samtykke' + ) + async def generate_db(self, interaction: discord.Interaction): + """ + Generer en ordsky basert på dine mest frekvente sagte ord siden du ga samtykke + + Parameters + ---------- + interaction (discord.Interaction): Slash command context object + """ + + await interaction.response.defer() + + # Clear cache first to ensure correct count + await self.batch_update_word_freqs() + + # Fetch word count from database + self.cursor.execute( + """ + SELECT word, frequency + FROM wordcloud_words + WHERE discord_user_id = %s + """, (interaction.user.id,) + ) + results = self.cursor.fetchall() + + if not results: + return await interaction.followup.send( + embed=embed_templates.error_warning(interaction, text='Fant ingen data om deg'), + ephemeral=False + ) + + # Fetch tracking start time metadata + self.cursor.execute( + """ + SELECT tracked_since_message_channel_id, tracked_since_message_id + FROM wordcloud_metadata + WHERE discord_user_id = %s + """, (interaction.user.id,) + ) + origin_msg_channel_id, origin_msg_id = self.cursor.fetchone() + origin_msg_channel = self.bot.get_channel(origin_msg_channel_id) + origin_msg = await origin_msg_channel.fetch_message(origin_msg_id) + origin_msg_timestamp = discord.utils.format_dt(origin_msg.created_at, style='f') + + # Converts the raw SQL results tuple to a string where + # each word occurs the x number of frequency that is provided + # Having to reconstruct the data this way is probably not optimal + text = " ".join(list(itertools.chain(*([word] * count for (word, count) in results)))) + + # Generate word cloud + generation_task = functools.partial(WordCloud.generate_wordcloud, text) + word_cloud = await self.bot.loop.run_in_executor(None, generation_task) + + word_cloud_file = discord.File(word_cloud, filename=f'wordcloud_{interaction.user.id}.png') + embed = discord.Embed(title='☁️ Her er ordskyen din! ☁️') + embed.description = f'Basert på de 4000 mest frekvente ordene dine siden {origin_msg_timestamp}\n' + \ + f'[Se melding]({origin_msg.jump_url})' + embed.set_image(url=f'attachment://wordcloud_{interaction.user.id}.png') + await interaction.followup.send(embed=embed, file=word_cloud_file) + + @app_commands.checks.bot_has_permissions(embed_links=True, attach_files=True) + @app_commands.checks.cooldown(1, 30) + @wordcloud_generate_group.command( + name='siste1000', + description='Generer en ordsky basert på dine mest frekvente sagte ord fra de siste 1000 meldingene dine' + ) + async def generate_last_1000(self, interaction: discord.Interaction): + """ + Generer en ordsky basert på dine mest frekvente sagte ord fra de siste 1000 meldingene dine + + Parameters + ---------- + interaction (discord.Interaction): Slash command context object + """ + + await interaction.response.defer() + + # Get member role for future channel permission checking + member_role = interaction.guild.get_role(779849617651138601) + + # Fetch last 1000 messages + all_messages = '' + for channel in interaction.guild.text_channels: + if not channel.permissions_for(interaction.user).send_messages: + continue + + # Avoid reading secret board member chats + if not channel.permissions_for(member_role).send_messages and not channel.permissions_for(interaction.guild.default_role).send_messages: + continue + + try: + async for message in channel.history(limit=200): + if message.author != interaction.user: + continue + + # Filter bot command messages + if not message.clean_content[:2].isalpha(): + continue + + # Filter URLs + filtered_msg = re.sub( + r'https?:\/\/(www\.)?[-a-zA-Z0-9@:%._\+~#=]{1,256}\.[a-zA-Z0-9()]{1,6}\b([-a-zA-Z0-9()@:%_\+.~#?&//=]*)', # noqa: E501 + '', + message.clean_content + ) + + all_messages += f'{filtered_msg} ' + except discord.errors.Forbidden: + continue + + if not all_messages: + return await interaction.followup.send( + embed=embed_templates.error_warning(interaction, text='Fant ingen data om deg'), + ephemeral=False + ) + + # Generate word cloud + generation_task = functools.partial( + WordCloud.generate_wordcloud, + all_messages, + max_words=1000, + allow_bigrams=True + ) + word_cloud = await self.bot.loop.run_in_executor(None, generation_task) + + word_cloud_file = discord.File(word_cloud, filename=f'wordcloud_{interaction.user.id}.png') + embed = discord.Embed(title='☁️ Her er ordskyen din! ☁️') + embed.description = 'Basert på de 1000 siste meldingene i alle kanaler og ' + \ + 'topp 1000 mest frekvente ord i dine meldinger' + embed.set_image(url=f'attachment://wordcloud_{interaction.user.id}.png') + await interaction.followup.send(embed=embed, file=word_cloud_file) + + +async def setup(bot: commands.Bot): + """ + Add the cog to the bot on extension load + + Parameters + ---------- + bot (commands.Bot): Bot instance + """ + + bot.add_listener(WordCloud(bot).word_freq_listener, 'on_message') + await bot.add_cog(WordCloud(bot))