fix: Derive playground values instead of storing them in state

app/src/pages/playground/PlaygroundInput.tsx

@@ -3,24 +3,19 @@ import React from "react";
 import { Flex, Text, View } from "@arizeai/components";
 
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
-import {
-  selectDerivedInputVariables,
-  selectInputVariableKeys,
-} from "@phoenix/store";
 import { assertUnreachable } from "@phoenix/typeUtils";
 
+import { useDerivedPlaygroundVariables } from "./useDerivedPlaygroundVariables";
 import { VariableEditor } from "./VariableEditor";
 
 export function PlaygroundInput() {
-  const variables = usePlaygroundContext(selectDerivedInputVariables);
-  const variableKeys = usePlaygroundContext(selectInputVariableKeys);
+  const { variableKeys, variablesMap } = useDerivedPlaygroundVariables();
   const setVariableValue = usePlaygroundContext(
     (state) => state.setVariableValue
   );
   const templateLanguage = usePlaygroundContext(
     (state) => state.templateLanguage
   );
-
   if (variableKeys.length === 0) {
     let templateSyntax = "";
     switch (templateLanguage) {
@@ -59,7 +54,7 @@ export function PlaygroundInput() {
             // change rapidly for a given variable
             key={i}
             label={variableKey}
-            value={variables[variableKey]}
+            value={variablesMap[variableKey]}
             onChange={(value) => setVariableValue(variableKey, value)}
           />
         );

app/src/pages/playground/PlaygroundOutput.tsx

@@ -9,11 +9,7 @@ import { useCredentialsContext } from "@phoenix/contexts/CredentialsContext";
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";
 import type { ToolCall } from "@phoenix/store";
-import {
-  ChatMessage,
-  generateMessageId,
-  selectDerivedInputVariables,
-} from "@phoenix/store";
+import { ChatMessage, generateMessageId } from "@phoenix/store";
 import { assertUnreachable } from "@phoenix/typeUtils";
 
 import {
@@ -26,6 +22,7 @@ import {
 import { isChatMessages } from "./playgroundUtils";
 import { TitleWithAlphabeticIndex } from "./TitleWithAlphabeticIndex";
 import { PlaygroundInstanceProps } from "./types";
+import { useDerivedPlaygroundVariables } from "./useDerivedPlaygroundVariables";
 
 interface PlaygroundOutputProps extends PlaygroundInstanceProps {}
 
@@ -221,7 +218,7 @@ function PlaygroundOutputText(props: PlaygroundInstanceProps) {
   const templateLanguage = usePlaygroundContext(
     (state) => state.templateLanguage
   );
-  const templateVariables = usePlaygroundContext(selectDerivedInputVariables);
+  const { variablesMap: templateVariables } = useDerivedPlaygroundVariables();
   const markPlaygroundInstanceComplete = usePlaygroundContext(
     (state) => state.markPlaygroundInstanceComplete
   );

app/src/pages/playground/__tests__/playgroundUtils.test.ts

@@ -31,7 +31,7 @@ const expectedPlaygroundInstanceWithIO: PlaygroundInstance = {
     provider: "OPENAI",
     modelName: "gpt-3.5-turbo",
   },
-  input: { variableKeys: [], variablesValueCache: {} },
+  input: { variablesValueCache: {} },
   tools: [],
   toolChoice: "auto",
   template: {

app/src/pages/playground/playgroundUtils.ts

@@ -1,3 +1,5 @@
+import { getTemplateLanguageUtils } from "@phoenix/components/templateEditor/templateEditorUtils";
+import { TemplateLanguage } from "@phoenix/components/templateEditor/types";
 import {
   DEFAULT_CHAT_ROLE,
   DEFAULT_MODEL_PROVIDER,
@@ -8,6 +10,7 @@ import {
   createPlaygroundInstance,
   generateMessageId,
 } from "@phoenix/store";
+import { assertUnreachable } from "@phoenix/typeUtils";
 import { safelyParseJSON } from "@phoenix/utils/jsonUtils";
 
 import {
@@ -240,3 +243,49 @@ export const isChatMessages = (
 ): messages is ChatMessage[] => {
   return chatMessagesSchema.safeParse(messages).success;
 };
+
+export const extractVariablesFromInstances = ({
+  instances,
+  templateLanguage,
+}: {
+  instances: PlaygroundInstance[];
+  templateLanguage: TemplateLanguage;
+}) => {
+  const variables = new Set<string>();
+  const utils = getTemplateLanguageUtils(templateLanguage);
+  instances.forEach((instance) => {
+    const instanceType = instance.template.__type;
+    // this double nested loop should be okay since we don't expect more than 4 instances
+    // and a handful of messages per instance
+    switch (instanceType) {
+      case "chat": {
+        // for each chat message in the instance
+        instance.template.messages.forEach((message) => {
+          // extract variables from the message content
+          const extractedVariables =
+            message.content == null
+              ? []
+              : utils.extractVariables(message.content);
+          extractedVariables.forEach((variable) => {
+            variables.add(variable);
+          });
+        });
+        break;
+      }
+      case "text_completion": {
+        const extractedVariables = utils.extractVariables(
+          instance.template.prompt
+        );
+        extractedVariables.forEach((variable) => {
+          variables.add(variable);
+        });
+        break;
+      }
+      default: {
+        assertUnreachable(instanceType);
+      }
+    }
+  });
+
+  return Array.from(variables);
+};

app/src/pages/playground/useDerivedPlaygroundVariables.tsx

@@ -0,0 +1,49 @@
+import { useMemo } from "react";
+
+import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
+import { isManualInput } from "@phoenix/store";
+
+import { extractVariablesFromInstances } from "./playgroundUtils";
+
+/**
+ * Get the variable keys from all instances in the playground, using the
+ * template language to determine the syntax of the variables
+ */
+export const useDerivedPlaygroundVariableKeys = () => {
+  const { instances, templateLanguage } = usePlaygroundContext((state) => ({
+    instances: state.instances,
+    templateLanguage: state.templateLanguage,
+  }));
+  const variableKeys = useMemo(() => {
+    return extractVariablesFromInstances({
+      instances,
+      templateLanguage,
+    });
+  }, [instances, templateLanguage]);
+
+  return variableKeys;
+};
+
+/**
+ * Get the variable values and keys from all instances in the playground
+ *
+ * Variables are recomputed whenever _anything_ in the playground instances change
+ * or when the template language changes. This can be optimized in the future.
+ */
+export const useDerivedPlaygroundVariables = () => {
+  const variableKeys = useDerivedPlaygroundVariableKeys();
+  const variableValueCache = usePlaygroundContext((state) =>
+    isManualInput(state.input) ? state.input.variablesValueCache : {}
+  );
+  const variablesMap = useMemo(() => {
+    return variableKeys.reduce(
+      (acc, key) => {
+        acc[key] = variableValueCache[key] || "";
+        return acc;
+      },
+      {} as Record<string, string>
+    );
+  }, [variableKeys, variableValueCache]);
+
+  return { variableKeys, variablesMap };
+};

app/src/store/playground/playgroundStore.tsx

@@ -2,13 +2,11 @@ import { create, StateCreator } from "zustand";
 import { devtools } from "zustand/middleware";
 
 import { TemplateLanguages } from "@phoenix/components/templateEditor/constants";
-import { getTemplateLanguageUtils } from "@phoenix/components/templateEditor/templateEditorUtils";
 import { TemplateLanguage } from "@phoenix/components/templateEditor/types";
 import {
   DEFAULT_CHAT_ROLE,
   DEFAULT_MODEL_PROVIDER,
 } from "@phoenix/constants/generativeConstants";
-import { assertUnreachable } from "@phoenix/typeUtils";
 
 import {
   GenAIOperationType,
@@ -100,7 +98,7 @@ export function createPlaygroundInstance(): PlaygroundInstance {
     // Default to auto tool choice as you are probably testing the LLM for it's ability to pick
     toolChoice: "auto",
     // TODO(apowell) - use datasetId if in dataset mode
-    input: { variablesValueCache: {}, variableKeys: [] },
+    input: { variablesValueCache: {} },
     output: undefined,
     activeRunId: null,
     isRunning: false,
@@ -135,13 +133,10 @@ export const createPlaygroundStore = (
     operationType: "chat",
     inputMode: "manual",
     input: {
-      // to get a record of visible variables and their values,
-      // use usePlaygroundContext(selectDerivedInputVariables). do not render variablesValueCache
-      // directly or users will see stale values.
-      variablesValueCache: {
-        question: "",
-      },
-      variableKeys: ["question"],
+      // variablesValueCache is used to store the values of variables for the
+      // manual input mode. It is indexed by the variable key. It keeps old
+      // values when variables are removed so that they can be restored.
+      variablesValueCache: {},
     },
     templateLanguage: TemplateLanguages.Mustache,
     setInputMode: (inputMode: PlaygroundInputMode) => set({ inputMode }),
@@ -252,7 +247,6 @@ export const createPlaygroundStore = (
           return instance;
         }),
       });
-      get().calculateVariables();
     },
     runPlaygroundInstances: () => {
       const instances = get().instances;
@@ -295,49 +289,6 @@ export const createPlaygroundStore = (
     },
     setTemplateLanguage: (templateLanguage: TemplateLanguage) => {
       set({ templateLanguage });
-      // Re-compute variables when the template language changes
-      get().calculateVariables();
-    },
-    calculateVariables: () => {
-      const instances = get().instances;
-      const variables = new Set<string>();
-      const utils = getTemplateLanguageUtils(get().templateLanguage);
-      instances.forEach((instance) => {
-        const instanceType = instance.template.__type;
-        // this double nested loop should be okay since we don't expect more than 4 instances
-        // and a handful of messages per instance
-        switch (instanceType) {
-          case "chat": {
-            // for each chat message in the instance
-            instance.template.messages.forEach((message) => {
-              // extract variables from the message content
-              const extractedVariables =
-                message.content == null
-                  ? []
-                  : utils.extractVariables(message.content);
-              extractedVariables.forEach((variable) => {
-                variables.add(variable);
-              });
-            });
-            break;
-          }
-          case "text_completion": {
-            const extractedVariables = utils.extractVariables(
-              instance.template.prompt
-            );
-            extractedVariables.forEach((variable) => {
-              variables.add(variable);
-            });
-            break;
-          }
-          default: {
-            assertUnreachable(instanceType);
-          }
-        }
-      });
-      set({
-        input: { ...get().input, variableKeys: [...Array.from(variables)] },
-      });
     },
     setVariableValue: (key: string, value: string) => {
       const input = get().input;
@@ -356,34 +307,3 @@ export const createPlaygroundStore = (
 };
 
 export type PlaygroundStore = ReturnType<typeof createPlaygroundStore>;
-
-/**
- * Selects the variable keys from the playground state
- * @param state the playground state
- * @returns the variable keys
- */
-export const selectInputVariableKeys = (state: PlaygroundState) => {
-  if (isManualInput(state.input)) {
-    return state.input.variableKeys;
-  }
-  return [];
-};
-
-/**
- * Selects the derived input variables from the playground state
- * @param state the playground state
- * @returns the derived input variables
- */
-export const selectDerivedInputVariables = (state: PlaygroundState) => {
-  if (isManualInput(state.input)) {
-    const input = state.input;
-    const variableKeys = input.variableKeys;
-    const variablesValueCache = input.variablesValueCache;
-    const valueMap: Record<string, string> = {};
-    variableKeys.forEach((key) => {
-      valueMap[key] = variablesValueCache?.[key] || "";
-    });
-    return valueMap;
-  }
-  return {};
-};

app/src/store/playground/types.ts

@@ -87,7 +87,6 @@ type DatasetInput = {
 
 type ManualInput = {
   variablesValueCache: Record<string, string | undefined>;
-  variableKeys: string[];
 };
 
 type PlaygroundInput = DatasetInput | ManualInput;
@@ -231,18 +230,16 @@ export interface PlaygroundState extends PlaygroundProps {
    */
   setTemplateLanguage: (templateLanguage: TemplateLanguage) => void;
   /**
-   * Calculate the variables used across all instances
+   * Set the value of a variable in the input
    */
-  calculateVariables: () => void;
-
   setVariableValue: (key: string, value: string) => void;
 }
 
 /**
  * Check if the input is manual
  */
 export const isManualInput = (input: PlaygroundInput): input is ManualInput => {
-  return "variablesValueCache" in input && "variableKeys" in input;
+  return "variablesValueCache" in input;
 };
 
 /**