Merge branch 'master' into DOCS-1050

wandb · Dec 11, 2024 · 5457651 · 5457651
2 parents 9891112 + 43e2b7b
commit 5457651
Show file tree

Hide file tree

Showing 16 changed files with 95 additions and 61 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -226,7 +226,7 @@ module = "weave_query.*"
 ignore_errors = true
 
 [tool.bumpversion]
-current_version = "0.51.24-dev0"
+current_version = "0.51.25-dev0"
 parse = """(?x)
     (?P<major>0|[1-9]\\d*)\\.
     (?P<minor>0|[1-9]\\d*)\\.

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/ChoiceView.tsx b/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}
     />
   );
 };
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/ChoicesDrawer.tsx b/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>

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/ChoicesView.tsx b/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) {

diff --git a/...js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/ChoicesViewCarousel.tsx b/...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">

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/MessagePanel.tsx b/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}

diff --git a/...ponents/PagePanelComponents/Home/Browse3/pages/ChatView/PlaygroundMessagePanelButtons.tsx b/...ponents/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);
           }

diff --git a/...mponents/PagePanelComponents/Home/Browse3/pages/ChatView/PlaygroundMessagePanelEditor.tsx b/...mponents/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, {

diff --git a/...s/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/PlaygroundChat.tsx b/...s/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>

diff --git a/...omponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/useChatCompletionFunctions.tsx b/...omponents/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;

diff --git a/...PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/useChatFunctions.tsx b/...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;
     });

diff --git a/...rc/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundContext.tsx b/...rc/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

diff --git a/...nelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundSettings/PlaygroundSettings.tsx b/...nelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundSettings/PlaygroundSettings.tsx
@@ -59,6 +59,12 @@ export const PlaygroundSettings: React.FC<PlaygroundSettingsProps> = ({
                 gap: '4px',
                 mt: 2,
               }}>
+              <ResponseFormatEditor
+                responseFormat={playgroundState.responseFormat}
+                setResponseFormat={value =>
+                  setPlaygroundStateField(idx, 'responseFormat', value)
+                }
+              />
               <FunctionEditor
                 playgroundState={playgroundState}
                 functions={playgroundState.functions}
@@ -71,22 +77,24 @@ export const PlaygroundSettings: React.FC<PlaygroundSettingsProps> = ({
                 }
               />
 
-              <ResponseFormatEditor
-                responseFormat={playgroundState.responseFormat}
-                setResponseFormat={value =>
-                  setPlaygroundStateField(idx, 'responseFormat', value)
+              <StopSequenceEditor
+                stopSequences={playgroundState.stopSequences}
+                setStopSequences={value =>
+                  setPlaygroundStateField(idx, 'stopSequences', value)
                 }
               />
 
+              {/* TODO: N times to run is not supported for all models */}
+              {/* TODO: rerun if this is not supported in the backend */}
               <PlaygroundSlider
-                min={0}
-                max={2}
-                step={0.01}
+                min={1}
+                max={100}
+                step={1}
                 setValue={value =>
-                  setPlaygroundStateField(idx, 'temperature', value)
+                  setPlaygroundStateField(idx, 'nTimes', value)
                 }
-                label="Temperature"
-                value={playgroundState.temperature}
+                label="Completion iterations"
+                value={playgroundState.nTimes}
               />
 
               <PlaygroundSlider
@@ -100,11 +108,15 @@ export const PlaygroundSettings: React.FC<PlaygroundSettingsProps> = ({
                 value={playgroundState.maxTokens}
               />
 
-              <StopSequenceEditor
-                stopSequences={playgroundState.stopSequences}
-                setStopSequences={value =>
-                  setPlaygroundStateField(idx, 'stopSequences', value)
+              <PlaygroundSlider
+                min={0}
+                max={2}
+                step={0.01}
+                setValue={value =>
+                  setPlaygroundStateField(idx, 'temperature', value)
                 }
+                label="Temperature"
+                value={playgroundState.temperature}
               />
 
               <PlaygroundSlider
@@ -137,6 +149,7 @@ export const PlaygroundSettings: React.FC<PlaygroundSettingsProps> = ({
                 label="Presence penalty"
                 value={playgroundState.presencePenalty}
               />
+
               <Box
                 sx={{
                   width: '100%',

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/types.ts b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/types.ts
@@ -20,9 +20,10 @@ export type PlaygroundState = {
   topP: number;
   frequencyPenalty: number;
   presencePenalty: number;
-  //   nTimes: number;
+  nTimes: number;
   maxTokensLimit: number;
   model: LLMMaxTokensKey;
+  selectedChoiceIndex: number;
 };
 
 export type PlaygroundStateKey = keyof PlaygroundState;