feat(prompts): Display dirty state on playground instance (#5961)

* feat(prompts): Derive dirty state from playground instance * Refactor dirty checking into hook * Clear prompt on compare * Manually track dirty state via store actions * Keep prompt selected when cloning instance * Fix types
Arize-ai · Jan 9, 2025 · b34c106 · b34c106
1 parent ef5ce56
commit b34c106
Show file tree

Hide file tree

Showing 13 changed files with 80 additions and 20 deletions.
diff --git a/app/src/pages/playground/ChatMessageToolCallsEditor.tsx b/app/src/pages/playground/ChatMessageToolCallsEditor.tsx
@@ -79,6 +79,7 @@ export function ChatMessageToolCallsEditor({
             ),
           },
         },
+        dirty: true,
       });
     },
     [messageId, playgroundInstanceId, templateMessages, updateInstance]

diff --git a/app/src/pages/playground/ModelSupportedParamsFetcher.tsx b/app/src/pages/playground/ModelSupportedParamsFetcher.tsx
@@ -16,16 +16,13 @@ export const ModelSupportedParamsFetcher = ({
 }: {
   instanceId: number;
 }) => {
-  const modelProvider = usePlaygroundContext(
-    (state) =>
-      state.instances.find((instance) => instance.id === instanceId)?.model
-        .provider
-  );
-  const modelName = usePlaygroundContext(
-    (state) =>
-      state.instances.find((instance) => instance.id === instanceId)?.model
-        .modelName
-  );
+  const instances = usePlaygroundContext((state) => state.instances);
+  const instance = instances.find((instance) => instance.id === instanceId);
+  if (!instance) {
+    throw new Error("Instance not found");
+  }
+  const modelProvider = instance.model.provider;
+  const modelName = instance.model.modelName;
   const modelConfigByProvider = usePreferencesContext(
     (state) => state.modelConfigByProvider
   );
@@ -88,6 +85,7 @@ export const ModelSupportedParamsFetcher = ({
         },
       }
     );
+  const promptId = instance.prompt?.id;
   useEffect(() => {
     updateModelSupportedInvocationParameters({
       instanceId,
@@ -100,6 +98,7 @@ export const ModelSupportedParamsFetcher = ({
     instanceId,
     updateModelSupportedInvocationParameters,
     modelConfigByProvider,
+    promptId,
   ]);
   return null;
 };
diff --git a/app/src/pages/playground/PlaygroundChatTemplate.tsx b/app/src/pages/playground/PlaygroundChatTemplate.tsx
@@ -135,6 +135,7 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
               messages: newMessages,
             },
           },
+          dirty: true,
         });
       }}
     >
@@ -343,6 +344,7 @@ function SortableMessageItem({
             ),
           },
         },
+        dirty: true,
       });
     },
     [message.id, playgroundInstanceId, template.messages, updateInstance]
@@ -396,6 +398,7 @@ function SortableMessageItem({
                       ),
                     },
                   },
+                  dirty: true,
                 });
               }}
             />
@@ -464,6 +467,7 @@ function SortableMessageItem({
                       ),
                     },
                   },
+                  dirty: true,
                 });
               }}
             />

diff --git a/app/src/pages/playground/PlaygroundChatTemplateFooter.tsx b/app/src/pages/playground/PlaygroundChatTemplateFooter.tsx
@@ -117,6 +117,7 @@ export function PlaygroundChatTemplateFooter({
                   }),
                 ],
               },
+              dirty: true,
             });
           }}
         >
@@ -143,6 +144,7 @@ export function PlaygroundChatTemplateFooter({
                 ],
               },
             },
+            dirty: true,
           });
         }}
       >

diff --git a/app/src/pages/playground/PlaygroundDatasetExamplesTable.tsx b/app/src/pages/playground/PlaygroundDatasetExamplesTable.tsx
@@ -465,6 +465,7 @@ export function PlaygroundDatasetExamplesTable({
             updateInstance({
               instanceId,
               patch: { experimentId: chatCompletion.experiment.id },
+              dirty: null,
             });
             break;
           case "ChatCompletionSubscriptionResult":
@@ -552,6 +553,7 @@ export function PlaygroundDatasetExamplesTable({
           patch: {
             experimentId: response.chatCompletionOverDataset.experimentId,
           },
+          dirty: null,
         });
         setExampleDataForInstance({
           instanceId,
@@ -581,6 +583,7 @@ export function PlaygroundDatasetExamplesTable({
         updateInstance({
           instanceId: instance.id,
           patch: { experimentId: null },
+          dirty: null,
         });
         if (activeRunId === null) {
           continue;
@@ -638,6 +641,7 @@ export function PlaygroundDatasetExamplesTable({
         updateInstance({
           instanceId: instance.id,
           patch: { experimentId: null },
+          dirty: null,
         });
         const variables = {
           input: getChatCompletionOverDatasetInput({

diff --git a/app/src/pages/playground/PlaygroundOutput.tsx b/app/src/pages/playground/PlaygroundOutput.tsx
@@ -245,6 +245,7 @@ export function PlaygroundOutput(props: PlaygroundOutputProps) {
           patch: {
             spanId: chatCompletion.span.id,
           },
+          dirty: null,
         });
         return;
       }
@@ -279,6 +280,7 @@ export function PlaygroundOutput(props: PlaygroundOutputProps) {
         patch: {
           spanId: response.chatCompletion.span.id,
         },
+        dirty: null,
       });
       if (errors) {
         notifyError({
@@ -315,6 +317,7 @@ export function PlaygroundOutput(props: PlaygroundOutputProps) {
       patch: {
         spanId: null,
       },
+      dirty: null,
     });
   }, [instanceId, updateInstance]);
 

diff --git a/app/src/pages/playground/PlaygroundOutputMoveButton.tsx b/app/src/pages/playground/PlaygroundOutputMoveButton.tsx
@@ -75,6 +75,7 @@ export const PlaygroundOutputMoveButton = ({
                 messages: [...instance.template.messages, ...messages],
               },
             },
+            dirty: true,
           });
           cleanupOutput();
         }}

diff --git a/app/src/pages/playground/PlaygroundTemplate.tsx b/app/src/pages/playground/PlaygroundTemplate.tsx
@@ -32,16 +32,13 @@ export function PlaygroundTemplate(props: PlaygroundTemplateProps) {
   const index = instances.findIndex((instance) => instance.id === instanceId);
   const prompt = instance?.prompt;
   const promptId = prompt?.id;
+  const dirty = instance?.dirty;
 
   const onChangePrompt = useCallback(
     async (promptId: string | null) => {
       if (!promptId) {
-        updateInstance({
-          instanceId,
-          patch: {
-            prompt: null,
-          },
-        });
+        const patch = { prompt: null };
+        updateInstance({ instanceId, patch, dirty: false });
         return;
       }
 
@@ -52,6 +49,7 @@ export function PlaygroundTemplate(props: PlaygroundTemplateProps) {
           patch: {
             ...response.instance,
           },
+          dirty: false,
         });
       }
     },
@@ -94,7 +92,11 @@ export function PlaygroundTemplate(props: PlaygroundTemplateProps) {
               <ModelSupportedParamsFetcher instanceId={instanceId} />
             </Suspense>
             <ModelConfigButton {...props} />
-            <SaveButton instanceId={instanceId} setDialog={setDialog} />
+            <SaveButton
+              instanceId={instanceId}
+              setDialog={setDialog}
+              dirty={dirty}
+            />
             {instances.length > 1 ? <DeleteButton {...props} /> : null}
           </Flex>
         }
@@ -143,9 +145,10 @@ function DeleteButton(props: PlaygroundInstanceProps) {
 type SaveButtonProps = {
   instanceId: number;
   setDialog: (dialog: React.ReactNode) => void;
+  dirty?: boolean;
 };
 
-function SaveButton({ instanceId, setDialog }: SaveButtonProps) {
+function SaveButton({ instanceId, setDialog, dirty }: SaveButtonProps) {
   const instance = usePlaygroundContext((state) =>
     state.instances.find((instance) => instance.id === instanceId)
   );
@@ -168,7 +171,7 @@ function SaveButton({ instanceId, setDialog }: SaveButtonProps) {
       <TooltipTrigger delay={100}>
         <TriggerWrap>
           <Button
-            // TODO(apowell): Make variant "primary" when instance is "dirty", aka different from selected prompt
+            variant={dirty ? "primary" : undefined}
             size="S"
             onPress={onSave}
           >

diff --git a/app/src/pages/playground/PlaygroundTool.tsx b/app/src/pages/playground/PlaygroundTool.tsx
@@ -104,6 +104,7 @@ export function PlaygroundTool({
               : t
           ),
         },
+        dirty: true,
       });
     },
     [instanceTools, playgroundInstanceId, tool.id, updateInstance]
@@ -165,6 +166,7 @@ export function PlaygroundTool({
                   tools: newTools,
                   toolChoice,
                 },
+                dirty: true,
               });
             }}
           />

diff --git a/app/src/pages/playground/PlaygroundTools.tsx b/app/src/pages/playground/PlaygroundTools.tsx
@@ -60,6 +60,7 @@ export function PlaygroundTools(props: PlaygroundToolsProps) {
                     patch: {
                       toolChoice: choice,
                     },
+                    dirty: true,
                   });
                 }}
                 toolNames={toolNames}

diff --git a/app/src/pages/playground/__tests__/playgroundUtils.test.ts b/app/src/pages/playground/__tests__/playgroundUtils.test.ts
@@ -68,6 +68,7 @@ const baseTestPlaygroundInstance: PlaygroundInstance = {
   tools: [],
   toolChoice: "auto",
   spanId: null,
+  dirty: false,
   template: {
     __type: "chat",
     messages: [],
@@ -86,6 +87,7 @@ const expectedPlaygroundInstanceWithIO: PlaygroundInstance = {
   tools: [],
   toolChoice: "auto",
   spanId: "fake-id",
+  dirty: false,
   template: {
     __type: "chat",
     // These id's are not 0, 1, 2, because we create a playground instance (including messages) at the top of the transformSpanAttributesToPlaygroundInstance function
@@ -994,6 +996,7 @@ describe("getVariablesMapFromInstances", () => {
     tools: [],
     toolChoice: "auto",
     spanId: null,
+    dirty: false,
     template: {
       __type: "chat",
       messages: [],

diff --git a/app/src/store/playground/playgroundStore.tsx b/app/src/store/playground/playgroundStore.tsx
@@ -116,6 +116,7 @@ export const DEFAULT_INSTANCE_PARAMS = () =>
     output: undefined,
     spanId: null,
     activeRunId: null,
+    dirty: false,
   }) satisfies Partial<PlaygroundInstance>;
 
 export function createPlaygroundInstance(): PlaygroundInstance {
@@ -313,6 +314,7 @@ export const createPlaygroundStore = (initialProps: InitialPlaygroundState) => {
             return {
               ...instance,
               ...patch,
+              dirty: true,
             };
           }
           return instance;
@@ -330,6 +332,7 @@ export const createPlaygroundStore = (initialProps: InitialPlaygroundState) => {
           if (instance.id === instanceId) {
             return {
               ...instance,
+              dirty: true,
               model: {
                 ...instance.model,
                 ...patch,
@@ -359,6 +362,7 @@ export const createPlaygroundStore = (initialProps: InitialPlaygroundState) => {
           ) {
             return {
               ...instance,
+              dirty: true,
               messages: [
                 ...instance.template.messages,
                 { role: DEFAULT_CHAT_ROLE, content: "{question}" },
@@ -369,14 +373,15 @@ export const createPlaygroundStore = (initialProps: InitialPlaygroundState) => {
         }),
       });
     },
-    updateInstance: ({ instanceId, patch }) => {
+    updateInstance: ({ instanceId, patch, dirty }) => {
       const instances = get().instances;
       set({
         instances: instances.map((instance) => {
           if (instance.id === instanceId) {
             return {
               ...instance,
               ...patch,
+              ...(dirty != undefined ? { dirty } : {}),
             };
           }
           return instance;
@@ -444,6 +449,7 @@ export const createPlaygroundStore = (initialProps: InitialPlaygroundState) => {
           if (instance.id === instanceId) {
             return {
               ...instance,
+              dirty: true,
               model: { ...instance.model, invocationParameters },
             };
           }
@@ -470,6 +476,7 @@ export const createPlaygroundStore = (initialProps: InitialPlaygroundState) => {
             if (instance.id === instanceId) {
               return {
                 ...instance,
+                dirty: true,
                 model: {
                   ...instance.model,
                   invocationParameters: instance.model.invocationParameters.map(
@@ -490,6 +497,7 @@ export const createPlaygroundStore = (initialProps: InitialPlaygroundState) => {
             if (instance.id === instanceId) {
               return {
                 ...instance,
+                dirty: true,
                 model: {
                   ...instance.model,
                   invocationParameters: [
@@ -517,6 +525,7 @@ export const createPlaygroundStore = (initialProps: InitialPlaygroundState) => {
           if (instance.id === instanceId) {
             return {
               ...instance,
+              dirty: true,
               model: {
                 ...instance.model,
                 invocationParameters:
@@ -532,6 +541,20 @@ export const createPlaygroundStore = (initialProps: InitialPlaygroundState) => {
         }),
       });
     },
+    setDirty: (instanceId: number, dirty: boolean) => {
+      const instances = get().instances;
+      set({
+        instances: instances.map((instance) => {
+          if (instance.id === instanceId) {
+            return {
+              ...instance,
+              dirty,
+            };
+          }
+          return instance;
+        }),
+      });
+    },
     ...initialProps,
   });
   return create(devtools(playgroundStore));
-Original file line number
+Diff line change
@@ Expand Up / @@ -79,6 +79,7 @@ export function ChatMessageToolCallsEditor({ @@
                 ),
               },
             },
+            dirty: true,
           });
         },
         [messageId, playgroundInstanceId, templateMessages, updateInstance]
@@ Expand Down @@