fix choices in playground

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/ChoiceView.tsx

@@ -7,12 +7,14 @@ type ChoiceViewProps = {
   choice: Choice;
   isStructuredOutput?: boolean;
   isNested?: boolean;
+  choiceIndex?: number;
 };
 
 export const ChoiceView = ({
   choice,
   isStructuredOutput,
   isNested,
+  choiceIndex,
 }: ChoiceViewProps) => {
   const {message} = choice;
   return (
@@ -21,7 +23,7 @@ export const ChoiceView = ({
       message={message}
       isStructuredOutput={isStructuredOutput}
       isNested={isNested}
-      isChoice
+      choiceIndex={choiceIndex}
     />
   );
 };

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/ChoicesDrawer.tsx

@@ -93,6 +93,7 @@ export const ChoicesDrawer = ({
               <ChoiceView
                 choice={c}
                 isStructuredOutput={isStructuredOutput}
+                choiceIndex={index}
                 isNested
               />
             </div>

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/ChoicesView.tsx

@@ -1,5 +1,6 @@
 import React, {useState} from 'react';
 
+import {usePlaygroundContext} from '../PlaygroundPage/PlaygroundContext';
 import {ChoicesDrawer} from './ChoicesDrawer';
 import {ChoicesViewCarousel} from './ChoicesViewCarousel';
 import {ChoiceView} from './ChoiceView';
@@ -14,11 +15,15 @@ export const ChoicesView = ({
   choices,
   isStructuredOutput,
 }: ChoicesViewProps) => {
+  const {setSelectedChoiceIndex: setGlobalSelectedChoiceIndex} =
+    usePlaygroundContext();
+
   const [isDrawerOpen, setIsDrawerOpen] = useState(false);
   const [localSelectedChoiceIndex, setLocalSelectedChoiceIndex] = useState(0);
 
   const handleSetSelectedChoiceIndex = (choiceIndex: number) => {
     setLocalSelectedChoiceIndex(choiceIndex);
+    setGlobalSelectedChoiceIndex(choiceIndex);
   };
 
   if (choices.length === 0) {

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/ChoicesViewCarousel.tsx

@@ -33,6 +33,7 @@ export const ChoicesViewCarousel = ({
       <ChoiceView
         choice={choices[selectedChoiceIndex]}
         isStructuredOutput={isStructuredOutput}
+        choiceIndex={selectedChoiceIndex}
       />
       <div className="flex items-center">
         <div className="flex-auto">

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/MessagePanel.tsx

@@ -15,7 +15,7 @@ type MessagePanelProps = {
   index: number;
   message: Message;
   isStructuredOutput?: boolean;
-  isChoice?: boolean;
+  choiceIndex?: number;
   isNested?: boolean;
   pendingToolResponseId?: string;
 };
@@ -24,7 +24,7 @@ export const MessagePanel = ({
   index,
   message,
   isStructuredOutput,
-  isChoice,
+  choiceIndex,
   isNested,
   // The id of the tool call response that is pending
   // If the tool call response is pending, the editor will be shown automatically
@@ -120,7 +120,7 @@ export const MessagePanel = ({
                 <PlaygroundMessagePanelEditor
                   message={message}
                   index={index}
-                  isChoice={isChoice ?? false}
+                  choiceIndex={choiceIndex}
                   editorHeight={editorHeight}
                   isNested={isNested ?? false}
                   pendingToolResponseId={pendingToolResponseId}
@@ -173,7 +173,7 @@ export const MessagePanel = ({
         <div className="flex w-full items-center justify-start opacity-0 group-hover:opacity-100">
           <PlaygroundMessagePanelButtons
             index={message.original_index ?? index}
-            isChoice={isChoice ?? false}
+            choiceIndex={choiceIndex}
             isTool={isTool}
             hasContent={hasContent}
             contentRef={contentRef}

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/PlaygroundMessagePanelButtons.tsx

@@ -5,7 +5,7 @@ import {usePlaygroundContext} from '../PlaygroundPage/PlaygroundContext';
 
 type PlaygroundMessagePanelButtonsProps = {
   index: number;
-  isChoice: boolean;
+  choiceIndex?: number;
   isTool: boolean;
   hasContent: boolean;
   contentRef: React.RefObject<HTMLDivElement>;
@@ -17,7 +17,7 @@ export const PlaygroundMessagePanelButtons: React.FC<
   PlaygroundMessagePanelButtonsProps
 > = ({
   index,
-  isChoice,
+  choiceIndex,
   isTool,
   hasContent,
   contentRef,
@@ -32,7 +32,7 @@ export const PlaygroundMessagePanelButtons: React.FC<
         variant="quiet"
         size="small"
         startIcon="randomize-reset-reload"
-        onClick={() => retry?.(index, isChoice)}
+        onClick={() => retry?.(index, choiceIndex)}
         tooltip={
           !hasContent
             ? 'We currently do not support retrying functions'
@@ -64,8 +64,8 @@ export const PlaygroundMessagePanelButtons: React.FC<
         size="small"
         startIcon="delete"
         onClick={() => {
-          if (isChoice) {
-            deleteChoice?.(index);
+          if (choiceIndex !== undefined) {
+            deleteChoice?.(index, choiceIndex);
           } else {
             deleteMessage?.(index, responseIndexes);
           }

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/PlaygroundMessagePanelEditor.tsx

@@ -13,15 +13,15 @@ type PlaygroundMessagePanelEditorProps = {
   pendingToolResponseId?: string;
   message: Message;
   index: number;
-  isChoice: boolean;
+  choiceIndex?: number;
   setEditorHeight: (height: number | null) => void;
 };
 
 export const PlaygroundMessagePanelEditor: React.FC<
   PlaygroundMessagePanelEditorProps
 > = ({
   index,
-  isChoice,
+  choiceIndex,
   setEditorHeight,
   editorHeight,
   isNested,
@@ -45,10 +45,10 @@ export const PlaygroundMessagePanelEditor: React.FC<
   }, [initialContent]);
 
   const handleSave = () => {
-    if (isChoice) {
-      editChoice?.(index, {
+    if (choiceIndex !== undefined) {
+      editChoice?.(choiceIndex, {
+        ...message,
         content: editedContent,
-        role: message.role,
       });
     } else {
       editMessage?.(index, {

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/PlaygroundChat.tsx

@@ -161,15 +161,21 @@ export const PlaygroundChat = ({
                           addMessage: newMessage => addMessage(idx, newMessage),
                           editChoice: (choiceIndex, newChoice) =>
                             editChoice(idx, choiceIndex, newChoice),
-                          retry: (messageIndex: number, isChoice?: boolean) =>
-                            handleRetry(idx, messageIndex, isChoice),
+                          retry: (messageIndex: number, choiceIndex?: number) =>
+                            handleRetry(idx, messageIndex, choiceIndex),
                           sendMessage: (
                             role: PlaygroundMessageRole,
                             content: string,
                             toolCallId?: string
                           ) => {
                             handleSend(role, idx, content, toolCallId);
                           },
+                          setSelectedChoiceIndex: (choiceIndex: number) =>
+                            setPlaygroundStateField(
+                              idx,
+                              'selectedChoiceIndex',
+                              choiceIndex
+                            ),
                         }}>
                         <CallChat call={state.traceCall as TraceCallSchema} />
                       </PlaygroundContext.Provider>

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/useChatCompletionFunctions.tsx

@@ -70,7 +70,7 @@ export const useChatCompletionFunctions = (
         if (callIndex !== undefined && callIndex !== index) {
           return state;
         }
-        const updatedState = appendChoicesToMessages(state);
+        const updatedState = appendChoiceToMessages(state);
         if (updatedState.traceCall?.inputs?.messages) {
           updatedState.traceCall.inputs.messages.push(newMessage);
         }
@@ -99,14 +99,14 @@ export const useChatCompletionFunctions = (
   const handleRetry = async (
     callIndex: number,
     messageIndex: number,
-    isChoice?: boolean
+    choiceIndex?: number
   ) => {
     try {
       setIsLoading(true);
       const updatedStates = playgroundStates.map((state, index) => {
         if (index === callIndex) {
-          if (isChoice) {
-            return appendChoicesToMessages(state);
+          if (choiceIndex !== undefined) {
+            return appendChoiceToMessages(state, choiceIndex);
           }
           const updatedState = JSON.parse(JSON.stringify(state));
           if (updatedState.traceCall?.inputs?.messages) {
@@ -203,17 +203,25 @@ const handleUpdateCallWithResponse = (
   };
 };
 
-const appendChoicesToMessages = (state: PlaygroundState): PlaygroundState => {
+const appendChoiceToMessages = (
+  state: PlaygroundState,
+  choiceIndex?: number
+): PlaygroundState => {
   const updatedState = JSON.parse(JSON.stringify(state));
   if (
     updatedState.traceCall?.inputs?.messages &&
     updatedState.traceCall.output?.choices
   ) {
-    updatedState.traceCall.output.choices.forEach((choice: any) => {
-      if (choice.message) {
-        updatedState.traceCall.inputs.messages.push(choice.message);
-      }
-    });
+    if (choiceIndex !== undefined) {
+      updatedState.traceCall.inputs.messages.push(
+        updatedState.traceCall.output.choices[choiceIndex].message
+      );
+    } else {
+      updatedState.traceCall.inputs.messages.push(
+        updatedState.traceCall.output.choices[updatedState.selectedChoiceIndex]
+          .message
+      );
+    }
     updatedState.traceCall.output.choices = undefined;
   }
   return updatedState;

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/useChatFunctions.tsx

@@ -114,16 +114,9 @@ export const useChatFunctions = (
         newTraceCall?.output &&
         Array.isArray((newTraceCall.output as TraceCallOutput).choices)
       ) {
-        // Delete the old choice
-        (newTraceCall.output as TraceCallOutput).choices!.splice(
-          choiceIndex,
-          1
-        );
-
-        // Add the new choice as a message
-        newTraceCall.inputs = newTraceCall.inputs ?? {};
-        newTraceCall.inputs.messages = newTraceCall.inputs.messages ?? [];
-        newTraceCall.inputs.messages.push(newChoice);
+        // Replace the choice
+        (newTraceCall.output as TraceCallOutput).choices![choiceIndex].message =
+          newChoice;
       }
       return newTraceCall;
     });

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundContext.tsx

@@ -10,14 +10,16 @@ export type PlaygroundContextType = {
   deleteMessage: (messageIndex: number, responseIndexes?: number[]) => void;
 
   editChoice: (choiceIndex: number, newChoice: Message) => void;
-  deleteChoice: (choiceIndex: number) => void;
+  deleteChoice: (messageIndex: number, choiceIndex: number) => void;
 
-  retry: (messageIndex: number, isChoice?: boolean) => void;
+  retry: (messageIndex: number, choiceIndex?: number) => void;
   sendMessage: (
     role: PlaygroundMessageRole,
     content: string,
     toolCallId?: string
   ) => void;
+
+  setSelectedChoiceIndex: (choiceIndex: number) => void;
 };
 
 const DEFAULT_CONTEXT: PlaygroundContextType = {
@@ -31,6 +33,7 @@ const DEFAULT_CONTEXT: PlaygroundContextType = {
 
   retry: () => {},
   sendMessage: () => {},
+  setSelectedChoiceIndex: () => {},
 };
 
 // Create context that can be undefined

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/types.ts

@@ -23,6 +23,7 @@ export type PlaygroundState = {
   //   nTimes: number;
   maxTokensLimit: number;
   model: LLMMaxTokensKey;
+  selectedChoiceIndex: number;
 };
 
 export type PlaygroundStateKey = keyof PlaygroundState;

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/usePlaygroundState.ts

@@ -37,6 +37,7 @@ const DEFAULT_PLAYGROUND_STATE = {
   //   nTimes: 1,
   maxTokensLimit: 16384,
   model: DEFAULT_MODEL,
+  selectedChoiceIndex: 0,
 };
 
 export const usePlaygroundState = () => {