feat(playground): allow up to 4 instances, distinguish alphabetically (…

…#4951)

* WIP

* feat(playground): allow up to 4 playground instances, distinguish alphabetically

app/src/components/AlphabeticIndexIcon.tsx

@@ -0,0 +1,33 @@
+import React, { useMemo } from "react";
+import { schemeSet2 } from "d3-scale-chromatic";
+import { transparentize } from "polished";
+import { css } from "@emotion/react";
+
+function indexToChar(index: number) {
+  // Wrap around using modulo if index exceeds 'C'
+  const charCode = 65 + index; // 'A' has ASCII code 65, 'B' is 66, 'C' is 67
+  return String.fromCharCode(charCode);
+}
+
+export function AlphabeticIndexIcon({ index }: { index: number }) {
+  const char = useMemo(() => indexToChar(index), [index]);
+  const color = useMemo(() => schemeSet2[index % 8], [index]);
+  const backgroundColor = useMemo(() => transparentize(0.8, color), [color]);
+  return (
+    <div
+      css={css`
+        color: ${color};
+        background-color: ${backgroundColor};
+        width: 24px;
+        height: 24px;
+        display: flex;
+        flex-direction: column;
+        justify-content: center;
+        align-items: center;
+        border-radius: var(--ac-global-rounding-small);
+      `}
+    >
+      {char}
+    </div>
+  );
+}

app/src/pages/playground/Playground.tsx

@@ -1,7 +1,7 @@
 import React, { Fragment } from "react";
 import { Panel, PanelGroup, PanelResizeHandle } from "react-resizable-panels";
 
-import { Button, Flex, Heading, View } from "@arizeai/components";
+import { Button, Flex, Heading, Icon, Icons, View } from "@arizeai/components";
 
 import { resizeHandleCSS } from "@phoenix/components/resize";
 import {
@@ -10,6 +10,7 @@ import {
 } from "@phoenix/contexts/PlaygroundContext";
 import { InitialPlaygroundState } from "@phoenix/store";
 
+import { NUM_MAX_PLAYGROUND_INSTANCES } from "./constants";
 import { PlaygroundInputTypeTypeRadioGroup } from "./PlaygroundInputModeRadioGroup";
 import { PlaygroundInstance } from "./PlaygroundInstance";
 import { PlaygroundRunButton } from "./PlaygroundRunButton";
@@ -35,6 +36,7 @@ export function Playground(props: InitialPlaygroundState) {
               <Button variant="default" size="compact">
                 API Keys
               </Button>
+              <AddPromptButton />
               <PlaygroundRunButton />
             </Flex>
           </Flex>
@@ -45,6 +47,25 @@ export function Playground(props: InitialPlaygroundState) {
   );
 }
 
+function AddPromptButton() {
+  const addInstance = usePlaygroundContext((state) => state.addInstance);
+  const numInstances = usePlaygroundContext((state) => state.instances.length);
+  return (
+    <Button
+      variant="default"
+      size="compact"
+      aria-label="add prompt"
+      icon={<Icon svg={<Icons.PlusCircleOutline />} />}
+      disabled={numInstances >= NUM_MAX_PLAYGROUND_INSTANCES}
+      onClick={() => {
+        addInstance();
+      }}
+    >
+      Prompt
+    </Button>
+  );
+}
+
 function PlaygroundInstances() {
   const instances = usePlaygroundContext((state) => state.instances);
   return (

app/src/pages/playground/PlaygroundInstance.tsx

@@ -22,7 +22,7 @@ export function PlaygroundInstance(props: PlaygroundInstanceProps) {
       <Panel defaultSize={50} order={1}>
         <PanelContent>
           <PlaygroundTemplate {...props} />
-          <PlaygroundTools />
+          <PlaygroundTools {...props} />
         </PanelContent>
       </Panel>
       <PanelResizeHandle

app/src/pages/playground/PlaygroundOutput.tsx

@@ -15,20 +15,31 @@ import {
   PlaygroundOutputSubscription$data,
   PlaygroundOutputSubscription$variables,
 } from "./__generated__/PlaygroundOutputSubscription.graphql";
+import { TitleWithAlphabeticIndex } from "./TitleWithAlphabeticIndex";
 import { PlaygroundInstanceProps } from "./types";
 
 interface PlaygroundOutputProps extends PlaygroundInstanceProps {}
 
 export function PlaygroundOutput(props: PlaygroundOutputProps) {
   const instanceId = props.playgroundInstanceId;
-  const instance = usePlaygroundContext((state) => state.instances[instanceId]);
+  const instance = usePlaygroundContext((state) =>
+    state.instances.find((instance) => instance.id === instanceId)
+  );
+  const index = usePlaygroundContext((state) =>
+    state.instances.findIndex((instance) => instance.id === instanceId)
+  );
   if (!instance) {
     return null;
   }
+
   const runId = instance.activeRunId;
   const hasRunId = runId !== null;
   return (
-    <Card title="Output" collapsible variant="compact">
+    <Card
+      title={<TitleWithAlphabeticIndex index={index} title="Output" />}
+      collapsible
+      variant="compact"
+    >
       {hasRunId ? (
         <PlaygroundOutputText key={runId} playgroundInstanceId={instanceId} />
       ) : (

app/src/pages/playground/PlaygroundTemplate.tsx

@@ -4,16 +4,17 @@ import {
   Button,
   Card,
   Content,
+  Flex,
   Icon,
   Icons,
   Tooltip,
   TooltipTrigger,
   TriggerWrap,
 } from "@arizeai/components";
 
+import { AlphabeticIndexIcon } from "@phoenix/components/AlphabeticIndexIcon";
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 
-import { NUM_MAX_PLAYGROUND_INSTANCES } from "./constants";
 import { PlaygroundChatTemplate } from "./PlaygroundChatTemplate";
 import { PlaygroundInstanceProps } from "./types";
 
@@ -23,24 +24,24 @@ export function PlaygroundTemplate(props: PlaygroundTemplateProps) {
   const instanceId = props.playgroundInstanceId;
   const instances = usePlaygroundContext((state) => state.instances);
   const instance = instances.find((instance) => instance.id === instanceId);
+  const index = instances.findIndex((instance) => instance.id === instanceId);
   if (!instance) {
     throw new Error(`Playground instance ${instanceId} not found`);
   }
   const { template } = instance;
 
   return (
     <Card
-      title="Prompt"
+      title={
+        <Flex direction="row" gap="size-100" alignItems="center">
+          <AlphabeticIndexIcon index={index} />
+          <span>Prompt</span>
+        </Flex>
+      }
       collapsible
       variant="compact"
       bodyStyle={{ padding: 0 }}
-      extra={
-        instances.length >= NUM_MAX_PLAYGROUND_INSTANCES ? (
-          <DeleteButton {...props} />
-        ) : (
-          <CompareButton />
-        )
-      }
+      extra={instances.length > 1 ? <DeleteButton {...props} /> : null}
     >
       {template.__type === "chat" ? (
         <PlaygroundChatTemplate {...props} />
@@ -51,20 +52,6 @@ export function PlaygroundTemplate(props: PlaygroundTemplateProps) {
   );
 }
 
-function CompareButton() {
-  const addInstance = usePlaygroundContext((state) => state.addInstance);
-  return (
-    <Button
-      variant="default"
-      size="compact"
-      icon={<Icon svg={<Icons.ArrowCompareOutline />} />}
-      onClick={() => {
-        addInstance();
-      }}
-    />
-  );
-}
-
 function DeleteButton(props: PlaygroundInstanceProps) {
   const deleteInstance = usePlaygroundContext((state) => state.deleteInstance);
   return (

app/src/pages/playground/PlaygroundTools.tsx

@@ -2,9 +2,24 @@ import React from "react";
 
 import { Card } from "@arizeai/components";
 
-export function PlaygroundTools() {
+import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
+
+import { TitleWithAlphabeticIndex } from "./TitleWithAlphabeticIndex";
+import { PlaygroundInstanceProps } from "./types";
+
+interface PlaygroundToolsProps extends PlaygroundInstanceProps {}
+export function PlaygroundTools(props: PlaygroundToolsProps) {
+  const index = usePlaygroundContext((state) =>
+    state.instances.findIndex(
+      (instance) => instance.id === props.playgroundInstanceId
+    )
+  );
   return (
-    <Card title="Tools" collapsible variant="compact">
+    <Card
+      title={<TitleWithAlphabeticIndex index={index} title="Tools" />}
+      collapsible
+      variant="compact"
+    >
       Tools go here
     </Card>
   );

app/src/pages/playground/TitleWithAlphabeticIndex.tsx

@@ -0,0 +1,23 @@
+import React from "react";
+
+import { Flex } from "@arizeai/components";
+
+import { AlphabeticIndexIcon } from "@phoenix/components/AlphabeticIndexIcon";
+
+/**
+ * Display the alphabetic index and title in a single line
+ */
+export function TitleWithAlphabeticIndex({
+  index,
+  title,
+}: {
+  index: number;
+  title: string;
+}) {
+  return (
+    <Flex direction="row" gap="size-100" alignItems="center">
+      <AlphabeticIndexIcon index={index} />
+      <span>{title}</span>
+    </Flex>
+  );
+}

app/src/pages/playground/constants.tsx

@@ -1,6 +1,6 @@
 import { ChatMessageRole } from "@phoenix/store";
 
-export const NUM_MAX_PLAYGROUND_INSTANCES = 2;
+export const NUM_MAX_PLAYGROUND_INSTANCES = 4;
 
 export const DEFAULT_CHAT_ROLE = "user";
 

app/src/store/playgroundStore.tsx

@@ -265,16 +265,16 @@ export const createPlaygroundStore = (
       set({ operationType });
     },
     addInstance: () => {
-      const instance = get().instances[0];
-      if (!instance) {
+      const instances = get().instances;
+      const firstInstance = get().instances[0];
+      if (!firstInstance) {
         return;
       }
-      // For now just hard-coded to two instances
       set({
         instances: [
-          instance,
+          ...instances,
           {
-            ...instance,
+            ...firstInstance,
             id: generateInstanceId(),
             activeRunId: null,
           },