refactor(prompts): Normalize playground instance messages (#6218)

* refactor(prompts): Normalize playground instance messages * Update use-zustand package to improve performance * Convert TemplateEditor, VariableEditor, and PlaygroundTool (editor) to uncontrolled components Workaround for uiwjs/react-codemirror#694 * Clean up variable names and semantics * Mark instance as dirty when updating message
Arize-ai · Feb 3, 2025 · b971d08 · b971d08
1 parent a36a384
commit b971d08
Show file tree

Hide file tree

Showing 23 changed files with 687 additions and 255 deletions.
diff --git a/app/package.json b/app/package.json
@@ -57,7 +57,7 @@
     "three": "~0.139.2",
     "three-stdlib": "^2.30.4",
     "use-deep-compare-effect": "^1.8.1",
-    "use-zustand": "^0.0.4",
+    "use-zustand": "^0.2.0",
     "zod": "^3.23.8",
     "zod-to-json-schema": "^3.23.3",
     "zustand": "^4.5.4"

diff --git a/app/pnpm-lock.yaml b/app/pnpm-lock.yaml
diff --git a/app/src/components/templateEditor/TemplateEditor.tsx b/app/src/components/templateEditor/TemplateEditor.tsx
@@ -1,4 +1,4 @@
-import React, { useMemo } from "react";
+import React, { useEffect, useMemo, useState } from "react";
 import { githubLight } from "@uiw/codemirror-theme-github";
 import { nord } from "@uiw/codemirror-theme-nord";
 import CodeMirror, {
@@ -16,8 +16,9 @@ import { MustacheLikeTemplating } from "./language/mustacheLike";
 import { TemplateLanguages } from "./constants";
 import { TemplateLanguage } from "./types";
 
-type TemplateEditorProps = ReactCodeMirrorProps & {
+type TemplateEditorProps = Omit<ReactCodeMirrorProps, "value"> & {
   templateLanguage: TemplateLanguage;
+  defaultValue: string;
 };
 
 const basicSetupOptions: BasicSetupOptions = {
@@ -28,10 +29,22 @@ const basicSetupOptions: BasicSetupOptions = {
   bracketMatching: false,
 };
 
+/**
+ * A template editor that is used to edit the template of a tool.
+ *
+ * This is an uncontrolled editor.
+ * You can only reset the value of the editor by triggering a re-mount, like with the `key` prop,
+ * or, when the readOnly prop is true, the editor will reset on all value changes.
+ * This is necessary because controlled react-codemirror editors incessantly reset
+ * cursor position when value is updated.
+ */
 export const TemplateEditor = ({
   templateLanguage,
+  defaultValue,
+  readOnly,
   ...props
 }: TemplateEditorProps) => {
+  const [value, setValue] = useState(() => defaultValue);
   const { theme } = useTheme();
   const codeMirrorTheme = theme === "light" ? githubLight : nord;
   const extensions = useMemo(() => {
@@ -51,12 +64,20 @@ export const TemplateEditor = ({
     return ext;
   }, [templateLanguage]);
 
+  useEffect(() => {
+    if (readOnly) {
+      setValue(defaultValue);
+    }
+  }, [readOnly, defaultValue]);
+
   return (
     <CodeMirror
       theme={codeMirrorTheme}
       extensions={extensions}
       basicSetup={basicSetupOptions}
+      readOnly={readOnly}
       {...props}
+      value={value}
     />
   );
 };

diff --git a/app/src/pages/playground/ChatMessageToolCallsEditor.tsx b/app/src/pages/playground/ChatMessageToolCallsEditor.tsx
@@ -8,32 +8,40 @@ import {
   llmProviderToolCallsSchema,
   openAIToolCallsJSONSchema,
 } from "@phoenix/schemas/toolCallSchemas";
-import { ChatMessage } from "@phoenix/store";
+import {
+  selectPlaygroundInstance,
+  selectPlaygroundInstanceMessage,
+} from "@phoenix/store/playground/selectors";
 import { safelyParseJSON } from "@phoenix/utils/jsonUtils";
 
-import { PlaygroundInstanceProps } from "./types";
-
 /**
  * Editor for message tool calls
  */
 export function ChatMessageToolCallsEditor({
   playgroundInstanceId,
-  toolCalls,
-  templateMessages,
   messageId,
-}: PlaygroundInstanceProps & {
-  toolCalls: ChatMessage["toolCalls"];
-  templateMessages: ChatMessage[];
+}: {
+  playgroundInstanceId: number;
   messageId: number;
 }) {
-  const updateInstance = usePlaygroundContext((state) => state.updateInstance);
-  const instance = usePlaygroundContext((state) =>
-    state.instances.find((instance) => instance.id === playgroundInstanceId)
+  const instanceSelector = useMemo(
+    () => selectPlaygroundInstance(playgroundInstanceId),
+    [playgroundInstanceId]
   );
-
+  const instance = usePlaygroundContext(instanceSelector);
   if (instance == null) {
-    throw new Error(`Playground instance ${playgroundInstanceId} not found`);
+    throw new Error(`Instance ${playgroundInstanceId} not found`);
+  }
+  const messageSelector = useMemo(
+    () => selectPlaygroundInstanceMessage(messageId),
+    [messageId]
+  );
+  const message = usePlaygroundContext(messageSelector);
+  if (message == null) {
+    throw new Error(`Message ${messageId} not found`);
   }
+  const toolCalls = message.toolCalls;
+  const updateMessage = usePlaygroundContext((state) => state.updateMessage);
   const [editorValue, setEditorValue] = useState(() =>
     JSON.stringify(toolCalls, null, 2)
   );
@@ -64,25 +72,15 @@ export function ChatMessageToolCallsEditor({
         return;
       }
       setLastValidToolCalls(toolCalls);
-      updateInstance({
+      updateMessage({
         instanceId: playgroundInstanceId,
+        messageId,
         patch: {
-          template: {
-            __type: "chat",
-            messages: templateMessages.map((m) =>
-              messageId === m.id
-                ? {
-                    ...m,
-                    toolCalls,
-                  }
-                : m
-            ),
-          },
+          toolCalls,
         },
-        dirty: true,
       });
     },
-    [messageId, playgroundInstanceId, templateMessages, updateInstance]
+    [playgroundInstanceId, messageId, updateMessage]
   );
 
   const toolCallsJSONSchema = useMemo((): JSONSchema7 | null => {

diff --git a/app/src/pages/playground/ModelConfigButton.tsx b/app/src/pages/playground/ModelConfigButton.tsx
@@ -31,7 +31,10 @@ import {
 import { useNotifySuccess } from "@phoenix/contexts";
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 import { usePreferencesContext } from "@phoenix/contexts/PreferencesContext";
-import { PlaygroundInstance } from "@phoenix/store";
+import {
+  PlaygroundInstance,
+  PlaygroundNormalizedInstance,
+} from "@phoenix/store";
 
 import { ModelConfigButtonDialogQuery } from "./__generated__/ModelConfigButtonDialogQuery.graphql";
 import { InvocationParametersFormFields } from "./InvocationParametersFormFields";
@@ -67,7 +70,7 @@ const modelConfigFormCSS = css`
 function AzureOpenAiModelConfigFormField({
   instance,
 }: {
-  instance: PlaygroundInstance;
+  instance: PlaygroundNormalizedInstance;
 }) {
   const updateModel = usePlaygroundContext((state) => state.updateModel);
   const updateModelConfig = useCallback(

diff --git a/app/src/pages/playground/Playground.tsx b/app/src/pages/playground/Playground.tsx
@@ -1,4 +1,4 @@
-import React, { Suspense, useCallback, useEffect } from "react";
+import React, { Fragment, Suspense, useCallback, useEffect } from "react";
 import { graphql, useLazyLoadQuery } from "react-relay";
 import { Panel, PanelGroup, PanelResizeHandle } from "react-resizable-panels";
 import { BlockerFunction, useBlocker, useSearchParams } from "react-router-dom";
@@ -205,6 +205,9 @@ const playgroundInputOutputPanelContentCSS = css`
  */
 const PLAYGROUND_PROMPT_PANEL_MIN_WIDTH = 632;
 
+const DEFAULT_EXPANDED_PROMPTS = ["prompts"];
+const DEFAULT_EXPANDED_PARAMS = ["input", "output"];
+
 function PlaygroundContent() {
   const instances = usePlaygroundContext((state) => state.instances);
   const templateLanguage = usePlaygroundContext(
@@ -246,7 +249,7 @@ function PlaygroundContent() {
   }, [isRunning, anyDirtyInstances]);
 
   return (
-    <>
+    <Fragment key="playground-content">
       <PanelGroup
         direction={
           isSingleInstance && !isDatasetMode ? "horizontal" : "vertical"
@@ -257,7 +260,7 @@ function PlaygroundContent() {
       >
         <Panel>
           <div css={playgroundPromptPanelContentCSS}>
-            <DisclosureGroup defaultExpandedKeys={["prompts"]}>
+            <DisclosureGroup defaultExpandedKeys={DEFAULT_EXPANDED_PROMPTS}>
               <Disclosure id="prompts" size="L">
                 <DisclosureTrigger
                   arrowPosition="start"
@@ -277,11 +280,10 @@ function PlaygroundContent() {
                       {instances.map((instance) => (
                         <View
                           flex="1 1 0px"
-                          key={instance.id}
+                          key={`${instance.id}-prompt`}
                           minWidth={PLAYGROUND_PROMPT_PANEL_MIN_WIDTH}
                         >
                           <PlaygroundTemplate
-                            key={instance.id}
                             playgroundInstanceId={instance.id}
                           />
                         </View>
@@ -301,7 +303,7 @@ function PlaygroundContent() {
             </Suspense>
           ) : (
             <div css={playgroundInputOutputPanelContentCSS}>
-              <DisclosureGroup defaultExpandedKeys={["input", "output"]}>
+              <DisclosureGroup defaultExpandedKeys={DEFAULT_EXPANDED_PARAMS}>
                 {templateLanguage !== TemplateLanguages.NONE ? (
                   <Disclosure id="input" size="L">
                     <DisclosureTrigger arrowPosition="start">
@@ -321,10 +323,9 @@ function PlaygroundContent() {
                   <DisclosurePanel>
                     <View padding="size-200" height="100%">
                       <Flex direction="row" gap="size-200">
-                        {instances.map((instance, i) => (
-                          <View key={i} flex="1 1 0px">
+                        {instances.map((instance) => (
+                          <View key={`${instance.id}-output`} flex="1 1 0px">
                             <PlaygroundOutput
-                              key={i}
                               playgroundInstanceId={instance.id}
                             />
                           </View>
@@ -348,6 +349,6 @@ function PlaygroundContent() {
           }
         />
       )}
-    </>
+    </Fragment>
   );
 }