add completion message to task and add message labeling

This commit is contained in:
Alex Ott
2022-12-30 02:02:46 -08:00
parent 9c15258fd1
commit 63e2120825
3 changed files with 317 additions and 29 deletions
+166
View File
@@ -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)
+150 -28
View File
@@ -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)
+1 -1
View File
@@ -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.