mirror of
https://github.com/wassname/Open-Assistant.git
synced 2026-07-21 12:20:08 +08:00
add completion message to task and add message labeling
This commit is contained in:
@@ -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)
|
||||
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user