Merge pull request #593 from LAION-AI/584-valid-labels-everywhere

Encapsulate valid label fetching to `FlaggableElement`

website/src/components/FlaggableElement.tsx

@@ -22,9 +22,11 @@ import {
   useId,
 } from "@chakra-ui/react";
 import { FlagIcon, QuestionMarkCircleIcon } from "@heroicons/react/20/solid";
-import { useState } from "react";
+import { useEffect, useState } from "react";
+import fetcher from "src/lib/fetcher";
 import poster from "src/lib/poster";
 import { colors } from "styles/Theme/colors";
+import useSWR from "swr";
 import useSWRMutation from "swr/mutation";
 
 interface textFlagLabels {
@@ -34,18 +36,28 @@ interface textFlagLabels {
 }
 
 export const FlaggableElement = (props) => {
+  const [labels, setLabels] = useState([]);
+  const [checkboxValues, setCheckboxValues] = useState([]);
+  const [sliderValues, setSliderValues] = useState([]);
+  const [isEditing, setIsEditing] = useBoolean();
   const backgroundColor = useColorModeValue("gray.200", "gray.700");
 
-  const [isEditing, setIsEditing] = useBoolean();
-  const flaggable_labels = props.flaggable_labels;
-  const TEXT_LABEL_FLAGS =
-    flaggable_labels?.valid_labels?.map((valid_label) => {
-      return {
-        attributeName: valid_label.name,
-        labelText: valid_label.display_text,
-        additionalExplanation: valid_label.help_text,
-      };
-    }) || [];
+  const { data, isLoading } = useSWR("/api/valid_labels", fetcher);
+  useEffect(() => {
+    if (isLoading) {
+      return;
+    }
+    const { valid_labels } = data;
+    const newLabels = valid_labels.map((valid_label) => ({
+      attributeName: valid_label.name,
+      labelText: valid_label.display_text,
+      additionalExplanation: valid_label.help_text,
+    }));
+    setSliderValues(new Array(newLabels.length).fill(1));
+    setCheckboxValues(new Array(newLabels.length).fill(false));
+    setLabels(newLabels);
+  }, [data, isLoading]);
+
   const { trigger } = useSWRMutation("/api/set_label", poster, {
     onSuccess: () => {
       setIsEditing.off();
@@ -54,7 +66,7 @@ export const FlaggableElement = (props) => {
 
   const submitResponse = () => {
     const label_map: Map<string, number> = new Map();
-    TEXT_LABEL_FLAGS.forEach((flag, i) => {
+    labels.forEach((flag, i) => {
       if (checkboxValues[i]) {
         label_map.set(flag.attributeName, sliderValues[i]);
       }
@@ -66,8 +78,6 @@ export const FlaggableElement = (props) => {
       text: props.text,
     });
   };
-  const [checkboxValues, setCheckboxValues] = useState(new Array(TEXT_LABEL_FLAGS.length).fill(false));
-  const [sliderValues, setSliderValues] = useState(new Array(TEXT_LABEL_FLAGS.length).fill(1));
 
   const handleCheckboxState = (isChecked, idx) => {
     setCheckboxValues(
@@ -112,7 +122,7 @@ export const FlaggableElement = (props) => {
           <PopoverCloseButton />
         </div>
         <PopoverBody>
-          {TEXT_LABEL_FLAGS.map((option, i) => (
+          {labels.map((option, i) => (
             <FlagCheckbox
               option={option}
               key={i}

website/src/components/Messages.tsx

@@ -2,29 +2,14 @@ import { Grid } from "@chakra-ui/react";
 import { forwardRef, useColorMode } from "@chakra-ui/react";
 import { useMemo } from "react";
 import { Message } from "src/types/Conversation";
-import { ValidLabel } from "src/types/Task";
 
 import { FlaggableElement } from "./FlaggableElement";
 
-export const Messages = ({
-  messages,
-  post_id,
-  valid_labels,
-}: {
-  messages: Message[];
-  post_id: string;
-  valid_labels: ValidLabel[];
-}) => {
+export const Messages = ({ messages, post_id }: { messages: Message[]; post_id: string }) => {
   const items = messages.map((messageProps: Message, i: number) => {
     const { message_id, text } = messageProps;
     return (
-      <FlaggableElement
-        text={text}
-        post_id={post_id}
-        message_id={message_id}
-        key={i + text}
-        flaggable_labels={valid_labels}
-      >
+      <FlaggableElement text={text} post_id={post_id} message_id={message_id} key={i + text}>
         <MessageView {...messageProps} />
       </FlaggableElement>
     );

website/src/components/Messages/MessageTable.tsx

@@ -1,11 +1,11 @@
 import { Stack, StackDivider } from "@chakra-ui/react";
 import { MessageTableEntry } from "src/components/Messages/MessageTableEntry";
 
-export function MessageTable({ messages, valid_labels }) {
+export function MessageTable({ messages }) {
   return (
     <Stack divider={<StackDivider />} spacing="4">
-      {messages.map((item, idx) => (
-        <MessageTableEntry item={item} idx={idx} key={item.message_id || item.id} valid_labels={valid_labels} />
+      {messages.map((item) => (
+        <MessageTableEntry item={item} key={item.message_id || item.id} />
       ))}
     </Stack>
   );

website/src/components/Messages/MessageTableEntry.tsx

@@ -1,25 +1,23 @@
 import { Avatar, HStack, LinkBox, LinkOverlay, useColorModeValue } from "@chakra-ui/react";
 import { boolean } from "boolean";
 import { FlaggableElement } from "src/components/FlaggableElement";
-import type { ValidLabel } from "src/types/Task";
 
 interface Message {
   text: string;
   id: string;
+  message_id: string;
   is_assistant: boolean;
 }
 interface MessageTableEntryProps {
   item: Message;
-  idx: number;
-  valid_labels: ValidLabel[];
 }
 export function MessageTableEntry(props: MessageTableEntryProps) {
-  const { item, valid_labels } = props;
+  const { item } = props;
   const backgroundColor = useColorModeValue("gray.50", "gray.800");
 
   return (
     <div>
-      <FlaggableElement text={item.text} post_id={item.id} key={`flag_${item.id}`} flaggable_labels={valid_labels}>
+      <FlaggableElement text={item.text} message_id={item.message_id} post_id={item.id} key={`flag_${item.id}`}>
         <HStack>
           <Avatar
             size="sm"

website/src/components/Messages/MessageWithChildren.tsx

@@ -68,7 +68,7 @@ export function MessageWithChildren(props: MessageWithChildrenProps) {
               {isFirst ? "Message" : depth === 1 ? "Children" : "Ancestor"}
             </Text>
             <Box width="fit-content" bg={backgroundColor} padding="4" borderRadius="xl" boxShadow="base">
-              <MessageTableEntry item={message} idx={1} valid_labels={[]} />
+              <MessageTableEntry item={message} />
             </Box>
           </Box>
         </>
@@ -95,7 +95,7 @@ export function MessageWithChildren(props: MessageWithChildrenProps) {
             <Stack {...MessageStackProps}>
               {children.map((item, idx) => (
                 <Box flex="1" key={`recursiveMessageWChildren_${idx}`}>
-                  <MessageTableEntry item={item} idx={idx * 2} valid_labels={[]} />
+                  <MessageTableEntry item={item} />
                 </Box>
               ))}
             </Stack>

website/src/components/Tasks/CreateTask.tsx

@@ -7,8 +7,6 @@ import { TaskSurveyProps } from "src/components/Tasks/Task";
 export const CreateTask = ({ task, taskType, onReplyChanged }: TaskSurveyProps<{ text: string }>) => {
   const [inputText, setInputText] = useState("");
 
-  const valid_labels = task.valid_labels;
-
   const textChangeHandler = (event: React.ChangeEvent<HTMLTextAreaElement>) => {
     const text = event.target.value;
     onReplyChanged({ content: { text }, state: "VALID" });
@@ -20,9 +18,7 @@ export const CreateTask = ({ task, taskType, onReplyChanged }: TaskSurveyProps<{
       <>
         <h5 className="text-lg font-semibold">{taskType.label}</h5>
         <p className="text-lg py-1">{taskType.overview}</p>
-        {task.conversation ? (
-          <Messages messages={task.conversation.messages} post_id={task.id} valid_labels={valid_labels} />
-        ) : null}
+        {task.conversation ? <Messages messages={task.conversation.messages} post_id={task.id} /> : null}
       </>
       <>
         <h5 className="text-lg font-semibold">{taskType.instruction}</h5>

website/src/components/Tasks/EvaluateTask.tsx

@@ -24,7 +24,6 @@ export const EvaluateTask = ({ task, onReplyChanged }: TaskSurveyProps<{ ranking
     onReplyChanged({ content: { ranking: newRanking }, state: "VALID" });
   };
 
-  const valid_labels = task.valid_labels;
   const sortables = task.replies ? "replies" : "prompts";
 
   return (
@@ -33,7 +32,7 @@ export const EvaluateTask = ({ task, onReplyChanged }: TaskSurveyProps<{ ranking
       <p className="text-lg py-1">
         Given the following {sortables}, sort them from best to worst, best being first, worst being last.
       </p>
-      <MessageTable messages={messages} valid_labels={valid_labels} />
+      <MessageTable messages={messages} />
       <Sortable items={task[sortables]} onChange={onRank} className="my-8" />
     </SurveyCard>
   );

website/src/components/Tasks/LabelTask.tsx

@@ -44,7 +44,6 @@ export const LabelTask = ({
                 message_id: task.message_id,
               },
             ]}
-            valid_labels={valid_labels}
           />
         ) : (
           <MessageView text={task.prompt} is_assistant={false} message_id={task.message_id} />

website/src/components/Tasks/Task.tsx

@@ -21,7 +21,7 @@ export interface TaskSurveyProps<T> {
   onReplyChanged: (state: TaskReplyState<T>) => void;
 }
 
-export const Task = ({ task, trigger, mutate }) => {
+export const Task = ({ frontendId, task, trigger, mutate }) => {
   const [taskStatus, setTaskStatus] = useState<TaskStatus>("NOT_SUBMITTABLE");
   const replyContent = useRef<TaskContent>(null);
   const [showUnchangedWarning, setShowUnchangedWarning] = useState(false);
@@ -39,7 +39,7 @@ export const Task = ({ task, trigger, mutate }) => {
 
   const rejectTask = (reason: string) => {
     sendRejection({
-      id: task.id,
+      id: frontendId,
       reason,
     });
   };
@@ -66,7 +66,7 @@ export const Task = ({ task, trigger, mutate }) => {
         break;
       case "SUBMITABLE": {
         trigger({
-          id: task.id,
+          id: frontendId,
           update_type: taskType.update_type,
           content: replyContent.current,
         });
@@ -107,7 +107,7 @@ export const Task = ({ task, trigger, mutate }) => {
         onSubmitAnyway={() => {
           if (taskStatus === "DEFAULT") {
             trigger({
-              id: task.id,
+              id: frontendId,
               update_type: taskType.update_type,
               content: replyContent.current,
             });

website/src/pages/api/new_task/[task_type].ts

@@ -14,7 +14,6 @@ const handler = withoutRole("banned", async (req, res, token) => {
   // Fetch the new task.
   const { task_type } = req.query;
   const task = await oasstApiClient.fetchTask(task_type as string, token);
-  const valid_labels = await oasstApiClient.fetch_valid_text();
 
   // Store the task and link it to the user..
   const registeredTask = await prisma.registeredTask.create({
@@ -28,9 +27,6 @@ const handler = withoutRole("banned", async (req, res, token) => {
     },
   });
 
-  // Add the valid labels that can be used to flag messages in this Task
-  registeredTask["valid_labels"] = valid_labels;
-
   // Send the results to the client.
   res.status(200).json(registeredTask);
 });

website/src/pages/api/set_label.ts

@@ -6,7 +6,7 @@ import { withoutRole } from "src/lib/auth";
  */
 const handler = withoutRole("banned", async (req, res, token) => {
   // Parse out the local message_id, task ID and the interaction contents.
-  const { message_id, label_map, text } = await JSON.parse(req.body);
+  const { message_id, post_id, label_map, text } = await JSON.parse(req.body);
 
   const interactionRes = await fetch(`${process.env.FASTAPI_URL}/api/v1/text_labels`, {
     method: "POST",
@@ -16,7 +16,7 @@ const handler = withoutRole("banned", async (req, res, token) => {
     },
     body: JSON.stringify({
       type: "text_labels",
-      message_id: message_id,
+      message_id: message_id || post_id,
       labels: label_map,
       text: text,
       user: {
@@ -26,6 +26,10 @@ const handler = withoutRole("banned", async (req, res, token) => {
       },
     }),
   });
+  if (interactionRes.status !== 204) {
+    const r = await interactionRes.json();
+    console.error(JSON.stringify(r));
+  }
   res.status(interactionRes.status).end();
 });
 

website/src/pages/api/valid_labels.ts

@@ -0,0 +1,12 @@
+import { withoutRole } from "src/lib/auth";
+import { oasstApiClient } from "src/lib/oasst_api_client";
+
+/**
+ * Returns the set of valid labels that can be applied to messages.
+ */
+const handler = withoutRole("banned", async (req, res) => {
+  const valid_labels = await oasstApiClient.fetch_valid_text();
+  res.status(200).json(valid_labels);
+});
+
+export default handler;

website/src/pages/create/assistant_reply.tsx

@@ -21,7 +21,7 @@ const AssistantReply = () => {
         <title>Reply as Assistant</title>
         <meta name="description" content="Reply as Assistant." />
       </Head>
-      <Task key={tasks[0].task.id} task={tasks[0].task} trigger={trigger} mutate={reset} />
+      <Task key={tasks[0].task.id} frontendId={tasks[0].id} task={tasks[0].task} trigger={trigger} mutate={reset} />
     </>
   );
 };

website/src/pages/create/initial_prompt.tsx

@@ -21,7 +21,7 @@ const InitialPrompt = () => {
         <title>Reply as Assistant</title>
         <meta name="description" content="Reply as Assistant." />
       </Head>
-      <Task key={tasks[0].task.id} task={tasks[0].task} trigger={trigger} mutate={reset} />
+      <Task key={tasks[0].task.id} frontendId={tasks[0].id} task={tasks[0].task} trigger={trigger} mutate={reset} />
     </>
   );
 };

website/src/pages/create/summarize_story.tsx

@@ -39,7 +39,7 @@ const SummarizeStory = () => {
   const submitResponse = () => {
     const text = inputText.trim();
     trigger({
-      id: tasks[0].task.id,
+      id: tasks[0].id,
       update_type: "text_reply_to_message",
       content: {
         text,

website/src/pages/create/user_reply.tsx

@@ -21,7 +21,7 @@ const UserReply = () => {
         <title>Reply as Assistant</title>
         <meta name="description" content="Reply as Assistant." />
       </Head>
-      <Task key={tasks[0].task.id} task={tasks[0].task} trigger={trigger} mutate={reset} />
+      <Task key={tasks[0].task.id} frontendId={tasks[0].id} task={tasks[0].task} trigger={trigger} mutate={reset} />
     </>
   );
 };

website/src/pages/evaluate/rank_assistant_replies.tsx

@@ -21,7 +21,7 @@ const RankAssistantReplies = () => {
         <title>Rank Assistant Replies</title>
         <meta name="description" content="Rank Assistant Replies." />
       </Head>
-      <Task key={tasks[0].task.id} task={tasks[0].task} trigger={trigger} mutate={reset} />
+      <Task key={tasks[0].task.id} frontendId={tasks[0].id} task={tasks[0].task} trigger={trigger} mutate={reset} />
     </>
   );
 };

website/src/pages/evaluate/rank_initial_prompts.tsx

@@ -21,7 +21,7 @@ const RankInitialPrompts = () => {
         <title>Rank Initial Prompts</title>
         <meta name="description" content="Rank initial prompts." />
       </Head>
-      <Task key={tasks[0].task.id} task={tasks[0].task} trigger={trigger} mutate={reset} />
+      <Task key={tasks[0].task.id} frontendId={tasks[0].id} task={tasks[0].task} trigger={trigger} mutate={reset} />
     </>
   );
 };

website/src/pages/evaluate/rank_user_replies.tsx

@@ -21,7 +21,7 @@ const RankUserReplies = () => {
         <title>Rank User Replies</title>
         <meta name="description" content="Rank User Replies." />
       </Head>
-      <Task key={tasks[0].task.id} task={tasks[0].task} trigger={trigger} mutate={reset} />
+      <Task key={tasks[0].task.id} frontendId={tasks[0].id} task={tasks[0].task} trigger={trigger} mutate={reset} />
     </>
   );
 };

website/src/pages/label/label_assistant_reply.tsx

@@ -21,7 +21,7 @@ const LabelAssistantReply = () => {
         <title>Label Assistant Reply</title>
         <meta name="description" content="Label Assistant Reply" />
       </Head>
-      <Task key={tasks[0].task.id} task={tasks[0].task} trigger={trigger} mutate={reset} />
+      <Task key={tasks[0].task.id} frontendId={tasks[0].id} task={tasks[0].task} trigger={trigger} mutate={reset} />
     </>
   );
 };

website/src/pages/label/label_initial_prompt.tsx

@@ -21,7 +21,7 @@ const LabelInitialPrompt = () => {
         <title>Label Initial Prompt</title>
         <meta name="description" content="Label Initial Prompt" />
       </Head>
-      <Task key={tasks[0].task.id} task={tasks[0].task} trigger={trigger} mutate={reset} />
+      <Task key={tasks[0].task.id} frontendId={tasks[0].id} task={tasks[0].task} trigger={trigger} mutate={reset} />
     </>
   );
 };

website/src/pages/label/label_prompter_reply.tsx

@@ -21,7 +21,7 @@ const LabelPrompterReply = () => {
         <title>Label Prompter Reply</title>
         <meta name="description" content="Label Prompter Reply" />
       </Head>
-      <Task key={tasks[0].task.id} task={tasks[0].task} trigger={trigger} mutate={reset} />
+      <Task key={tasks[0].task.id} frontendId={tasks[0].id} task={tasks[0].task} trigger={trigger} mutate={reset} />
     </>
   );
 };

website/src/pages/messages/[id]/index.tsx

@@ -48,7 +48,7 @@ const MessageDetail = ({ id }) => {
                 </Text>
                 <Box bg={backgroundColor} padding="4" borderRadius="xl" boxShadow="base" width="fit-content">
                   {" "}
-                  <MessageTableEntry item={parent} idx={1} valid_labels={[]} />
+                  <MessageTableEntry item={parent} />
                 </Box>
               </Box>
             </>

website/src/pages/messages/index.tsx

@@ -52,11 +52,7 @@ const MessagesDashboard = () => {
             borderRadius="xl"
             className="p-6 shadow-sm"
           >
-            {receivedMessages ? (
-              <MessageTable messages={messages} valid_labels={[]} />
-            ) : (
-              <CircularProgress isIndeterminate />
-            )}
+            {receivedMessages ? <MessageTable messages={messages} /> : <CircularProgress isIndeterminate />}
           </Box>
         </Box>
         <Box>
@@ -70,11 +66,7 @@ const MessagesDashboard = () => {
             borderRadius="xl"
             className="p-6 shadow-sm"
           >
-            {receivedUserMessages ? (
-              <MessageTable messages={userMessages} valid_labels={[]} />
-            ) : (
-              <CircularProgress isIndeterminate />
-            )}
+            {receivedUserMessages ? <MessageTable messages={userMessages} /> : <CircularProgress isIndeterminate />}
           </Box>
         </Box>
       </SimpleGrid>

website/src/types/Task.ts

@@ -31,5 +31,4 @@ export interface TaskResponse<Task extends BaseTask> {
   id: string;
   userId: string;
   task: Task;
-  valid_labels: ValidLabel[];
 }