diff --git a/discord-bot/bot/api_client.py b/discord-bot/bot/api_client.py index cec1900f..9f319869 100644 --- a/discord-bot/bot/api_client.py +++ b/discord-bot/bot/api_client.py @@ -1,5 +1,5 @@ # -*- coding: utf-8 -*- -import asyncio +"""API Client for interacting with the OASST backend.""" import enum import typing as t from typing import Optional, Type @@ -7,11 +7,12 @@ from uuid import UUID import aiohttp from loguru import logger - from oasst_shared.schemas import protocol as protocol_schema class TaskType(str, enum.Enum): + """Task types.""" + summarize_story = "summarize_story" rate_summary = "rate_summary" initial_prompt = "initial_prompt" @@ -30,6 +31,7 @@ class OasstApiClient: """Create a new OasstApiClient. Args: + ---- backend_url (str): The base backend URL. api_key (str): The API key to use for authentication. """ @@ -38,7 +40,7 @@ class OasstApiClient: self.backend_url = backend_url self.api_key = api_key - self.task_models_map: dict[str, Type[protocol_schema.Task]] = { + self.task_models_map: dict[TaskType, Type[protocol_schema.Task]] = { TaskType.summarize_story: protocol_schema.SummarizeStoryTask, TaskType.rate_summary: protocol_schema.RateSummaryTask, TaskType.initial_prompt: protocol_schema.InitialPromptTask, @@ -58,17 +60,13 @@ class OasstApiClient: return await response.json() def _parse_task(self, data: dict[str, t.Any]) -> protocol_schema.Task: - task_type = data.get("type") - - if not isinstance(task_type, str): - logger.error(f"task type must be a `str`: {task_type}") - raise ValueError(f"task type must be a `str`: {task_type}") + task_type = TaskType(data.get("type")) model = self.task_models_map.get(task_type) if not model: logger.error(f"Unsupported task type: {task_type}") raise ValueError(f"Unsupported task type: {task_type}") - return self.task_models_map[task_type].parse_obj(data) + return self.task_models_map[task_type].parse_obj(data) # type: ignore async def fetch_task( self, task_type: protocol_schema.TaskRequestType, user: Optional[protocol_schema.User] = None @@ -76,8 +74,8 @@ class OasstApiClient: """Fetch a task from the backend.""" logger.debug(f"Fetching task {task_type} for user {user}") req = protocol_schema.TaskRequest(type=task_type.value, user=user) - resp = await self.post(f"/api/v1/tasks/", data=req.dict()) - print("resp", resp) + resp = await self.post("/api/v1/tasks/", data=req.dict()) + logger.debug(f"Fetch task response: {resp}") return self._parse_task(resp) async def fetch_random_task(self, user: Optional[protocol_schema.User] = None) -> protocol_schema.Task: @@ -107,24 +105,3 @@ class OasstApiClient: async def close(self): logger.debug("Closing OasstApiClient session") await self.session.close() - - -async def main(): - api = OasstApiClient("http://localhost:8080", "test") - try: - task = await api.fetch_task(protocol_schema.TaskRequestType.initial_prompt, None) - print(task) - finally: - - await api.close() - # session = aiohttp.ClientSession() - # try: - # resp = await session.post("http://localhost:8080/api/v1/tasks/", json={"type": "initial_prompt", "user": None}) - # resp.raise_for_status() - # print(await resp.text()) - # finally: - # await session.close() - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/discord-bot/bot/bot.py b/discord-bot/bot/bot.py index de8ceacf..e189b765 100644 --- a/discord-bot/bot/bot.py +++ b/discord-bot/bot/bot.py @@ -5,8 +5,8 @@ import hikari import lightbulb import miru -from bot.config import Config from bot.api_client import OasstApiClient +from bot.config import Config config = Config.from_env() diff --git a/discord-bot/bot/extensions/_example.py b/discord-bot/bot/extensions/_example.py index 330f5909..37783e43 100644 --- a/discord-bot/bot/extensions/_example.py +++ b/discord-bot/bot/extensions/_example.py @@ -2,7 +2,8 @@ # -*- coding: utf-8 -*- """Example plugin for reference. -Because this file starts with an `_`, it cannot be loaded by the bot. To see the example plugin in action, rename this file to `example.py`. +Because this file starts with an `_`, it cannot be loaded by the bot. +To see the example plugin in action, rename this file to `example.py`. """ import asyncio diff --git a/discord-bot/bot/extensions/tasks.py b/discord-bot/bot/extensions/tasks.py index dfe51160..71f47f52 100644 --- a/discord-bot/bot/extensions/tasks.py +++ b/discord-bot/bot/extensions/tasks.py @@ -6,13 +6,13 @@ import typing as t from datetime import datetime, timedelta import hikari - import lightbulb import lightbulb.decorators import miru -from bot.utils import format_time from oasst_shared.schemas.protocol import TaskRequestType +from bot.utils import format_time + plugin = lightbulb.Plugin("TaskPlugin") MAX_TASK_TIME = 60 * 60 @@ -42,9 +42,13 @@ async def task_thread(ctx: lightbulb.SlashContext): await ctx.respond(f"Please complete the task in the thread: {thread.mention}") # Send the task in the thread - # TODO: Request task from the backend await thread.send( - f"Please complete the task.\nSample Task\n\nSelf destruct {format_time(datetime.now() + timedelta(seconds=MAX_TASK_TIME), 'R')}" + f"""\ +Please complete the task. +Sample Task + +Self destruct {format_time(datetime.now() + timedelta(seconds=MAX_TASK_TIME), 'R')} +""" ) # Wait for the user to respond @@ -55,7 +59,6 @@ async def task_thread(ctx: lightbulb.SlashContext): predicate=lambda e: e.author.id == ctx.author.id and e.channel_id == thread.id, ) await ctx.respond(f"Received message: {event.message.content}") - # TODO: Send the message to the backend except asyncio.TimeoutError: await ctx.respond("You took too long to respond.") finally: @@ -75,16 +78,16 @@ async def task_thread(ctx: lightbulb.SlashContext): @lightbulb.implements(lightbulb.SlashCommand, lightbulb.PrefixCommand) async def task_dm(ctx: lightbulb.Context): """Request a task from the backend.""" - typ: str = ctx.options.type + await ctx.respond("Please complete the task in your DMs") - # Create a thread for the task - - await ctx.respond(f"Please complete the task in your DMs") - - # Send the task in the thread - # TODO: Request task from the backend + # Send the task in the dm await ctx.author.send( - f"Please complete the task.\nSample Task ({typ})\n\nSelf destruct {format_time(datetime.now() + timedelta(seconds=MAX_TASK_TIME), 'R')}" + f"""\ +Please complete the task. +Sample Task + +Self destruct {format_time(datetime.now() + timedelta(seconds=MAX_TASK_TIME), 'R')} +""" ) # Wait for the user to respond @@ -95,7 +98,6 @@ async def task_dm(ctx: lightbulb.Context): predicate=lambda e: e.author.id == ctx.author.id, ) await ctx.respond(f"Received message: {event.message.content}") - # TODO: Send the message to the backend except asyncio.TimeoutError: await ctx.respond("You took too long to respond.") @@ -113,7 +115,6 @@ class TaskModal(miru.Modal): async def callback(self, context: miru.ModalContext) -> None: await context.respond(f"Received response: {self.response.value}", flags=hikari.MessageFlag.EPHEMERAL) - # TODO: Send the message to the backend class ModalView(miru.View): @@ -146,7 +147,7 @@ async def task_modal(ctx: lightbulb.SlashContext): """Request a task from the backend.""" # typ: str = ctx.options.type view = ModalView( - modal_title=f"Assistant Response", + modal_title="Assistant Response", task="Please explain the moon landing to a six year old.", timeout=MAX_TASK_TIME, ) @@ -211,6 +212,8 @@ class RatingView(miru.View): class SelectRating(miru.View): + """View for rating a task with a select menu.""" + @miru.select( options=[ hikari.SelectMenuOption( @@ -249,7 +252,6 @@ class SelectRating(miru.View): @lightbulb.implements(lightbulb.SlashCommand) async def rating_task(ctx: lightbulb.SlashContext): """Rate stuff.""" - # Message Based rating await ctx.respond( "List the responses in order of best to worst response (1,2,3,4,5)", flags=hikari.MessageFlag.EPHEMERAL diff --git a/discord-bot/bot/extensions/work.py b/discord-bot/bot/extensions/work.py index e6ea3d7c..5c191481 100644 --- a/discord-bot/bot/extensions/work.py +++ b/discord-bot/bot/extensions/work.py @@ -6,13 +6,13 @@ import typing as t from datetime import datetime import hikari - import lightbulb import lightbulb.decorators import miru -from bot.api_client import OasstApiClient, TaskType 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 plugin = lightbulb.Plugin("WorkPlugin") @@ -52,7 +52,6 @@ async def _handle_task(ctx: lightbulb.SlashContext, task_type: TaskRequestType) If they select one, present the task steps until a `task_done` task is received. Finally, ask the user if they want to perform another task (of the same type). """ - oasst_api: OasstApiClient = ctx.bot.d.oasst_api # Continue to complete tasks until the user doesn't want to do another @@ -165,8 +164,8 @@ async def _send_task( ) -> tuple[t.Literal["accept", "next", "cancel"] | None, str]: """Send a task to the user. - Returns the user's choice and the message ID of the task message.""" - + Returns the user's choice and the message ID of the task message. + """ # 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. @@ -204,9 +203,7 @@ async def _send_task( 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?ixlib=rb-4.0.3&ixid=MnwxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8&auto=format&fit=crop&w=1073&q=80", - ) + .set_image("https://images.unsplash.com/photo-1455390582262-044cdead277a?w=512") .set_footer(text=f"OASST Assistant | {task.id}") ) @@ -215,12 +212,10 @@ def _rank_initial_prompt_embed(task: protocol_schema.RankInitialPromptsTask) -> embed = ( hikari.Embed( title="Rank Initial Prompt", - description=f"Rank the following tasks from best to worst (1,2,3,4,5)", + 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?ixlib=rb-4.0.3&ixid=MnwxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8&auto=format&fit=crop&w=1073&q=80", - ) + .set_image("https://images.unsplash.com/photo-1455390582262-044cdead277a?w=512") .set_footer(text=f"OASST Assistant | {task.id}") ) @@ -258,6 +253,11 @@ class TaskAcceptView(miru.View): class ChoiceView(miru.View): + """View with two buttons: yes and no. + + The view stops once one of the buttons is pressed and the choice is stored in the `choice` attribute. + """ + choice: bool | None = None @miru.button(label="Yes", custom_id="yes", style=hikari.ButtonStyle.SUCCESS) diff --git a/discord-bot/message_templates.py b/discord-bot/message_templates.py index df3ef1ac..dcb84c94 100644 --- a/discord-bot/message_templates.py +++ b/discord-bot/message_templates.py @@ -1,16 +1,20 @@ # -*- coding: utf-8 -*- +"""Message templates for the discord bot.""" import jinja2 +import typing from loguru import logger class MessageTemplates: - def __init__(self, template_dir="./templates"): - self.env = jinja2.Environment( + """Create message templates for the discord bot.""" + + def __init__(self, template_dir: str = "./templates"): + self.env = jinja2.Environment( # noqa: S701 loader=jinja2.FileSystemLoader(template_dir), autoescape=jinja2.select_autoescape(disabled_extensions=("msg",), default=False, default_for_string=False), ) - def render(self, template_name, **kwargs): + def render(self, template_name: str, **kwargs: typing.Any): template = self.env.get_template(template_name) txt = template.render(kwargs) logger.debug(txt) diff --git a/discord-bot/noxfile.py b/discord-bot/noxfile.py index 37226787..f85fc60c 100644 --- a/discord-bot/noxfile.py +++ b/discord-bot/noxfile.py @@ -27,6 +27,7 @@ def lint_code(session: Session): @nox.session(reuse_venv=True) def typecheck_code(session: Session): + """Typecheck the codebase.""" session.install("-r", "requirements.txt", "-U") session.install("pyright", "-U")