From 5e6faf0a5bb9482e0975ae26b110a21392237aeb Mon Sep 17 00:00:00 2001
From: Mikyo King <mikyo@arize.com>
Date: Tue, 8 Oct 2024 23:37:12 -0600
Subject: [PATCH] feat(playground): make chat messages match style in traces
 (#4931)

* feat(playground): make chat messages match style in traces

* small fixes

* WIP
---
 app/src/hooks/useChatMessageStyles.ts         | 35 +++++++++
 app/src/pages/playground/Playground.tsx       |  6 +-
 .../playground/PlaygroundChatTemplate.tsx     | 75 +++++++++++++------
 .../pages/playground/SpanPlaygroundPage.tsx   |  8 +-
 app/src/pages/trace/SpanDetails.tsx           | 30 +-------
 app/src/store/playgroundStore.tsx             | 24 +++---
 6 files changed, 112 insertions(+), 66 deletions(-)
 create mode 100644 app/src/hooks/useChatMessageStyles.ts

diff --git a/app/src/hooks/useChatMessageStyles.ts b/app/src/hooks/useChatMessageStyles.ts
new file mode 100644
index 0000000000..3ccfb43b54
--- /dev/null
+++ b/app/src/hooks/useChatMessageStyles.ts
@@ -0,0 +1,35 @@
+import { useMemo } from "react";
+
+import { ViewStyleProps } from "@arizeai/components";
+
+export function useChatMessageStyles(
+  role: string
+): Pick<ViewStyleProps, "backgroundColor" | "borderColor"> {
+  return useMemo<ViewStyleProps>(() => {
+    if (role === "user" || role === "human") {
+      return {
+        backgroundColor: "grey-100",
+        borderColor: "grey-500",
+      };
+    } else if (role === "assistant" || role === "ai") {
+      return {
+        backgroundColor: "blue-100",
+        borderColor: "blue-700",
+      };
+    } else if (role === "system") {
+      return {
+        backgroundColor: "indigo-100",
+        borderColor: "indigo-700",
+      };
+    } else if (["function", "tool"].includes(role)) {
+      return {
+        backgroundColor: "yellow-100",
+        borderColor: "yellow-700",
+      };
+    }
+    return {
+      backgroundColor: "grey-100",
+      borderColor: "grey-700",
+    };
+  }, [role]);
+}
diff --git a/app/src/pages/playground/Playground.tsx b/app/src/pages/playground/Playground.tsx
index 0b45b1cd1e..2cd0dd93a0 100644
--- a/app/src/pages/playground/Playground.tsx
+++ b/app/src/pages/playground/Playground.tsx
@@ -1,4 +1,4 @@
-import React from "react";
+import React, { Fragment } from "react";
 import { Panel, PanelGroup, PanelResizeHandle } from "react-resizable-panels";
 
 import { Button, Flex, Heading, View } from "@arizeai/components";
@@ -48,12 +48,12 @@ function PlaygroundInstances() {
     <Flex direction="row" alignItems="stretch" height="100%">
       <PanelGroup direction="horizontal">
         {instances.map((instance, i) => (
-          <>
+          <Fragment key={i}>
             {i !== 0 && <PanelResizeHandle css={resizeHandleCSS} />}
             <Panel defaultSize={50}>
               <PlaygroundInstance key={i} playgroundInstanceId={instance.id} />
             </Panel>
-          </>
+          </Fragment>
         ))}
       </PanelGroup>
     </Flex>
diff --git a/app/src/pages/playground/PlaygroundChatTemplate.tsx b/app/src/pages/playground/PlaygroundChatTemplate.tsx
index 7263eb7b47..021d3f9516 100644
--- a/app/src/pages/playground/PlaygroundChatTemplate.tsx
+++ b/app/src/pages/playground/PlaygroundChatTemplate.tsx
@@ -1,9 +1,10 @@
-import React from "react";
+import React, { PropsWithChildren } from "react";
 import { css } from "@emotion/react";
 
-import { Card, TextArea } from "@arizeai/components";
+import { Card, CardProps, TextArea } from "@arizeai/components";
 
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
+import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";
 
 import { MessageRolePicker } from "./MessageRolePicker";
 import { PlaygroundInstanceProps } from "./types";
@@ -35,7 +36,8 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
       {template.messages.map((message, index) => {
         return (
           <li key={index}>
-            <Card
+            <ChatMessageCard
+              role={message.role}
               title={
                 <MessageRolePicker
                   includeLabel={false}
@@ -55,31 +57,58 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
                   }}
                 />
               }
-              variant="compact"
-              backgroundColor="light"
-              borderColor="light"
             >
-              <TextArea
-                height={100}
-                value={message.content}
-                onChange={(val) => {
-                  updateInstance({
-                    instanceId: id,
-                    patch: {
-                      template: {
-                        __type: "chat",
-                        messages: template.messages.map((message, i) =>
-                          i === index ? { ...message, content: val } : message
-                        ),
+              <div
+                css={css`
+                  // TODO: remove these styles once the codemiror editor is added
+                  .ac-textfield {
+                    border: none !important;
+                    border-radius: 0;
+                    textarea {
+                      padding: var(--ac-global-dimension-size-200);
+                    }
+                  }
+                `}
+              >
+                <TextArea
+                  value={message.content}
+                  height={200}
+                  variant="quiet"
+                  onChange={(val) => {
+                    updateInstance({
+                      instanceId: id,
+                      patch: {
+                        template: {
+                          __type: "chat",
+                          messages: template.messages.map((message, i) =>
+                            i === index ? { ...message, content: val } : message
+                          ),
+                        },
                       },
-                    },
-                  });
-                }}
-              />
-            </Card>
+                    });
+                  }}
+                />
+              </div>
+            </ChatMessageCard>
           </li>
         );
       })}
     </ul>
   );
 }
+
+function ChatMessageCard(
+  props: PropsWithChildren<{ title: CardProps["title"]; role: string }>
+) {
+  const styles = useChatMessageStyles(props.role);
+  return (
+    <Card
+      title={props.title}
+      variant="compact"
+      {...styles}
+      bodyStyle={{ padding: 0 }}
+    >
+      {props.children}
+    </Card>
+  );
+}
diff --git a/app/src/pages/playground/SpanPlaygroundPage.tsx b/app/src/pages/playground/SpanPlaygroundPage.tsx
index 8a34f996fb..3dbdf58b9a 100644
--- a/app/src/pages/playground/SpanPlaygroundPage.tsx
+++ b/app/src/pages/playground/SpanPlaygroundPage.tsx
@@ -1,6 +1,8 @@
 import React, { useMemo } from "react";
 import { useLoaderData } from "react-router";
 
+import { createPlaygroundInstance } from "@phoenix/store";
+
 import { spanPlaygroundPageLoaderQuery$data } from "./__generated__/spanPlaygroundPageLoaderQuery.graphql";
 import { Playground } from "./Playground";
 import { transformSpanAttributesToPlaygroundInstance } from "./playgroundUtils";
@@ -25,7 +27,11 @@ export function SpanPlaygroundPage() {
 
   return (
     <Playground
-      instances={playgroundInstance != null ? [playgroundInstance] : undefined}
+      instances={
+        playgroundInstance != null
+          ? [playgroundInstance]
+          : [createPlaygroundInstance()]
+      }
     />
   );
 }
diff --git a/app/src/pages/trace/SpanDetails.tsx b/app/src/pages/trace/SpanDetails.tsx
index ab9af7fffc..6da594418a 100644
--- a/app/src/pages/trace/SpanDetails.tsx
+++ b/app/src/pages/trace/SpanDetails.tsx
@@ -40,7 +40,6 @@ import {
   TooltipTrigger,
   View,
   ViewProps,
-  ViewStyleProps,
 } from "@arizeai/components";
 import {
   DocumentAttributePostfixes,
@@ -65,6 +64,7 @@ import { SpanKindLabel } from "@phoenix/components/trace/SpanKindLabel";
 import { useNotifySuccess, useTheme } from "@phoenix/contexts";
 import { useFeatureFlag } from "@phoenix/contexts/FeatureFlagsContext";
 import { usePreferencesContext } from "@phoenix/contexts/PreferencesContext";
+import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";
 import {
   AttributeDocument,
   AttributeEmbedding,
@@ -1292,33 +1292,7 @@ function LLMMessage({ message }: { message: AttributeMessage }) {
     message[MessageAttributePostfixes.function_call_arguments_json] &&
     message[MessageAttributePostfixes.function_call_name];
   const role = message[MessageAttributePostfixes.role] || "unknown";
-  const messageStyles = useMemo<ViewStyleProps>(() => {
-    if (role === "user") {
-      return {
-        backgroundColor: "grey-100",
-        borderColor: "grey-500",
-      };
-    } else if (role === "assistant") {
-      return {
-        backgroundColor: "blue-100",
-        borderColor: "blue-700",
-      };
-    } else if (role === "system") {
-      return {
-        backgroundColor: "indigo-100",
-        borderColor: "indigo-700",
-      };
-    } else if (["function", "tool"].includes(role)) {
-      return {
-        backgroundColor: "yellow-100",
-        borderColor: "yellow-700",
-      };
-    }
-    return {
-      backgroundColor: "grey-100",
-      borderColor: "grey-700",
-    };
-  }, [role]);
+  const messageStyles = useChatMessageStyles(role);
 
   return (
     <MarkdownDisplayProvider>
diff --git a/app/src/store/playgroundStore.tsx b/app/src/store/playgroundStore.tsx
index af39ff8a12..368aa2b09c 100644
--- a/app/src/store/playgroundStore.tsx
+++ b/app/src/store/playgroundStore.tsx
@@ -193,6 +193,18 @@ const DEFAULT_TEXT_COMPLETION_TEMPLATE: PlaygroundTextCompletionTemplate = {
   prompt: "{{question}}",
 };
 
+export function createPlaygroundInstance(): PlaygroundInstance {
+  return {
+    id: generateInstanceId(),
+    template: DEFAULT_CHAT_COMPLETION_TEMPLATE,
+    tools: {},
+    input: { variables: {} },
+    output: {},
+    activeRunId: null,
+    isRunning: false,
+  };
+}
+
 export const createPlaygroundStore = (
   initialProps?: InitialPlaygroundState
 ) => {
@@ -200,17 +212,7 @@ export const createPlaygroundStore = (
     operationType: "chat",
     inputMode: "manual",
     setInputMode: (inputMode: PlaygroundInputMode) => set({ inputMode }),
-    instances: [
-      {
-        id: generateInstanceId(),
-        template: DEFAULT_CHAT_COMPLETION_TEMPLATE,
-        tools: {},
-        input: { variables: {} },
-        output: {},
-        activeRunId: null,
-        isRunning: false,
-      },
-    ],
+    instances: [createPlaygroundInstance()],
     setOperationType: (operationType: GenAIOperationType) => {
       if (operationType === "chat") {
         // TODO: this is incorrect, it should only change the template