From 54cc88bb1f8f7d1bc401e69f30cfc27737a4226d Mon Sep 17 00:00:00 2001
From: Keith Stevens <fozziethebeat@gmail.com>
Date: Tue, 10 Jan 2023 14:08:43 +0900
Subject: [PATCH] Ensure FlaggableElement always has the set of labels

---
 website/cypress/e2e/evaluate/rank_assistant_replies.cy.ts | 2 +-
 website/cypress/e2e/evaluate/rank_user_replies.cy.ts      | 2 +-
 website/src/components/Messages/MessageTable.tsx          | 5 +++--
 website/src/components/Messages/MessageTableEntry.tsx     | 6 ++++--
 website/src/components/Tasks/EvaluateTask.tsx             | 3 ++-
 website/src/pages/label/label_assistant_reply.tsx         | 3 ++-
 website/src/pages/label/label_prompter_reply.tsx          | 3 ++-
 7 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/website/cypress/e2e/evaluate/rank_assistant_replies.cy.ts b/website/cypress/e2e/evaluate/rank_assistant_replies.cy.ts
index c7b85695..3093bd56 100644
--- a/website/cypress/e2e/evaluate/rank_assistant_replies.cy.ts
+++ b/website/cypress/e2e/evaluate/rank_assistant_replies.cy.ts
@@ -1,7 +1,7 @@
 describe("ranking prompter replies", () => {
   it("completes the current task on submit and on request shows a new task", () => {
     cy.signInWithEmail("cypress@example.com");
-    cy.visit("/evaluate/rank_user_replies");
+    cy.visit("/evaluate/rank_assistant_replies");
 
     cy.get('[data-cy="task-id"').then((taskIdElement) => {
       const taskId = taskIdElement.text();
diff --git a/website/cypress/e2e/evaluate/rank_user_replies.cy.ts b/website/cypress/e2e/evaluate/rank_user_replies.cy.ts
index c448a4c7..55487f1d 100644
--- a/website/cypress/e2e/evaluate/rank_user_replies.cy.ts
+++ b/website/cypress/e2e/evaluate/rank_user_replies.cy.ts
@@ -1,7 +1,7 @@
 describe("ranking assistant replies", () => {
   it("completes the current task on submit and on request shows a new task", () => {
     cy.signInWithEmail("cypress@example.com");
-    cy.visit("/evaluate/rank_assistant_replies");
+    cy.visit("/evaluate/rank_user_replies");
 
     cy.get('[data-cy="task-id"').then((taskIdElement) => {
       const taskId = taskIdElement.text();
diff --git a/website/src/components/Messages/MessageTable.tsx b/website/src/components/Messages/MessageTable.tsx
index 95ccc540..33ffc4a4 100644
--- a/website/src/components/Messages/MessageTable.tsx
+++ b/website/src/components/Messages/MessageTable.tsx
@@ -1,11 +1,12 @@
 import { Stack, StackDivider } from "@chakra-ui/react";
 import { MessageTableEntry } from "src/components/Messages/MessageTableEntry";
 
-export function MessageTable({ messages }) {
+export function MessageTable({ messages, valid_labels }) {
+  console.log(messages);
   return (
     <Stack divider={<StackDivider />} spacing="4">
       {messages.map((item, idx) => (
-        <MessageTableEntry item={item} idx={idx} key={item.id} />
+        <MessageTableEntry item={item} idx={idx} key={item.message_id} valid_labels={valid_labels} />
       ))}
     </Stack>
   );
diff --git a/website/src/components/Messages/MessageTableEntry.tsx b/website/src/components/Messages/MessageTableEntry.tsx
index 9fad7262..0f58efad 100644
--- a/website/src/components/Messages/MessageTableEntry.tsx
+++ b/website/src/components/Messages/MessageTableEntry.tsx
@@ -2,6 +2,7 @@ import { Avatar, HStack, LinkBox, useColorModeValue } from "@chakra-ui/react";
 import { boolean } from "boolean";
 import NextLink from "next/link";
 import { FlaggableElement } from "src/components/FlaggableElement";
+import type { ValidLabel } from "src/components/Messages";
 
 interface Message {
   text: string;
@@ -11,13 +12,14 @@ interface Message {
 interface MessageTableEntryProps {
   item: Message;
   idx: number;
+  valid_labels: ValidLabel[];
 }
 export function MessageTableEntry(props: MessageTableEntryProps) {
-  const { item, idx } = props;
+  const { item, idx, valid_labels } = props;
   const bgColor = useColorModeValue(idx % 2 === 0 ? "bg-slate-800" : "bg-black", "bg-sky-900");
 
   return (
-    <FlaggableElement text={item.text} post_id={item.id} key={`flag_${item.id}`}>
+    <FlaggableElement text={item.text} post_id={item.id} key={`flag_${item.id}`} flaggable_labels={valid_labels}>
       <HStack>
         <Avatar
           name={`${boolean(item.is_assistant) ? "Assitant" : "User"}`}
diff --git a/website/src/components/Tasks/EvaluateTask.tsx b/website/src/components/Tasks/EvaluateTask.tsx
index d0a1f404..61ed3889 100644
--- a/website/src/components/Tasks/EvaluateTask.tsx
+++ b/website/src/components/Tasks/EvaluateTask.tsx
@@ -33,6 +33,7 @@ export const EvaluateTask = ({ tasks, trigger, onSkipTask, onNextTask, mainBgCla
     messages = messages.map((message, index) => ({ ...message, id: index }));
   }
 
+  const valid_labels = tasks[0].valid_labels;
   const sortables = tasks[0].task.replies ? "replies" : "prompts";
 
   return (
@@ -42,7 +43,7 @@ export const EvaluateTask = ({ tasks, trigger, onSkipTask, onNextTask, mainBgCla
         <p className="text-lg py-1">
           Given the following {sortables}, sort them from best to worst, best being first, worst being last.
         </p>
-        {messages ? <MessageTable messages={messages} /> : null}
+        {messages ? <MessageTable messages={messages} valid_labels={valid_labels} /> : null}
         <Sortable items={tasks[0].task[sortables]} onChange={setRanking} className="my-8" />
       </SurveyCard>
 
diff --git a/website/src/pages/label/label_assistant_reply.tsx b/website/src/pages/label/label_assistant_reply.tsx
index a0f961f7..89b612ca 100644
--- a/website/src/pages/label/label_assistant_reply.tsx
+++ b/website/src/pages/label/label_assistant_reply.tsx
@@ -19,6 +19,7 @@ const LabelAssistantReply = () => {
   }
 
   const task = tasks[0].task;
+  const valid_labels = tasks[0].valid_labels;
   const messages: Message[] = [
     ...task.conversation.messages,
     { text: task.reply, is_assistant: true, message_id: task.message_id },
@@ -28,7 +29,7 @@ const LabelAssistantReply = () => {
     <LabelTask
       title="Label Assistant Reply"
       desc="Given the following discussion, provide labels for the final prompt"
-      messages={<MessageTable messages={messages} />}
+      messages={<MessageTable messages={messages} valid_labels={valid_labels} />}
       inputs={<LabelSliderGroup labelIDs={task.valid_labels} onChange={setSliderValues} />}
       controls={
         <TaskControls
diff --git a/website/src/pages/label/label_prompter_reply.tsx b/website/src/pages/label/label_prompter_reply.tsx
index 2fd3d76a..812a3fcc 100644
--- a/website/src/pages/label/label_prompter_reply.tsx
+++ b/website/src/pages/label/label_prompter_reply.tsx
@@ -19,6 +19,7 @@ const LabelPrompterReply = () => {
   }
 
   const task = tasks[0].task;
+  const valid_labels = tasks[0].valid_labels;
   const messages: Message[] = [
     ...task.conversation.messages,
     { text: task.reply, is_assistant: false, message_id: task.message_id },
@@ -28,7 +29,7 @@ const LabelPrompterReply = () => {
     <LabelTask
       title="Label Prompter Reply"
       desc="Given the following discussion, provide labels for the final prompt"
-      messages={<MessageTable messages={messages} />}
+      messages={<MessageTable messages={messages} valid_labels={valid_labels} />}
       inputs={<LabelSliderGroup labelIDs={task.valid_labels} onChange={setSliderValues} />}
       controls={
         <TaskControls