From edc62bc7ced4c14a16d4c6eca84273e6af8d8db2 Mon Sep 17 00:00:00 2001 From: Kaveen Kumarasinghe Date: Sun, 18 Dec 2022 22:47:55 -0500 Subject: [PATCH] Auto-expire interactions (1 hour for gpt messages), (10 minutes for drawings) --- cogs/draw_image_generation.py | 33 ++++++++++++++++++------- cogs/gpt_3_commands_and_converser.py | 36 +++++++++++++++++----------- cogs/image_prompt_optimizer.py | 2 +- 3 files changed, 47 insertions(+), 24 deletions(-) diff --git a/cogs/draw_image_generation.py b/cogs/draw_image_generation.py index ced39ae..874c4e6 100644 --- a/cogs/draw_image_generation.py +++ b/cogs/draw_image_generation.py @@ -207,20 +207,35 @@ class DrawDallEService(commands.Cog, name="DrawDallEService"): class SaveView(discord.ui.View): - def __init__(self, image_urls, cog, converser_cog, no_retry=False): - super().__init__() + def __init__(self, image_urls, cog, converser_cog, no_retry=False, only_save=None): + super().__init__(timeout=600 if not only_save else None) # 10 minute timeout for Retry, Save + self.image_urls = image_urls self.cog = cog + self.no_retry = no_retry self.converser_cog = converser_cog for x in range(1, len(image_urls) + 1): self.add_item(SaveButton(x, image_urls[x - 1])) - if not no_retry: - self.add_item(RedoButton(self.cog, converser_cog=self.converser_cog)) - for x in range(1, len(image_urls) + 1): - self.add_item( - VaryButton( - x, image_urls[x - 1], self.cog, converser_cog=self.converser_cog + if not only_save: + if not no_retry: + self.add_item(RedoButton(self.cog, converser_cog=self.converser_cog)) + for x in range(1, len(image_urls) + 1): + self.add_item( + VaryButton( + x, image_urls[x - 1], self.cog, converser_cog=self.converser_cog + ) ) - ) + + # On the timeout event, override it and we want to clear the items. + async def on_timeout(self): + # Save all the SaveButton items, then clear all the items, then add back the SaveButton items, then + # update the message + self.clear_items() + + # Create a new view with the same params as this one, but pass only_save=True + new_view = SaveView(self.image_urls, self.cog, self.converser_cog, self.no_retry, only_save=True) + + # Set the view of the message to the new view + await self.message.edit(view=new_view) class VaryButton(discord.ui.Button): diff --git a/cogs/gpt_3_commands_and_converser.py b/cogs/gpt_3_commands_and_converser.py index 6dd95c4..cb9fe49 100644 --- a/cogs/gpt_3_commands_and_converser.py +++ b/cogs/gpt_3_commands_and_converser.py @@ -379,13 +379,10 @@ class GPT3ComCon(commands.Cog, name="GPT3ComCon"): if len(response_text) > self.TEXT_CUTOFF: await self.paginate_and_send(response_text, message) else: - response_message = await message.reply(response_text) + response_message = await message.reply(response_text, view=RedoView(self)) redo_users[message.author.id] = RedoUser( prompt, message, response_message ) - RedoButtonView.bot = self - await response_message.edit(view=RedoButtonView()) - else: # We have response_text available, this is the original message that we want to edit await response_message.edit(content=response_text) @@ -547,15 +544,26 @@ class GPT3ComCon(commands.Cog, name="GPT3ComCon"): # Send the request to the model await self.encapsulated_send(message, prompt) +class RedoView(discord.ui.View): + def __init__(self, converser_cog): + super().__init__(timeout=3600) # 1 hour interval to redo. + self.converser_cog = converser_cog + self.add_item(RedoButton(self.converser_cog)) + + async def on_timeout(self): + # Remove the button from the view/message + self.clear_items() + # Send a message to the user saying the view has timed out + await self.message.edit( + view=None, + ) + +class RedoButton(discord.ui.Button["RedoView"]): + def __init__(self, converser_cog): + super().__init__(style=discord.ButtonStyle.danger, label="Retry") + self.converser_cog = converser_cog -class RedoButtonView( - discord.ui.View -): # Create a class called MyView that subclasses discord.ui.View - @discord.ui.button( - label="Retry", - style=discord.ButtonStyle.danger, - ) # Create a button with the label "😎 Click me!" with color Blurple - async def button_callback(self, button, interaction): + async def callback(self, interaction: discord.Interaction): msg = await interaction.response.send_message( "Retrying your original request...", ephemeral=True ) @@ -564,7 +572,7 @@ class RedoButtonView( deletion = Deletion( msg, (datetime.datetime.now() + datetime.timedelta(seconds=10)).timestamp() ) - await self.bot.deletion_queue.put(deletion) + await self.converser_cog.deletion_queue.put(deletion) # Get the user user_id = interaction.user.id @@ -573,4 +581,4 @@ class RedoButtonView( message = redo_users[user_id].message prompt = redo_users[user_id].prompt response_message = redo_users[user_id].response - await self.bot.encapsulated_send(message, prompt, response_message) + await self.converser_cog.encapsulated_send(message, prompt, response_message) diff --git a/cogs/image_prompt_optimizer.py b/cogs/image_prompt_optimizer.py index 93a4584..ebd9575 100644 --- a/cogs/image_prompt_optimizer.py +++ b/cogs/image_prompt_optimizer.py @@ -113,7 +113,7 @@ class ImgPromptOptimizer(commands.Cog, name="ImgPromptOptimizer"): class OptimizeView(discord.ui.View): def __init__(self, converser_cog, image_service_cog, deletion_queue): - super().__init__() + super().__init__(timeout=None) self.cog = converser_cog self.image_service_cog = image_service_cog self.deletion_queue = deletion_queue