diff --git a/backend/oasst_backend/api/v1/tasks.py b/backend/oasst_backend/api/v1/tasks.py index 40965022..dac2a9bd 100644 --- a/backend/oasst_backend/api/v1/tasks.py +++ b/backend/oasst_backend/api/v1/tasks.py @@ -17,7 +17,7 @@ router = APIRouter() def generate_task(request: protocol_schema.TaskRequest) -> protocol_schema.Task: - match (request.type): + match request.type: case protocol_schema.TaskRequestType.random: logger.info("Frontend requested a random task.") while request.type == protocol_schema.TaskRequestType.random: diff --git a/bot/__main__.py b/bot/__main__.py index 1c456849..362b16f0 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -11,5 +11,6 @@ if __name__ == "__main__": bot_channel_name=settings.BOT_CHANNEL_NAME, backend_url=settings.BACKEND_URL, api_key=settings.API_KEY, + owner_id=settings.OWNER_ID, ) bot.run() diff --git a/bot/bot.py b/bot/bot.py index 2d809646..e6e90770 100644 --- a/bot/bot.py +++ b/bot/bot.py @@ -1,9 +1,11 @@ # -*- coding: utf-8 -*- import asyncio -from typing import Any +from typing import Optional, Union import discord from api_client import ApiClient, TaskType +from discord import app_commands +from loguru import logger from oasst_shared.schemas import protocol as protocol_schema @@ -24,47 +26,74 @@ def generate_rating_view(lo: int, hi: int, response_handler) -> discord.ui.View: return view -class ModifiedClient(discord.Client): - def __init__(self, *, intents: discord.Intents, **options: Any): - super().__init__(intents=intents, **options) - - async def setup_hook(self): - print("setup") - - class OpenAssistantBot: - def __init__(self, bot_token: str, bot_channel_name: str, backend_url: str, api_key: str): + def __init__( + self, + bot_token: str, + bot_channel_name: str, + backend_url: str, + api_key: str, + owner_id: Optional[Union[int, str]] = None, + ): intents = discord.Intents.default() intents.message_content = True + + if isinstance(owner_id, str): + owner_id = int(owner_id) + self.owner_id = owner_id + self.bot_token = bot_token - client = ModifiedClient(intents=intents) + client = discord.Client(intents=intents) self.client = client + self.bot_channel: discord.TextChannel = None self.backend = ApiClient(backend_url, api_key) self.reply_handlers = {} # handlers by msg_id + self.tree = app_commands.CommandTree(self.client, fallback_to_global=True) + + self.auto_archive_minutes = 60 # ToDo: add to bot config @client.event async def on_ready(): self.bot_channel = self.get_text_channel_by_name(bot_channel_name) - client.loop.create_task(self.background_timer(), name="OpenAssistantBot.background_timer()") - print(f"{client.user} is now running!") + logger.info(f"{client.user} is now running!") @client.event async def on_message(message: discord.Message): # ignore own messages - if message.author == client.user: - return + if message.author != client.user: + await self.handle_message(message) - await self.handle_message(message) + @self.tree.command() + async def tutorial(interaction: discord.Interaction): + """Start the Open-Assistant tutorial via DMs.""" + await interaction.response.send_message(f"tutorial command by {interaction.user.name}") + + @self.tree.command() + async def help(interaction: discord.Interaction): + """Sends the user a list of all available commands""" + await interaction.response.send_message(f"help command by {interaction.user.name}") + + @self.tree.command() + async def work(interaction: discord.Interaction): + """Request a new personalized task""" + await interaction.response.send_message(f"work command by {interaction.user.name}") + + async def print_separtor(self, title: str) -> discord.Message: + msg: discord.Message = await self.bot_channel.send(f"\n:point_right: {title} :point_left:\n") + return msg async def generate_summarize_story(self, task: protocol_schema.SummarizeStoryTask): text = f"Summarize to the following story:\n{task.story}" msg: discord.Message = await self.bot_channel.send(text) + await self.bot_channel.create_thread( + message=discord.Object(msg.id), name="Summaries", auto_archive_duration=self.auto_archive_minutes + ) async def on_reply(message: discord.Message): - print("on_summarize_story_reply", message) - await message.reply("thx, on_summarize_story_reply") + logger.info("on_summarize_story_reply", message) + await message.add_reaction("✅") self.reply_handlers[msg.id] = on_reply @@ -81,15 +110,15 @@ class OpenAssistantBot: text = "\n".join(s) async def rating_response_handler(score, interaction: discord.Interaction): - print("rating_response_handler", score) + logger.info("rating_response_handler", score) await interaction.response.send_message(f"got your feedback: {score}") view = generate_rating_view(task.scale.min, task.scale.max, rating_response_handler) msg: discord.Message = await self.bot_channel.send(text, view=view) async def on_reply(message: discord.Message): - print("on_summary_reply", message) - await message.reply("thx, on_summary_reply") + logger.info("on_summary_reply", message) + await message.add_reaction("") self.reply_handlers[msg.id] = on_reply @@ -100,10 +129,13 @@ class OpenAssistantBot: if task.hint: text += f"\nHint: {task.hint}" msg: discord.Message = await self.bot_channel.send(text) + await self.bot_channel.create_thread( + message=discord.Object(msg.id), name="Prompts", auto_archive_duration=self.auto_archive_minutes + ) async def on_reply(message: discord.Message): - print("on_initial_prompt_reply", message) - await message.reply("thx, on_initial_prompt_reply") + logger.info("on_initial_prompt_reply", message) + await message.add_reaction("✅") self.reply_handlers[msg.id] = on_reply @@ -112,36 +144,46 @@ class OpenAssistantBot: def _render_message(self, message: protocol_schema.ConversationMessage) -> str: """Render a message to the user.""" if message.is_assistant: - return f"Assistant: {message.text}" - return f"User: {message.text}" + return f":robot: Assistant:\n{message.text}" + else: + return f":person_red_hair: User:\n**{message.text}**" async def generate_user_reply(self, task: protocol_schema.UserReplyTask): - s = ["Please provide a reply to the assistant.", "Here is the conversation so far:"] + s = ["Please provide a reply to the assistant.", "Here is the conversation so far:\n"] for message in task.conversation.messages: s.append(self._render_message(message)) + s.append("") if task.hint: s.append(f"Hint: {task.hint}") text = "\n".join(s) msg: discord.Message = await self.bot_channel.send(text) + await self.bot_channel.create_thread( + message=discord.Object(msg.id), name="User responses", auto_archive_duration=self.auto_archive_minutes + ) async def on_reply(message: discord.Message): - print("on_user_reply_reply", message) - await message.reply("thx, on_user_reply_reply") + logger.info("on_user_reply_reply", message) + await message.add_reaction("✅") self.reply_handlers[msg.id] = on_reply return msg async def generate_assistant_reply(self, task: protocol_schema.AssistantReplyTask): - s = ["Act as the assistant and reply to the user.", "Here is the conversation so far:"] + s = ["Act as the assistant and reply to the user.", "Here is the conversation so far\n:"] for message in task.conversation.messages: s.append(self._render_message(message)) + s.append("") + s.append(":robot: Assistant: { human, pls help me! ... }") text = "\n".join(s) msg: discord.Message = await self.bot_channel.send(text) + await self.bot_channel.create_thread( + message=discord.Object(msg.id), name="Agent responses", auto_archive_duration=self.auto_archive_minutes + ) async def on_reply(message: discord.Message): - print("on_assistant_reply_reply", message) - await message.reply("thx, on_assistant_reply_reply") + logger.info("on_assistant_reply_reply", message) + await message.add_reaction("✅") self.reply_handlers[msg.id] = on_reply @@ -151,12 +193,17 @@ class OpenAssistantBot: s = ["Rank the following prompts:"] for idx, prompt in enumerate(task.prompts, start=1): s.append(f"{idx}: {prompt}") + s.append("") + s.append(':scroll: Reply with the numbers of best to worst prompts separated by commas (example: "4,1,3,2").') text = "\n".join(s) msg: discord.Message = await self.bot_channel.send(text) + await self.bot_channel.create_thread( + message=discord.Object(msg.id), name="User responses", auto_archive_duration=self.auto_archive_minutes + ) async def on_reply(message: discord.Message): - print("on_rank_initial_prompts_reply", message) - await message.reply("thx, on_rank_initial_prompts_reply") + logger.info("on_rank_initial_prompts_reply", message) + await message.add_reaction("✅") self.reply_handlers[msg.id] = on_reply @@ -166,14 +213,21 @@ class OpenAssistantBot: s = ["Here is the conversation so far:"] for message in task.conversation.messages: s.append(self._render_message(message)) + s.append("") s.append("Rank the following replies:") for idx, reply in enumerate(task.replies, start=1): s.append(f"{idx}: {reply}") + s.append("") + s.append(':scroll: Reply with the numbers of best to worst prompts separated by commas (example: "4,1,3,2").') text = "\n".join(s) msg: discord.Message = await self.bot_channel.send(text) + await self.bot_channel.create_thread( + message=discord.Object(msg.id), name="User responses", auto_archive_duration=self.auto_archive_minutes + ) async def on_reply(message: discord.Message): - print("on_rrank_conversation_reply", message) + logger.info("on_rank_conversation_reply", message) + await message.add_reaction("✅") message self.reply_handlers[msg.id] = on_reply @@ -181,9 +235,11 @@ class OpenAssistantBot: return msg async def next_task(self): - # task = self.backend.fetch_task(protocol_schema.TaskRequestType.rate_summary, user=None) + # task = self.backend.fetch_task(protocol_schema.TaskRequestType.user_reply, user=None) task = self.backend.fetch_random_task(user=None) + await self.print_separtor("New Task") + msg: discord.Message = None match task.type: case TaskType.summarize_story: @@ -211,26 +267,70 @@ class OpenAssistantBot: if self.bot_channel: try: await self.next_task() - except Exception as e: - print(e.with_traceback()) + except Exception: + logger.exception("fetching next task failed") await asyncio.sleep(30) - def run(self): - """Run bot loop blocking.""" - self.client.run(self.bot_token) + async def _sync(self, command: str, message: discord.Message): + + logger.info(f"sync tree command received: {command}") + + if command == "sync.copy_global": + await self.tree.copy_global_to(guild=message.guild) + synced = await self.tree.sync(guild=message.guild) + elif command == "sync.clear_guild": + self.tree.clear_commands(guild=message.guild) + synced = await self.tree.sync(guild=message.guild) + elif command == "sync.guild": + synced = await self.tree.sync(guild=message.guild) + else: + synced = await self.tree.sync() + + logger.info(f"Synced {len(synced)} commands") + await message.reply(f"Synced {len(synced)} commands") + + async def handle_command(self, message: discord.Message, is_owner: bool): + command_text: str = message.content + command_text = command_text[1:] + match command_text: + case "sync" | "sync.guild" | "sync.copy_global" | "sync.clear_guild" | "sync.clear_guild": + if is_owner: + await self._sync(command_text, message) + case _: + await message.reply(f"unknown command: {command_text}") async def handle_message(self, message: discord.Message): user_id = message.author.id user_display_name = message.author.name + command_prefix = "!" + if ( + message.channel.type == discord.ChannelType.private + and message.type == discord.MessageType.default + and message.content.startswith(command_prefix) + ): + is_owner = self.owner_id and user_id == self.owner_id + await self.handle_command(message, is_owner) + + if isinstance(message.channel, discord.Thread): + handler = self.reply_handlers.get(message.channel.id) + if handler: + await handler(message) + if message.reference: handler = self.reply_handlers.get(message.reference.message_id) if handler: await handler(message) - print(user_id, user_display_name, message.content, type(message.content)) + logger.debug( + f"{message.type} {message.channel.type} from ({user_display_name}) {user_id}: {message.content} ({type(message.content)})" + ) def get_text_channel_by_name(self, channel_name) -> discord.TextChannel: for channel in self.client.get_all_channels(): if channel.type == discord.ChannelType.text and channel.name == channel_name: return channel + + def run(self): + """Run bot loop blocking.""" + self.client.run(self.bot_token) diff --git a/bot/bot_settings.py b/bot/bot_settings.py index 3323b2fe..b7a46aa6 100644 --- a/bot/bot_settings.py +++ b/bot/bot_settings.py @@ -7,7 +7,7 @@ class BotSettings(BaseSettings): API_KEY: str = "any_key" BOT_TOKEN: str BOT_CHANNEL_NAME: str = "bot" - TEST_GUILD: str = None + OWNER_ID: int = None settings = BotSettings(_env_file=".env")