feat(playground): add / delete messages (#4947)

Arize-ai · Oct 10, 2024 · b65ddeb · b65ddeb
1 parent 104b9cf
commit b65ddeb
Show file tree

Hide file tree

Showing 5 changed files with 132 additions and 39 deletions.
diff --git a/app/src/components/dnd/DragHandle.tsx b/app/src/components/dnd/DragHandle.tsx
@@ -12,7 +12,7 @@ function DragHandle() {
       css={css`
         cursor: grab;
         background-color: var(--ac-global-color-grey-200);
-        border: 1px solid var(--ac-global-color-grey-500);
+        border: 1px solid var(--ac-global-color-grey-400);
         color: var(--ac-global-text-color-900);
         display: flex;
         align-items: center;

diff --git a/app/src/pages/playground/Playground.tsx b/app/src/pages/playground/Playground.tsx
@@ -17,35 +17,38 @@ import { PlaygroundRunButton } from "./PlaygroundRunButton";
 export function Playground(props: InitialPlaygroundState) {
   return (
     <PlaygroundProvider {...props}>
-      <View
-        borderBottomColor="dark"
-        borderBottomWidth="thin"
-        padding="size-200"
-      >
-        <Flex
-          direction="row"
-          justifyContent="space-between"
-          alignItems="center"
+      <Flex direction="column" height="100%">
+        <View
+          borderBottomColor="dark"
+          borderBottomWidth="thin"
+          padding="size-200"
+          flex="none"
         >
-          <Heading level={1}>Playground</Heading>
-          <Flex direction="row" gap="size-100" alignItems="center">
-            <PlaygroundInputTypeTypeRadioGroup />
-            <Button variant="default" size="compact">
-              API Keys
-            </Button>
-            <PlaygroundRunButton />
+          <Flex
+            direction="row"
+            justifyContent="space-between"
+            alignItems="center"
+          >
+            <Heading level={1}>Playground</Heading>
+            <Flex direction="row" gap="size-100" alignItems="center">
+              <PlaygroundInputTypeTypeRadioGroup />
+              <Button variant="default" size="compact">
+                API Keys
+              </Button>
+              <PlaygroundRunButton />
+            </Flex>
           </Flex>
-        </Flex>
-      </View>
-      <PlaygroundInstances />
+        </View>
+        <PlaygroundInstances />
+      </Flex>
     </PlaygroundProvider>
   );
 }
 
 function PlaygroundInstances() {
   const instances = usePlaygroundContext((state) => state.instances);
   return (
-    <Flex direction="row" alignItems="stretch" height="100%">
+    <Flex direction="row" alignItems="stretch" height="100%" flex="1 1 auto">
       <PanelGroup direction="horizontal">
         {instances.map((instance, i) => (
           <Fragment key={i}>

diff --git a/app/src/pages/playground/PlaygroundChatTemplate.tsx b/app/src/pages/playground/PlaygroundChatTemplate.tsx
@@ -4,14 +4,23 @@ import { DragDropProvider } from "@dnd-kit/react";
 import { useSortable } from "@dnd-kit/react/sortable";
 import { css } from "@emotion/react";
 
-import { Card, TextArea } from "@arizeai/components";
+import {
+  Button,
+  Card,
+  Flex,
+  Icon,
+  Icons,
+  TextArea,
+  View,
+} from "@arizeai/components";
 
 import { DragHandle } from "@phoenix/components/dnd/DragHandle";
 import { move } from "@phoenix/components/dnd/helpers/move";
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";
 import {
   ChatMessage,
+  generateMessageId,
   PlaygroundChatTemplate as PlaygroundChatTemplateType,
 } from "@phoenix/store";
 
@@ -79,6 +88,43 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
           );
         })}
       </ul>
+      <View
+        paddingStart="size-200"
+        paddingEnd="size-200"
+        paddingTop="size-100"
+        paddingBottom="size-100"
+        borderTopColor="dark"
+        borderTopWidth="thin"
+      >
+        <Flex direction="row" justifyContent="end">
+          <Button
+            variant="default"
+            aria-label="add message"
+            size="compact"
+            icon={<Icon svg={<Icons.PlusOutline />} />}
+            onClick={() => {
+              updateInstance({
+                instanceId: id,
+                patch: {
+                  template: {
+                    __type: "chat",
+                    messages: [
+                      ...template.messages,
+                      {
+                        id: generateMessageId(),
+                        role: "user",
+                        content: "",
+                      },
+                    ],
+                  },
+                },
+              });
+            }}
+          >
+            Message
+          </Button>
+        </Flex>
+      </View>
     </DragDropProvider>
   );
 }
@@ -128,7 +174,30 @@ function SortableMessageItem({
             }}
           />
         }
-        extra={<DragHandle ref={handleRef} />}
+        extra={
+          <Flex direction="row" gap="size-100">
+            <Button
+              aria-label="Delete message"
+              icon={<Icon svg={<Icons.TrashOutline />} />}
+              variant="default"
+              size="compact"
+              onClick={() => {
+                updateInstance({
+                  instanceId: playgroundInstanceId,
+                  patch: {
+                    template: {
+                      __type: "chat",
+                      messages: template.messages.filter(
+                        (msg) => msg.id !== message.id
+                      ),
+                    },
+                  },
+                });
+              }}
+            />
+            <DragHandle ref={handleRef} />
+          </Flex>
+        }
       >
         <div
           css={css`
@@ -141,6 +210,10 @@ function SortableMessageItem({
               }
             }
           `}
+          onKeyDownCapture={(e) => {
+            // Don't bubble up any keyboard events from the text area as to not interfere with DND
+            e.stopPropagation();
+          }}
         >
           <TextArea
             value={message.content}

diff --git a/app/src/pages/playground/PlaygroundInstance.tsx b/app/src/pages/playground/PlaygroundInstance.tsx
@@ -1,4 +1,4 @@
-import React from "react";
+import React, { PropsWithChildren } from "react";
 import { Panel, PanelGroup, PanelResizeHandle } from "react-resizable-panels";
 import { css } from "@emotion/react";
 
@@ -14,30 +14,47 @@ import { PlaygroundTemplate } from "./PlaygroundTemplate";
 import { PlaygroundTools } from "./PlaygroundTools";
 import { PlaygroundInstanceProps } from "./types";
 
-const panelContentCSS = css`
-  padding: var(--ac-global-dimension-size-200);
-  overflow: auto;
-  display: flex;
-  flex-direction: column;
-  gap: var(--ac-global-dimension-size-200);
-`;
-
 export function PlaygroundInstance(props: PlaygroundInstanceProps) {
   const numInstances = usePlaygroundContext((state) => state.instances.length);
   const isSingleInstance = numInstances == 1;
   return (
     <PanelGroup direction={isSingleInstance ? "horizontal" : "vertical"}>
-      <Panel defaultSize={50} order={1} css={panelContentCSS}>
-        <PlaygroundTemplate {...props} />
-        <PlaygroundTools />
+      <Panel defaultSize={50} order={1}>
+        <PanelContent>
+          <PlaygroundTemplate {...props} />
+          <PlaygroundTools />
+        </PanelContent>
       </Panel>
       <PanelResizeHandle
         css={isSingleInstance ? resizeHandleCSS : compactResizeHandleCSS}
       />
-      <Panel defaultSize={50} order={2} css={panelContentCSS}>
-        <PlaygroundInput />
-        <PlaygroundOutput {...props} />
+      <Panel defaultSize={50} order={2}>
+        <PanelContent>
+          <PlaygroundInput />
+          <PlaygroundOutput {...props} />
+        </PanelContent>
       </Panel>
     </PanelGroup>
   );
 }
+
+const PanelContent = (props: PropsWithChildren) => (
+  <div
+    css={css`
+      height: 100%;
+      padding: var(--ac-global-dimension-size-200);
+      overflow: auto;
+    `}
+  >
+    <div
+      css={css`
+        display: flex;
+        flex-direction: column;
+        gap: var(--ac-global-dimension-size-200);
+        padding-bottom: var(--ac-global-dimension-size-400);
+      `}
+    >
+      {props.children}
+    </div>
+  </div>
+);
diff --git a/app/src/pages/playground/PlaygroundTemplate.tsx b/app/src/pages/playground/PlaygroundTemplate.tsx
@@ -30,7 +30,7 @@ export function PlaygroundTemplate(props: PlaygroundTemplateProps) {
 
   return (
     <Card
-      title="Template"
+      title="Prompt"
       collapsible
       variant="compact"
       bodyStyle={{ padding: 0 }}