wassname/Open-Assistant
1
0
Fork 0
mirror of https://github.com/wassname/Open-Assistant.git synced 2026-07-10 00:20:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

clean up code

Browse Source
This commit is contained in:
AlexanderHOtt
2022-12-29 14:39:02 -08:00
parent 9fd2e76917
commit 221d3396f7
7 changed files with 51 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
discord-bot/bot/api_client.py
+9 -32
View File
@@ -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())
discord-bot/bot/bot.py
+1 -1
View File
@@ -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()
discord-bot/bot/extensions/_example.py
+2 -1
View File
@@ -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
discord-bot/bot/extensions/tasks.py
+19 -17
View File
@@ -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
discord-bot/bot/extensions/work.py
+12 -12
View File
@@ -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)
discord-bot/message_templates.py
+7 -3
View File
@@ -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)
discord-bot/noxfile.py
+1
View File
@@ -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")