import asyncio from datetime import datetime, timedelta import discord from redbot.core import Config from redbot.core import commands from chatter.chatterbot import ChatBot from chatter.chatterbot.trainers import ListTrainer class Chatter: """ This cog trains a chatbot that will talk like members of your Guild """ def __init__(self, bot): self.bot = bot self.config = Config.get_conf(self, identifier=6710497116116101114) default_global = {} default_guild = { "whitelist": None, "days": 1 } self.chatbot = ChatBot( "ChatterBot", storage_adapter='chatter.chatterbot.storage.SQLStorageAdapter', database='./database.sqlite3' ) self.chatbot.set_trainer(ListTrainer) self.config.register_global(**default_global) self.config.register_guild(**default_guild) self.loop = asyncio.get_event_loop() async def _get_conversation(self, ctx, in_channel: discord.TextChannel = None): """ Compiles all conversation in the Guild this bot can get it's hands on Currently takes a stupid long time Returns a list of text """ out = [] after = datetime.today() - timedelta(days=(await self.config.guild(ctx.guild).days())) for channel in ctx.guild.text_channels: if in_channel: channel = in_channel await ctx.send("Gathering {}".format(channel.mention)) user = None try: async for message in channel.history(limit=None, reverse=True, after=after): if user == message.author: out[-1] += "\n" + message.clean_content else: user = message.author out.append(message.clean_content) except discord.Forbidden: pass except discord.HTTPException: pass if in_channel: break return out def _train(self, data): try: self.chatbot.train(data) except: return False return True @commands.group() async def chatter(self, ctx: commands.Context): """ Base command for this cog. Check help for the commands list. """ if ctx.invoked_subcommand is None: await ctx.send_help() @chatter.command() async def age(self, ctx: commands.Context, days: int): """ Sets the number of days to look back Will train on 1 day otherwise """ await self.config.guild(ctx.guild).days.set(days) await ctx.send("Success") @chatter.command() async def backup(self, ctx, backupname): """ Backup your training data to a json for later use """ await ctx.send("Backing up data, this may take a while") future = await self.loop.run_in_executor(None, self.chatbot.trainer.export_for_training, './{}.json'.format(backupname)) if future: await ctx.send("Backup successful!") else: await ctx.send("Error occurred :(") @chatter.command() async def train(self, ctx: commands.Context, channel: discord.TextChannel): """ Trains the bot based on language in this guild """ conversation = await self._get_conversation(ctx, channel) if not conversation: await ctx.send("Failed to gather training data") return await ctx.send("Gather successful! Training begins now\n(**This will take a long time, be patient**)") embed = discord.Embed(title="Loading") embed.set_image(url="http://www.loop.universaleverything.com/animations/1295.gif") temp_message = await ctx.send(embed=embed) future = await self.loop.run_in_executor(None, self._train, conversation) try: await temp_message.delete() except: pass if future: await ctx.send("Training successful!") else: await ctx.send("Error occurred :(") async def on_message(self, message): """ Credit to https://github.com/Twentysix26/26-Cogs/blob/master/cleverbot/cleverbot.py for on_message recognition of @bot """ author = message.author channel = message.channel if message.author.id != self.bot.user.id: to_strip = "@" + author.guild.me.display_name + " " text = message.clean_content if not text.startswith(to_strip): return text = text.replace(to_strip, "", 1) async with channel.typing(): future = await self.loop.run_in_executor(None, self.chatbot.get_response, text) if future: await channel.send(str(future)) else: await channel.send(':thinking:')