Merge pull request LAION-AI#575 from LAION-AI/website-types

Centralize task types

website/src/components/Buttons/Skip.tsx

@@ -1,10 +1,6 @@
 import {
   Button,
   ButtonProps,
-  Menu,
-  MenuButton,
-  MenuItem,
-  MenuList,
   Modal,
   ModalBody,
   ModalCloseButton,
@@ -16,7 +12,6 @@ import {
   useDisclosure,
 } from "@chakra-ui/react";
 import { useState } from "react";
-import { FaChevronDown } from "react-icons/fa";
 
 interface SkipButtonProps extends ButtonProps {
   onSkip: (reason: string) => void;

website/src/components/Messages.tsx

@@ -1,15 +1,10 @@
 import { Grid } from "@chakra-ui/react";
 import { useColorMode } from "@chakra-ui/react";
 import { useMemo } from "react";
+import { Message } from "src/types/Conversation";
 
 import { FlaggableElement } from "./FlaggableElement";
 
-export interface Message {
-  text: string;
-  is_assistant: boolean;
-  message_id: string;
-}
-
 export interface ValidLabel {
   name: string;
   display_text: string;

website/src/components/Tasks/CreateTask.tsx

@@ -4,13 +4,13 @@ import { Messages } from "src/components/Messages";
 import { TaskControls } from "src/components/Survey/TaskControls";
 import { TrackedTextarea } from "src/components/Survey/TrackedTextarea";
 import { TwoColumnsWithCards } from "src/components/Survey/TwoColumnsWithCards";
-import { TaskType } from "./TaskTypes";
+import { TaskInfo } from "./TaskTypes";
 
 export interface CreateTaskProps {
   // we need a task type
   // eslint-disable-next-line  @typescript-eslint/no-explicit-any
   tasks: any[];
-  taskType: TaskType;
+  taskType: TaskInfo;
   trigger: (update: { id: string; update_type: string; content: { text: string } }) => void;
   onSkipTask: (task: { id: string }, reason: string) => void;
   onNextTask: () => void;

website/src/components/Tasks/TaskTypes.tsx

@@ -4,7 +4,7 @@ export enum TaskCategory {
   Label = "Label",
 }
 
-export interface TaskType {
+export interface TaskInfo {
   label: string;
   desc: string;
   category: TaskCategory;
@@ -14,7 +14,7 @@ export interface TaskType {
   instruction?: string;
 }
 
-export const TaskTypes: TaskType[] = [
+export const TaskTypes: TaskInfo[] = [
   // create
   {
     label: "Create Initial Prompts",

website/src/hooks/tasks/useCreateReply.ts

@@ -0,0 +1,8 @@
+import { TaskType } from "src/types/Task";
+import { CreateAssistantReplyTask, CreateInitialPromptTask, CreatePrompterReplyTask } from "src/types/Tasks";
+
+import { useGenericTaskAPI } from "./useGenericTaskAPI";
+
+export const useCreateAssistantReply = () => useGenericTaskAPI<CreateAssistantReplyTask>(TaskType.assistant_reply);
+export const useCreatePrompterReply = () => useGenericTaskAPI<CreatePrompterReplyTask>(TaskType.prompter_reply);
+export const useCreateInitialPrompt = () => useGenericTaskAPI<CreateInitialPromptTask>(TaskType.initial_prompt);

website/src/hooks/tasks/useGenericTaskAPI.tsx

@@ -1,18 +1,11 @@
 import { useState } from "react";
 import fetcher from "src/lib/fetcher";
 import poster from "src/lib/poster";
+import { BaseTask, TaskResponse } from "src/types/Task";
 import useSWRImmutable from "swr/immutable";
 import useSWRMutation from "swr/mutation";
 
-// TODO: type & centralize types for all tasks
-
-export interface TaskResponse<TaskType> {
-  id: string;
-  userId: string;
-  task: TaskType;
-}
-
-export const useGenericTaskAPI = <TaskType,>(taskApiEndpoint: string) => {
+export const useGenericTaskAPI = <TaskType extends BaseTask>(taskApiEndpoint: string) => {
   type ConcreteTaskResponse = TaskResponse<TaskType>;
 
   const [tasks, setTasks] = useState<ConcreteTaskResponse[]>([]);

website/src/hooks/tasks/useLabelingTask.ts

@@ -0,0 +1,32 @@
+import { BaseTask, TaskResponse, TaskType } from "src/types/Task";
+import { LabelAssistantReplyTask, LabelInitialPromptTask, LabelPrompterReplyTask } from "src/types/Tasks";
+
+import { useGenericTaskAPI } from "./useGenericTaskAPI";
+
+const useLabelingTask = <Task extends BaseTask>(
+  endpoint: TaskType.label_assistant_reply | TaskType.label_prompter_reply | TaskType.label_initial_prompt
+) => {
+  const { tasks, isLoading, trigger, reset, error } = useGenericTaskAPI<Task>(endpoint);
+
+  const submit = (id: string, message_id: string, text: string, validLabels: string[], labelWeights: number[]) => {
+    console.assert(validLabels.length === labelWeights.length);
+    const labels = Object.fromEntries(validLabels.map((label, i) => [label, labelWeights[i]]));
+
+    return trigger({ id, update_type: "text_labels", content: { labels, text, message_id } });
+  };
+
+  return { tasks, isLoading, submit, reset, error };
+};
+
+export type LabelAssistantReplyTaskResponse = TaskResponse<LabelAssistantReplyTask>;
+
+export const useLabelAssistantReplyTask = () =>
+  useLabelingTask<LabelAssistantReplyTask>(TaskType.label_assistant_reply);
+
+export type LabelInitialPromptTaskResponse = TaskResponse<LabelInitialPromptTask>;
+
+export const useLabelInitialPromptTask = () => useLabelingTask<LabelInitialPromptTask>(TaskType.label_initial_prompt);
+
+export type LabelPrompterReplyTaskResponse = TaskResponse<LabelPrompterReplyTask>;
+
+export const useLabelPrompterReplyTask = () => useLabelingTask<LabelPrompterReplyTask>(TaskType.label_prompter_reply);

website/src/hooks/tasks/useRankReplies.ts

@@ -0,0 +1,12 @@
+import { TaskType } from "src/types/Task";
+import { RankAssistantRepliesTask, RankInitialPromptsTask, RankPrompterRepliesTask } from "src/types/Tasks";
+
+import { useGenericTaskAPI } from "./useGenericTaskAPI";
+
+export const useRankAssistantRepliesTask = () =>
+  useGenericTaskAPI<RankAssistantRepliesTask>(TaskType.rank_assistant_replies);
+
+export const useRankPrompterRepliesTask = () =>
+  useGenericTaskAPI<RankPrompterRepliesTask>(TaskType.rank_prompter_replies);
+
+export const useRankInitialPromptsTask = () => useGenericTaskAPI<RankInitialPromptsTask>(TaskType.rank_initial_prompts);

website/src/pages/create/assistant_reply.tsx

@@ -3,7 +3,7 @@ import { useColorMode } from "@chakra-ui/react";
 import Head from "next/head";
 import { LoadingScreen } from "src/components/Loading/LoadingScreen";
 import { Task } from "src/components/Tasks/Task";
-import { useCreateAssistantReply } from "src/hooks/tasks/create/useCreateReply";
+import { useCreateAssistantReply } from "src/hooks/tasks/useCreateReply";
 
 const AssistantReply = () => {
   const { tasks, isLoading, reset, trigger } = useCreateAssistantReply();

website/src/pages/create/initial_prompt.tsx

@@ -3,7 +3,7 @@ import { useColorMode } from "@chakra-ui/react";
 import Head from "next/head";
 import { LoadingScreen } from "src/components/Loading/LoadingScreen";
 import { Task } from "src/components/Tasks/Task";
-import { useCreateInitialPrompt } from "src/hooks/tasks/create/useCreateInitialPrompt";
+import { useCreateInitialPrompt } from "src/hooks/tasks/useCreateReply";
 
 const InitialPrompt = () => {
   const { tasks, isLoading, reset, trigger } = useCreateInitialPrompt();

website/src/pages/create/user_reply.tsx

@@ -3,7 +3,7 @@ import Head from "next/head";
 import { Container } from "src/components/Container";
 import { LoadingScreen } from "src/components/Loading/LoadingScreen";
 import { Task } from "src/components/Tasks/Task";
-import { useCreatePrompterReply } from "src/hooks/tasks/create/useCreateReply";
+import { useCreatePrompterReply } from "src/hooks/tasks/useCreateReply";
 
 const UserReply = () => {
   const { tasks, isLoading, reset, trigger } = useCreatePrompterReply();

website/src/pages/evaluate/rank_assistant_replies.tsx

@@ -3,7 +3,7 @@ import Head from "next/head";
 import { Container } from "src/components/Container";
 import { LoadingScreen } from "src/components/Loading/LoadingScreen";
 import { Task } from "src/components/Tasks/Task";
-import { useRankAssistantRepliesTask } from "src/hooks/tasks/evaluate/useRankReplies";
+import { useRankAssistantRepliesTask } from "src/hooks/tasks/useRankReplies";
 
 const RankAssistantReplies = () => {
   const { tasks, isLoading, reset, trigger } = useRankAssistantRepliesTask();

website/src/pages/evaluate/rank_initial_prompts.tsx

@@ -3,7 +3,7 @@ import Head from "next/head";
 import { Container } from "src/components/Container";
 import { LoadingScreen } from "src/components/Loading/LoadingScreen";
 import { Task } from "src/components/Tasks/Task";
-import { useRankInitialPromptsTask } from "src/hooks/tasks/evaluate/useRankInitialPrompts";
+import { useRankInitialPromptsTask } from "src/hooks/tasks/useRankReplies";
 
 const RankInitialPrompts = () => {
   const { tasks, isLoading, reset, trigger } = useRankInitialPromptsTask();

website/src/pages/evaluate/rank_user_replies.tsx

@@ -3,7 +3,7 @@ import Head from "next/head";
 import { Container } from "src/components/Container";
 import { LoadingScreen } from "src/components/Loading/LoadingScreen";
 import { Task } from "src/components/Tasks/Task";
-import { useRankPrompterRepliesTask } from "src/hooks/tasks/evaluate/useRankReplies";
+import { useRankPrompterRepliesTask } from "src/hooks/tasks/useRankReplies";
 
 const RankUserReplies = () => {
   const { tasks, isLoading, reset, trigger } = useRankPrompterRepliesTask();

website/src/pages/label/label_assistant_reply.tsx

@@ -1,13 +1,10 @@
 import { useState } from "react";
 import { LoadingScreen } from "src/components/Loading/LoadingScreen";
-import { Message } from "src/components/Messages";
 import { MessageTable } from "src/components/Messages/MessageTable";
 import { TaskControls } from "src/components/Survey/TaskControls";
 import { LabelSliderGroup, LabelTask } from "src/components/Tasks/LabelTask";
-import {
-  LabelAssistantReplyTaskResponse,
-  useLabelAssistantReplyTask,
-} from "src/hooks/tasks/labeling/useLabelAssistantReply";
+import { LabelAssistantReplyTaskResponse, useLabelAssistantReplyTask } from "src/hooks/tasks/useLabelingTask";
+import { Message } from "src/types/Conversation";
 
 const LabelAssistantReply = () => {
   const [sliderValues, setSliderValues] = useState<number[]>([]);

website/src/pages/label/label_initial_prompt.tsx

@@ -3,10 +3,7 @@ import { LoadingScreen } from "src/components/Loading/LoadingScreen";
 import { MessageView } from "src/components/Messages";
 import { TaskControls } from "src/components/Survey/TaskControls";
 import { LabelSliderGroup, LabelTask } from "src/components/Tasks/LabelTask";
-import {
-  LabelInitialPromptTaskResponse,
-  useLabelInitialPromptTask,
-} from "src/hooks/tasks/labeling/useLabelInitialPrompt";
+import { LabelInitialPromptTaskResponse, useLabelInitialPromptTask } from "src/hooks/tasks/useLabelingTask";
 
 const LabelInitialPrompt = () => {
   const [sliderValues, setSliderValues] = useState<number[]>([]);

website/src/pages/label/label_prompter_reply.tsx

@@ -1,13 +1,10 @@
 import { useState } from "react";
 import { LoadingScreen } from "src/components/Loading/LoadingScreen";
-import { Message } from "src/components/Messages";
 import { MessageTable } from "src/components/Messages/MessageTable";
 import { TaskControls } from "src/components/Survey/TaskControls";
 import { LabelSliderGroup, LabelTask } from "src/components/Tasks/LabelTask";
-import {
-  LabelPrompterReplyTaskResponse,
-  useLabelPrompterReplyTask,
-} from "src/hooks/tasks/labeling/useLabelPrompterReply";
+import { LabelPrompterReplyTaskResponse, useLabelPrompterReplyTask } from "src/hooks/tasks/useLabelingTask";
+import { Message } from "src/types/Conversation";
 
 const LabelPrompterReply = () => {
   const [sliderValues, setSliderValues] = useState<number[]>([]);