From 44002b78e61e9a700d01274966eaf52324e9793e Mon Sep 17 00:00:00 2001
From: Mikyo King <mikyo@arize.com>
Date: Thu, 10 Oct 2024 15:26:53 -0600
Subject: [PATCH 1/2] WIP

---
 app/src/components/AlphabeticIndexIcon.tsx    | 30 +++++++++++++++++++
 .../pages/playground/PlaygroundTemplate.tsx   | 10 ++++++-
 2 files changed, 39 insertions(+), 1 deletion(-)
 create mode 100644 app/src/components/AlphabeticIndexIcon.tsx

diff --git a/app/src/components/AlphabeticIndexIcon.tsx b/app/src/components/AlphabeticIndexIcon.tsx
new file mode 100644
index 0000000000..01e35a2254
--- /dev/null
+++ b/app/src/components/AlphabeticIndexIcon.tsx
@@ -0,0 +1,30 @@
+import React, { useMemo } from "react";
+import { transparentize } from "polished";
+import { css } from "@emotion/react";
+
+import { useWordColor } from "@phoenix/hooks/useWordColor";
+
+function indexToChar(index: number) {
+  // Wrap around using modulo if index exceeds 'C'
+  const charCode = 65 + (index % 3); // '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 = useWordColor(char);
+  const backgroundColor = useMemo(() => transparentize(0.5, color), [color]);
+  return (
+    <div
+      css={css`
+        color: ${color};
+        background-color: ${backgroundColor};
+        width: 24px;
+        height: 24px;
+        border-radius: var(--ac-global-rounding-small);
+      `}
+    >
+      {char}
+    </div>
+  );
+}
diff --git a/app/src/pages/playground/PlaygroundTemplate.tsx b/app/src/pages/playground/PlaygroundTemplate.tsx
index 643fca727b..495f6652ef 100644
--- a/app/src/pages/playground/PlaygroundTemplate.tsx
+++ b/app/src/pages/playground/PlaygroundTemplate.tsx
@@ -4,6 +4,7 @@ import {
   Button,
   Card,
   Content,
+  Flex,
   Icon,
   Icons,
   Tooltip,
@@ -11,6 +12,7 @@ import {
   TriggerWrap,
 } from "@arizeai/components";
 
+import { AlphabeticIndexIcon } from "@phoenix/components/AlphabeticIndexIcon";
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 
 import { NUM_MAX_PLAYGROUND_INSTANCES } from "./constants";
@@ -23,6 +25,7 @@ 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`);
   }
@@ -30,7 +33,12 @@ export function PlaygroundTemplate(props: PlaygroundTemplateProps) {
 
   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 }}

From b4747974d570485fb736479a19632cbe8d548d14 Mon Sep 17 00:00:00 2001
From: Mikyo King <mikyo@arize.com>
Date: Thu, 10 Oct 2024 16:02:11 -0600
Subject: [PATCH 2/2] feat(playground): allow up to 4 playground instances,
 distinguish alphabetically

---
 app/src/components/AlphabeticIndexIcon.tsx    | 13 +++++++----
 app/src/pages/playground/Playground.tsx       | 23 ++++++++++++++++++-
 .../pages/playground/PlaygroundInstance.tsx   |  2 +-
 app/src/pages/playground/PlaygroundOutput.tsx | 15 ++++++++++--
 .../pages/playground/PlaygroundTemplate.tsx   | 23 +------------------
 app/src/pages/playground/PlaygroundTools.tsx  | 19 +++++++++++++--
 .../playground/TitleWithAlphabeticIndex.tsx   | 23 +++++++++++++++++++
 app/src/pages/playground/constants.tsx        |  2 +-
 app/src/store/playgroundStore.tsx             | 10 ++++----
 9 files changed, 91 insertions(+), 39 deletions(-)
 create mode 100644 app/src/pages/playground/TitleWithAlphabeticIndex.tsx

diff --git a/app/src/components/AlphabeticIndexIcon.tsx b/app/src/components/AlphabeticIndexIcon.tsx
index 01e35a2254..a0d944cc41 100644
--- a/app/src/components/AlphabeticIndexIcon.tsx
+++ b/app/src/components/AlphabeticIndexIcon.tsx
@@ -1,19 +1,18 @@
 import React, { useMemo } from "react";
+import { schemeSet2 } from "d3-scale-chromatic";
 import { transparentize } from "polished";
 import { css } from "@emotion/react";
 
-import { useWordColor } from "@phoenix/hooks/useWordColor";
-
 function indexToChar(index: number) {
   // Wrap around using modulo if index exceeds 'C'
-  const charCode = 65 + (index % 3); // 'A' has ASCII code 65, 'B' is 66, 'C' is 67
+  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 = useWordColor(char);
-  const backgroundColor = useMemo(() => transparentize(0.5, color), [color]);
+  const color = useMemo(() => schemeSet2[index % 8], [index]);
+  const backgroundColor = useMemo(() => transparentize(0.8, color), [color]);
   return (
     <div
       css={css`
@@ -21,6 +20,10 @@ export function AlphabeticIndexIcon({ index }: { index: number }) {
         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);
       `}
     >
diff --git a/app/src/pages/playground/Playground.tsx b/app/src/pages/playground/Playground.tsx
index bed0f4edd2..14adc0962e 100644
--- a/app/src/pages/playground/Playground.tsx
+++ b/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 (
diff --git a/app/src/pages/playground/PlaygroundInstance.tsx b/app/src/pages/playground/PlaygroundInstance.tsx
index 37be6bc84e..8e6bf27de8 100644
--- a/app/src/pages/playground/PlaygroundInstance.tsx
+++ b/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
diff --git a/app/src/pages/playground/PlaygroundOutput.tsx b/app/src/pages/playground/PlaygroundOutput.tsx
index 2ff982c6b6..58f6e589c0 100644
--- a/app/src/pages/playground/PlaygroundOutput.tsx
+++ b/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} />
       ) : (
diff --git a/app/src/pages/playground/PlaygroundTemplate.tsx b/app/src/pages/playground/PlaygroundTemplate.tsx
index 495f6652ef..f5448d2ea5 100644
--- a/app/src/pages/playground/PlaygroundTemplate.tsx
+++ b/app/src/pages/playground/PlaygroundTemplate.tsx
@@ -15,7 +15,6 @@ import {
 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";
 
@@ -42,13 +41,7 @@ export function PlaygroundTemplate(props: PlaygroundTemplateProps) {
       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} />
@@ -59,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 (
diff --git a/app/src/pages/playground/PlaygroundTools.tsx b/app/src/pages/playground/PlaygroundTools.tsx
index b9e5d81bbe..805f4e1823 100644
--- a/app/src/pages/playground/PlaygroundTools.tsx
+++ b/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>
   );
diff --git a/app/src/pages/playground/TitleWithAlphabeticIndex.tsx b/app/src/pages/playground/TitleWithAlphabeticIndex.tsx
new file mode 100644
index 0000000000..af6bba26dc
--- /dev/null
+++ b/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>
+  );
+}
diff --git a/app/src/pages/playground/constants.tsx b/app/src/pages/playground/constants.tsx
index 90ea5993b0..a4b1bb2b59 100644
--- a/app/src/pages/playground/constants.tsx
+++ b/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";
 
diff --git a/app/src/store/playgroundStore.tsx b/app/src/store/playgroundStore.tsx
index 8c9222cd59..0aa0331301 100644
--- a/app/src/store/playgroundStore.tsx
+++ b/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,
           },