From 63e2120825b9b3565f32b3b329d386d1baea472e Mon Sep 17 00:00:00 2001 From: Alex Ott <66271487+AlexanderHOtt@users.noreply.github.com> Date: Fri, 30 Dec 2022 02:02:46 -0800 Subject: [PATCH] add completion message to task and add message labeling --- discord-bot/bot/extensions/text_labels.py | 166 ++++++++++++++++++++ discord-bot/bot/extensions/work.py | 178 ++++++++++++++++++---- discord-bot/bot/utils.py | 2 +- 3 files changed, 317 insertions(+), 29 deletions(-) create mode 100644 discord-bot/bot/extensions/text_labels.py diff --git a/discord-bot/bot/extensions/text_labels.py b/discord-bot/bot/extensions/text_labels.py new file mode 100644 index 00000000..ab7d109d --- /dev/null +++ b/discord-bot/bot/extensions/text_labels.py @@ -0,0 +1,166 @@ +# -*- coding: utf-8 -*- +"""Hot reload plugin.""" +import hikari +import lightbulb +import miru +from datetime import datetime + +import typing as t + +plugin = lightbulb.Plugin( + "HotReloadPlugin", +) + +from bot.utils import EMPTY + +DISCORD_GRAY = 0x2F3136 + + +def clamp(num: float) -> float: + """Clamp a number between 0 and 1.""" + return min(max(0.0, num), 1.0) + + +class LabelModal(miru.Modal): + """Modal for submitting text labels.""" + + def __init__(self, label: str, content: str, *args: t.Any, **kwargs: t.Any): + super().__init__(*args, **kwargs) + self.label = label + self.original_content = content + + # Add the text of the message to the modal + self.content = miru.TextInput( + label="Text", style=hikari.TextInputStyle.PARAGRAPH, value=content, required=True, row=1 + ) + self.add_item(self.content) + + value = miru.TextInput(label="Value", placeholder="Enter a value between 0 and 1", required=True, row=2) + + async def callback(self, context: miru.ModalContext) -> None: + val = float(self.value.value) if self.value.value else 0.0 + val = clamp(val) + + edited = self.content.value != self.original_content + await context.respond( + f"Sending {self.label}=`{val}` for `{self.content.value}` (edited={edited}) to the backend.", + flags=hikari.MessageFlag.EPHEMERAL, + ) + + # Send a notification to the log channel + embed = ( + hikari.Embed( + title="Message Label", + description=f"{context.author.mention} labeled a message as `{self.label}`.", + timestamp=datetime.now().astimezone(), + color=0x00FF00, + ) + .set_author(name=context.author.username, icon=context.author.avatar_url) + .add_field("Total Labeled Message", "0", inline=True) + .add_field("Server Ranking", "0/0", inline=True) + .add_field("Global Ranking", "0/0", inline=True) + .set_footer(f"Message ID: TODO") + ) + channel = await context.bot.rest.fetch_channel(1058299131115872297) + assert isinstance(channel, hikari.TextableChannel) + await channel.send(EMPTY, embed=embed) + + +class LabelSelect(miru.View): + """Select menu for selecting a label. + + The current labels are: + - contains toxic language + - encourages illegal activity + - good quality + - bad quality + - is spam + """ + + def __init__(self, content: str, *args: t.Any, **kwargs: t.Any): + super().__init__(*args, **kwargs) + self.content = content + + @miru.select( + options=[ + hikari.SelectMenuOption( + label="Toxic Language", + value="toxic_language", + description="The message contains toxic language.", + is_default=False, + emoji=None, + ), + hikari.SelectMenuOption( + label="Illegal Activity", + value="illegal_activity", + description="The message encourages illegal activity.", + is_default=False, + emoji=None, + ), + hikari.SelectMenuOption( + label="Good Quality", + value="good_quality", + description="The message is good quality.", + is_default=False, + emoji=None, + ), + hikari.SelectMenuOption( + label="Bad Quality", + value="bad_quality", + description="The message is bad quality.", + is_default=False, + emoji=None, + ), + hikari.SelectMenuOption( + label="Spam", + value="spam", + description="The message is spam.", + is_default=False, + emoji=None, + ), + ], + min_values=1, + max_values=1, + ) + async def label_select(self, select: miru.Select, ctx: miru.ViewContext) -> None: + """Handle the select menu.""" + label = select.values[0] + modal = LabelModal(label, self.content, title=f"Text Label: {label}", timeout=60) + await modal.send(ctx.interaction) + await modal.wait() + + self.stop() + + +@plugin.command +@lightbulb.command("Label Message", "Label a message") +@lightbulb.implements(lightbulb.MessageCommand) +async def label_message_text(ctx: lightbulb.MessageContext): + """Label a message.""" + msg: hikari.Message = ctx.options.target + # Exit if the message is empty + if not msg.content: + await ctx.respond("Cannot label an empty message.", flags=hikari.MessageFlag.EPHEMERAL) + return + + # Send the select menu + # The modal will be opened from the select menu interaction + embed = hikari.Embed(title="Label Message", description="Select a label for the message.", color=DISCORD_GRAY) + label_select_view = LabelSelect( + msg.content, + timeout=60, + ) + resp = await ctx.respond(EMPTY, embed=embed, components=label_select_view, flags=hikari.MessageFlag.EPHEMERAL) + + await label_select_view.start(await resp.message()) + await label_select_view.wait() + + +def load(bot: lightbulb.BotApp): + """Add the plugin to the bot.""" + bot.add_plugin(plugin) + + +def unload(bot: lightbulb.BotApp): + """Remove the plugin to the bot.""" + bot.remove_plugin(plugin) diff --git a/discord-bot/bot/extensions/work.py b/discord-bot/bot/extensions/work.py index 1175e4a5..86462f87 100644 --- a/discord-bot/bot/extensions/work.py +++ b/discord-bot/bot/extensions/work.py @@ -13,7 +13,7 @@ from oasst_shared.schemas import protocol as protocol_schema from oasst_shared.schemas.protocol import TaskRequestType from bot.api_client import OasstApiClient, TaskType -from bot.utils import ZWJ +from bot.utils import EMPTY plugin = lightbulb.Plugin("WorkPlugin") @@ -106,6 +106,24 @@ async def _handle_task(ctx: lightbulb.SlashContext, task_type: TaskRequestType) else: logger.fatal(f"Unexpected task type received: {new_task.type}") + # Send a message in the log channel that the task is complete + # TODO: Maybe do something with the msg ID + done_embed = ( + hikari.Embed( + title="Task Completion", + description=f"`{task.type}` completed by {ctx.author.mention}", + color=hikari.Color(0x00FF00), + timestamp=datetime.now().astimezone(), + ) + .add_field("Total Tasks", "0", inline=True) + .add_field("Server Ranking", "0/0", inline=True) + .add_field("Global Ranking", "0/0", inline=True) + .set_footer(f"Task ID: {task.id}") + ) + channel = await ctx.bot.rest.fetch_channel(1058299131115872297) + assert isinstance(channel, hikari.TextableChannel) + await channel.send(EMPTY, embed=done_embed) + # ask the user if they want to do another task choice_view = ChoiceView(timeout=MAX_TASK_ACCEPT_TIME) msg = await ctx.author.send("Would you like another task?", components=choice_view) @@ -169,7 +187,6 @@ async def _send_task( # The clean way to do this would be to attach a `to_embed` method to the task classes # but the tasks aren't discord specific so that doesn't really make sense. - view = TaskAcceptView(timeout=MAX_TASK_ACCEPT_TIME) embed: hikari.UndefinedOr[hikari.Embed] = hikari.UNDEFINED # Create an embed based on the task's type @@ -183,11 +200,38 @@ async def _send_task( logger.info("sending rank initial prompt task") embed = _rank_initial_prompt_embed(task) + elif task.type == TaskRequestType.rank_user_replies: + assert isinstance(task, protocol_schema.RankUserRepliesTask) + logger.info("sending rank user reply task") + embed = _rank_user_reply_embed(task) + + elif task.type == TaskRequestType.rank_assistant_replies: + assert isinstance(task, protocol_schema.RankAssistantRepliesTask) + logger.info("sending rank assistant reply task") + embed = _rank_assistant_reply_embed(task) + + elif task.type == TaskRequestType.user_reply: + assert isinstance(task, protocol_schema.UserReplyTask) + logger.info("sending user reply task") + embed = _user_reply_embed(task) + + elif task.type == TaskRequestType.assistant_reply: + assert isinstance(task, protocol_schema.AssistantReplyTask) + logger.info("sending assistant reply task") + embed = _assistant_reply_embed(task) + + elif task.type == TaskRequestType.summarize_story: + raise NotImplementedError + elif task.type == TaskRequestType.rate_summary: + raise NotImplementedError + else: logger.error(f"unknown task type {task.type}") + raise ValueError(f"unknown task type {task.type}") + view = TaskAcceptView(timeout=MAX_TASK_ACCEPT_TIME) msg = await ctx.author.send( - ZWJ, + EMPTY, embed=embed, components=view, ) @@ -200,31 +244,6 @@ async def _send_task( return view.choice, str(msg.id) -def _initial_prompt_embed(task: protocol_schema.InitialPromptTask) -> hikari.Embed: - return ( - hikari.Embed(title="Initial Prompt", description=f"Hint: {task.hint}", timestamp=datetime.now().astimezone()) - .set_image("https://images.unsplash.com/photo-1455390582262-044cdead277a?w=512") - .set_footer(text=f"OASST Assistant | {task.id}") - ) - - -def _rank_initial_prompt_embed(task: protocol_schema.RankInitialPromptsTask) -> hikari.Embed: - embed = ( - hikari.Embed( - title="Rank Initial Prompt", - description="Rank the following tasks from best to worst (1,2,3,4,5)", - timestamp=datetime.now().astimezone(), - ) - .set_image("https://images.unsplash.com/photo-1455390582262-044cdead277a?w=512") - .set_footer(text=f"OASST Assistant | {task.id}") - ) - - for i, prompt in enumerate(task.prompts): - embed.add_field(name=f"Prompt {i + 1}", value=prompt, inline=False) - - return embed - - class TaskAcceptView(miru.View): """View with three buttons: accept, next, and cancel. @@ -271,6 +290,109 @@ class ChoiceView(miru.View): self.stop() +################################################################ +# Template Embeds # +################################################################ + +# TODO: Maybe implement a better way of creating embeds, like `from_json` or something + + +def _initial_prompt_embed(task: protocol_schema.InitialPromptTask) -> hikari.Embed: + return ( + hikari.Embed(title="Initial Prompt", description=f"Hint: {task.hint}", timestamp=datetime.now().astimezone()) + .set_image("https://images.unsplash.com/photo-1455390582262-044cdead277a?w=512") + .set_footer(text=f"OASST Assistant | {task.id}") + ) + + +def _rank_initial_prompt_embed(task: protocol_schema.RankInitialPromptsTask) -> hikari.Embed: + embed = ( + hikari.Embed( + title="Rank Initial Prompt", + description="Rank the following tasks from best to worst (1,2,3,4,5)", + timestamp=datetime.now().astimezone(), + ) + .set_image("https://images.unsplash.com/photo-1455390582262-044cdead277a?w=512") + .set_footer(text=f"OASST Assistant | {task.id}") + ) + + for i, prompt in enumerate(task.prompts): + embed.add_field(name=f"Prompt {i + 1}", value=prompt, inline=False) + + return embed + + +def _rank_user_reply_embed(task: protocol_schema.RankUserRepliesTask) -> hikari.Embed: + embed = ( + hikari.Embed( + title="Rank User Reply", + description="Rank the following tasks from best to worst. e.g. 1,2,5,3,4", + timestamp=datetime.now().astimezone(), + ) + .set_image("https://images.unsplash.com/photo-1455390582262-044cdead277a?w=512") # TODO: update image + .set_footer(text=f"OASST Assistant | {task.id}") + ) + + for i, reply in enumerate(task.replies): + embed.add_field(name=f"Reply {i + 1}", value=reply, inline=False) + + return embed + + +def _rank_assistant_reply_embed(task: protocol_schema.RankAssistantRepliesTask) -> hikari.Embed: + embed = ( + hikari.Embed( + title="Rank Assistant Reply", + description="Rank the following tasks from best to worst. e.g. 1,2,5,3,4", + timestamp=datetime.now().astimezone(), + ) + .set_image("https://images.unsplash.com/photo-1455390582262-044cdead277a?w=512") # TODO: update image + .set_footer(text=f"OASST Assistant | {task.id}") + ) + + for i, reply in enumerate(task.replies): + embed.add_field(name=f"Reply {i + 1}", value=reply, inline=False) + + return embed + + +def _user_reply_embed(task: protocol_schema.UserReplyTask) -> hikari.Embed: + embed = ( + hikari.Embed( + title="User Reply", + description=f"""\ + Send the next message in the conversation as if you were the user. + {'Hint: ' if task.hint else ''} + """, + timestamp=datetime.now().astimezone(), + ) + # .set_image("https://images.unsplash.com/photo-1455390582262-044cdead277a?w=512") # TODO: change image + .set_footer(text=f"OASST Assistant | {task.id}") + ) + + for message in task.conversation.messages: + embed.add_field(name="Assistant" if message.is_assistant else "User", value=message.text, inline=False) + + return embed + + +def _assistant_reply_embed(task: protocol_schema.AssistantReplyTask) -> hikari.Embed: + embed = ( + hikari.Embed( + title="User Reply", + description="Send the next message in the conversation as if you were the user.", + timestamp=datetime.now().astimezone(), + ) + # .set_image("https://images.unsplash.com/photo-1455390582262-044cdead277a?w=512") # TODO: change image + .set_footer(text=f"OASST Assistant | {task.id}") + ) + + for message in task.conversation.messages: + embed.add_field(name="Assistant" if message.is_assistant else "User", value=message.text, inline=False) + + return embed + + def load(bot: lightbulb.BotApp): """Add the plugin to the bot.""" bot.add_plugin(plugin) diff --git a/discord-bot/bot/utils.py b/discord-bot/bot/utils.py index 1ff6ef1f..1ce99560 100644 --- a/discord-bot/bot/utils.py +++ b/discord-bot/bot/utils.py @@ -23,7 +23,7 @@ def format_time(dt: datetime, fmt: t.Literal["t", "T", "D", "f", "F", "R"]) -> s raise ValueError(f"`fmt` must be 't', 'T', 'D', 'f', 'F' or 'R', not {fmt}") -ZWJ = "\u200d" +EMPTY = "\u200d" """Zero-width joiner. This appears as an empty message in Discord.