feat(playground): Implement editable input variable textareas (#4987) (…

…#5006)

* feat(playground): Implement editable input variable textareas (#4987)

* docs(playground): Improve commenting around playground store usage

* refactor(playground): Switch to mostly unstyled codemirror for variable editing

* fix(playground): Fix styling issues and flickering when rendering playground variable inputs

* refactor(playground): Add and use type guards for playground input type

* docs(playground): update comment on variable cache mechanics

* fix(playground): fix types and tests post rebase

app/src/components/code/CodeWrap.tsx

@@ -1,10 +1,18 @@
 import React, { ReactNode } from "react";
 
-import { View } from "@arizeai/components";
+import { View, ViewStyleProps } from "@arizeai/components";
 
-export function CodeWrap({ children }: { children: ReactNode }) {
+export function CodeWrap({
+  children,
+  ...props
+}: { children: ReactNode } & ViewStyleProps) {
   return (
-    <View borderColor="light" borderWidth="thin" borderRadius="small">
+    <View
+      borderColor="light"
+      borderWidth="thin"
+      borderRadius="small"
+      {...props}
+    >
       {children}
     </View>
   );

app/src/pages/playground/Playground.tsx

@@ -22,6 +22,7 @@ import { InitialPlaygroundState } from "@phoenix/store";
 
 import { NUM_MAX_PLAYGROUND_INSTANCES } from "./constants";
 import { PlaygroundCredentialsDropdown } from "./PlaygroundCredentialsDropdown";
+import { PlaygroundInput } from "./PlaygroundInput";
 import { PlaygroundInputTypeTypeRadioGroup } from "./PlaygroundInputModeRadioGroup";
 import { PlaygroundOutput } from "./PlaygroundOutput";
 import { PlaygroundRunButton } from "./PlaygroundRunButton";
@@ -121,7 +122,6 @@ const playgroundInputOutputPanelContentCSS = css`
 
 function PlaygroundContent() {
   const instances = usePlaygroundContext((state) => state.instances);
-  const inputs = usePlaygroundContext((state) => state.input);
   const numInstances = instances.length;
   const isSingleInstance = numInstances === 1;
 
@@ -170,13 +170,9 @@ function PlaygroundContent() {
               id="input"
               extra={<PlaygroundInputTypeTypeRadioGroup />}
             >
-              <pre>
-                {JSON.stringify(
-                  "variables" in inputs ? inputs.variables : "inputs go here",
-                  null,
-                  2
-                )}
-              </pre>
+              <View padding="size-200">
+                <PlaygroundInput />
+              </View>
             </AccordionItem>
             <AccordionItem title="Output" id="output">
               <View padding="size-200" height="100%">

app/src/pages/playground/PlaygroundInput.tsx

@@ -1,11 +1,69 @@
 import React from "react";
 
-import { Card } from "@arizeai/components";
+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 { VariableEditor } from "./VariableEditor";
 
 export function PlaygroundInput() {
+  const variables = usePlaygroundContext(selectDerivedInputVariables);
+  const variableKeys = usePlaygroundContext(selectInputVariableKeys);
+  const setVariableValue = usePlaygroundContext(
+    (state) => state.setVariableValue
+  );
+  const templateLanguage = usePlaygroundContext(
+    (state) => state.templateLanguage
+  );
+
+  if (variableKeys.length === 0) {
+    let templateSyntax = "";
+    switch (templateLanguage) {
+      case "f-string": {
+        templateSyntax = "{input name}";
+        break;
+      }
+      case "mustache": {
+        templateSyntax = "{{input name}}";
+        break;
+      }
+      default:
+        assertUnreachable(templateLanguage);
+    }
+    return (
+      <View padding="size-100">
+        <Flex direction="column" justifyContent="center" alignItems="center">
+          <Text color="text-700">
+            Add variable inputs to your prompt using{" "}
+            <Text color="text-900">{templateSyntax}</Text> within your prompt
+            template.
+          </Text>
+        </Flex>
+      </View>
+    );
+  }
+
   return (
-    <Card title="Input" collapsible variant="compact">
-      Input goes here
-    </Card>
+    <Flex direction="column" gap="size-200" width="100%">
+      {variableKeys.map((variableKey, i) => {
+        return (
+          <VariableEditor
+            // using the index as the key actually prevents the UI from
+            // flickering; if we use the variable key directly, it will
+            // re-mount the entire editor and cause a flicker because key may
+            // change rapidly for a given variable
+            key={i}
+            label={variableKey}
+            value={variables[variableKey]}
+            onChange={(value) => setVariableValue(variableKey, value)}
+          />
+        );
+      })}
+    </Flex>
   );
 }

app/src/pages/playground/VariableEditor.tsx

@@ -0,0 +1,51 @@
+import React from "react";
+import { githubLight } from "@uiw/codemirror-theme-github";
+import { nord } from "@uiw/codemirror-theme-nord";
+import ReactCodeMirror, {
+  BasicSetupOptions,
+  EditorView,
+} from "@uiw/react-codemirror";
+
+import { Field } from "@arizeai/components";
+
+import { CodeWrap } from "@phoenix/components/code";
+import { useTheme } from "@phoenix/contexts";
+
+type VariableEditorProps = {
+  label?: string;
+  value?: string;
+  onChange?: (value: string) => void;
+};
+
+const basicSetupOptions: BasicSetupOptions = {
+  lineNumbers: false,
+  highlightActiveLine: false,
+  foldGutter: false,
+  highlightActiveLineGutter: false,
+  bracketMatching: false,
+  syntaxHighlighting: false,
+};
+
+const extensions = [EditorView.lineWrapping];
+
+export const VariableEditor = ({
+  label,
+  value,
+  onChange,
+}: VariableEditorProps) => {
+  const { theme } = useTheme();
+  const codeMirrorTheme = theme === "light" ? githubLight : nord;
+  return (
+    <Field label={label}>
+      <CodeWrap width="100%">
+        <ReactCodeMirror
+          theme={codeMirrorTheme}
+          basicSetup={basicSetupOptions}
+          value={value}
+          extensions={extensions}
+          onChange={onChange}
+        />
+      </CodeWrap>
+    </Field>
+  );
+};

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

@@ -26,7 +26,7 @@ const expectedPlaygroundInstanceWithIO: PlaygroundInstance = {
     provider: "OPENAI",
     modelName: "gpt-4o",
   },
-  input: { variables: {} },
+  input: { variableKeys: [], variablesValueCache: {} },
   tools: [],
   toolChoice: "auto",
   template: {

app/src/store/playground/playgroundStore.tsx

@@ -9,6 +9,7 @@ import { assertUnreachable } from "@phoenix/typeUtils";
 import {
   GenAIOperationType,
   InitialPlaygroundState,
+  isManualInput,
   PlaygroundChatTemplate,
   PlaygroundInputMode,
   PlaygroundInstance,
@@ -93,7 +94,8 @@ export function createPlaygroundInstance(): PlaygroundInstance {
     model: { provider: "OPENAI", modelName: "gpt-4o" },
     tools: [],
     toolChoice: "auto",
-    input: { variables: {} },
+    // TODO(apowell) - use datasetId if in dataset mode
+    input: { variablesValueCache: {}, variableKeys: [] },
     output: undefined,
     activeRunId: null,
     isRunning: false,
@@ -128,15 +130,13 @@ export const createPlaygroundStore = (
     operationType: "chat",
     inputMode: "manual",
     input: {
-      // TODO(apowell): When implementing variable forms, we should maintain a separate
-      // map of variableName to variableValue. This will allow us to "cache" variable values
-      // as the user types and will prevent data loss if users accidentally change the variable name
-      variables: {
-        // TODO(apowell): This is hardcoded based on the default chat template
-        // Instead we should calculate this based on the template on store creation
-        // Not a huge deal since this will be overridden on the first keystroke
+      // 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"],
     },
     templateLanguage: TemplateLanguages.Mustache,
     setInputMode: (inputMode: PlaygroundInputMode) => set({ inputMode }),
@@ -295,7 +295,7 @@ export const createPlaygroundStore = (
     },
     calculateVariables: () => {
       const instances = get().instances;
-      const variables: Record<string, string> = {};
+      const variables = new Set<string>();
       const utils = getTemplateLanguageUtils(get().templateLanguage);
       instances.forEach((instance) => {
         const instanceType = instance.template.__type;
@@ -310,7 +310,7 @@ export const createPlaygroundStore = (
                 message.content
               );
               extractedVariables.forEach((variable) => {
-                variables[variable] = "";
+                variables.add(variable);
               });
             });
             break;
@@ -320,7 +320,7 @@ export const createPlaygroundStore = (
               instance.template.prompt
             );
             extractedVariables.forEach((variable) => {
-              variables[variable] = "";
+              variables.add(variable);
             });
             break;
           }
@@ -329,11 +329,55 @@ export const createPlaygroundStore = (
           }
         }
       });
-      set({ input: { variables: { ...variables } } });
+      set({
+        input: { ...get().input, variableKeys: [...Array.from(variables)] },
+      });
+    },
+    setVariableValue: (key: string, value: string) => {
+      const input = get().input;
+      if (isManualInput(input)) {
+        set({
+          input: {
+            ...input,
+            variablesValueCache: { ...input.variablesValueCache, [key]: value },
+          },
+        });
+      }
     },
     ...initialProps,
   });
   return create(devtools(playgroundStore));
 };
 
 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

@@ -49,7 +49,8 @@ type DatasetInput = {
 };
 
 type ManualInput = {
-  variables: Record<string, string>;
+  variablesValueCache: Record<string, string | undefined>;
+  variableKeys: string[];
 };
 
 type PlaygroundInput = DatasetInput | ManualInput;
@@ -190,4 +191,22 @@ export interface PlaygroundState extends PlaygroundProps {
    * Calculate the variables used across all instances
    */
   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;
+};
+
+/**
+ * Check if the input is a dataset
+ */
+export const isDatasetInput = (
+  input: PlaygroundInput
+): input is DatasetInput => {
+  return "datasetId" in input;
+};