feat: Playground UI prototype (#4778)

* docs: initial playground PRD

* feat: playground prototype

* add rudamentary messages UI

* WIP

* routing

* add a playground button

* cleanup

* cleanup routes

* Update app/src/pages/playground/PlaygroundTemplate.tsx

* Update app/src/pages/playground/spanPlaygroundPageLoader.ts

* cleanup store

app/src/Routes.tsx

@@ -5,6 +5,7 @@ import { createBrowserRouter } from "react-router-dom";
 import { datasetLoaderQuery$data } from "./pages/dataset/__generated__/datasetLoaderQuery.graphql";
 import { embeddingLoaderQuery$data } from "./pages/embedding/__generated__/embeddingLoaderQuery.graphql";
 import { Layout } from "./pages/Layout";
+import { spanPlaygroundPageLoaderQuery$data } from "./pages/playground/__generated__/spanPlaygroundPageLoaderQuery.graphql";
 import { projectLoaderQuery$data } from "./pages/project/__generated__/projectLoaderQuery.graphql";
 import {
   APIsPage,
@@ -40,6 +41,8 @@ import {
   ResetPasswordPage,
   ResetPasswordWithTokenPage,
   SettingsPage,
+  SpanPlaygroundPage,
+  spanPlaygroundPageLoader,
   TracePage,
   TracingRoot,
 } from "./pages";
@@ -157,11 +160,25 @@ const router = createBrowserRouter(
           </Route>
           <Route
             path="/playground"
-            element={<PlaygroundPage />}
             handle={{
               crumb: () => "Playground",
             }}
-          />
+          >
+            <Route index element={<PlaygroundPage />} />
+            <Route
+              path="spans/:spanId" // TODO: Make it possible to go back to the span
+              element={<SpanPlaygroundPage />}
+              loader={spanPlaygroundPageLoader}
+              handle={{
+                crumb: (data: spanPlaygroundPageLoaderQuery$data) => {
+                  if (data.span.__typename === "Span") {
+                    return `span ${data.span.context.spanId}`;
+                  }
+                  return "span unknown";
+                },
+              }}
+            />
+          </Route>
           <Route
             path="/apis"
             element={<APIsPage />}

app/src/pages/playground/Playground.tsx

@@ -1,25 +1,16 @@
 import React from "react";
 import { Panel, PanelGroup, PanelResizeHandle } from "react-resizable-panels";
-import { css } from "@emotion/react";
 
 import { Button, Flex, Heading, View } from "@arizeai/components";
 
 import { resizeHandleCSS } from "@phoenix/components/resize";
-import { PlaygroundProvider } from "@phoenix/contexts/PlaygroundContext";
+import {
+  PlaygroundProvider,
+  usePlaygroundContext,
+} from "@phoenix/contexts/PlaygroundContext";
 
-import { PlaygroundInput } from "./PlaygroundInput";
+import { PlaygroundInstance } from "./PlaygroundInstance";
 import { PlaygroundOperationTypeRadioGroup } from "./PlaygroundOperationTypeRadioGroup";
-import { PlaygroundOutput } from "./PlaygroundOutput";
-import { PlaygroundTemplate } from "./PlaygroundTemplate";
-import { PlaygroundTools } from "./PlaygroundTools";
-
-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 Playground() {
   return (
@@ -39,17 +30,25 @@ export function Playground() {
           <Button variant="default">API Keys</Button>
         </Flex>
       </View>
+      <PlaygroundInstances />
+    </PlaygroundProvider>
+  );
+}
+
+function PlaygroundInstances() {
+  const instances = usePlaygroundContext((state) => state.instances);
+  return (
+    <Flex direction="row" alignItems="stretch" height="100%">
       <PanelGroup direction="horizontal">
-        <Panel defaultSize={50} order={1} css={panelContentCSS}>
-          <PlaygroundTemplate />
-          <PlaygroundTools />
-        </Panel>
-        <PanelResizeHandle css={resizeHandleCSS} />
-        <Panel defaultSize={50} order={2} css={panelContentCSS}>
-          <PlaygroundInput />
-          <PlaygroundOutput />
-        </Panel>
+        {instances.map((instance, i) => (
+          <>
+            {i !== 0 && <PanelResizeHandle css={resizeHandleCSS} />}
+            <Panel defaultSize={50}>
+              <PlaygroundInstance key={i} playgroundInstanceId={instance.id} />
+            </Panel>
+          </>
+        ))}
       </PanelGroup>
-    </PlaygroundProvider>
+    </Flex>
   );
 }

app/src/pages/playground/PlaygroundChatTemplate.tsx

@@ -0,0 +1,53 @@
+import React from "react";
+
+import { Card, TextArea } from "@arizeai/components";
+
+import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
+
+import { PlaygroundInstanceProps } from "./types";
+
+interface PlaygroundChatTemplateProps extends PlaygroundInstanceProps {}
+export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
+  const id = props.playgroundInstanceId;
+  // TODO: remove the hard coding of the first instance
+  const instances = usePlaygroundContext((state) => state.instances);
+  const updateInstance = usePlaygroundContext((state) => state.updateInstance);
+  const playground = instances.find((instance) => instance.id === id);
+  if (!playground) {
+    throw new Error(`Playground instance ${id} not found`);
+  }
+  const { template } = playground;
+  if (template.__type !== "chat") {
+    throw new Error(`Invalid template type ${template.__type}`);
+  }
+
+  return (
+    <ul>
+      {template.messages.map((message, index) => {
+        return (
+          <li key={index}>
+            <Card title={message.role} variant="compact">
+              <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
+                        ),
+                      },
+                    },
+                  });
+                }}
+              />
+            </Card>
+          </li>
+        );
+      })}
+    </ul>
+  );
+}

app/src/pages/playground/PlaygroundInput.tsx

@@ -2,9 +2,16 @@ import React from "react";
 
 import { Card } from "@arizeai/components";
 
+import { PlaygroundInputTypeTypeRadioGroup } from "./PlaygroundInputModeRadioGroup";
+
 export function PlaygroundInput() {
   return (
-    <Card title="Input" collapsible variant="compact">
+    <Card
+      title="Input"
+      collapsible
+      variant="compact"
+      extra={<PlaygroundInputTypeTypeRadioGroup />}
+    >
       Input goes here
     </Card>
   );

app/src/pages/playground/PlaygroundInputModeRadioGroup.tsx

@@ -0,0 +1,32 @@
+import React from "react";
+
+import { Radio, RadioGroup } from "@arizeai/components";
+
+import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
+
+/**
+ * A store connected radio group that toggles between manual and dataset input types.
+ */
+export function PlaygroundInputTypeTypeRadioGroup() {
+  const inputMode = usePlaygroundContext((state) => state.inputMode);
+  const setInputMode = usePlaygroundContext((state) => state.setInputMode);
+  return (
+    <RadioGroup
+      value={inputMode}
+      variant="inline-button"
+      size="compact"
+      onChange={(value) => {
+        if (value === "manual" || value === "dataset") {
+          setInputMode(value);
+        }
+      }}
+    >
+      <Radio label="manual" value={"manual"}>
+        Manual
+      </Radio>
+      <Radio label="Dataset" value={"dataset"}>
+        dataset
+      </Radio>
+    </RadioGroup>
+  );
+}

app/src/pages/playground/PlaygroundInstance.tsx

@@ -0,0 +1,43 @@
+import React from "react";
+import { Panel, PanelGroup, PanelResizeHandle } from "react-resizable-panels";
+import { css } from "@emotion/react";
+
+import {
+  compactResizeHandleCSS,
+  resizeHandleCSS,
+} from "@phoenix/components/resize";
+import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
+
+import { PlaygroundInput } from "./PlaygroundInput";
+import { PlaygroundOutput } from "./PlaygroundOutput";
+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>
+      <PanelResizeHandle
+        css={isSingleInstance ? resizeHandleCSS : compactResizeHandleCSS}
+      />
+      <Panel defaultSize={50} order={2} css={panelContentCSS}>
+        <PlaygroundInput />
+        <PlaygroundOutput />
+      </Panel>
+    </PanelGroup>
+  );
+}

app/src/pages/playground/PlaygroundTemplate.tsx

@@ -1,18 +1,70 @@
 import React from "react";
 
-import { Card } from "@arizeai/components";
+import { Button, Card, Icon, Icons } from "@arizeai/components";
 
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 
-export function PlaygroundTemplate() {
-  const operationType = usePlaygroundContext((state) => state.operationType);
+import { NUM_MAX_PLAYGROUND_INSTANCES } from "./constants";
+import { PlaygroundChatTemplate } from "./PlaygroundChatTemplate";
+import { PlaygroundInstanceProps } from "./types";
+
+interface PlaygroundTemplateProps extends PlaygroundInstanceProps {}
+
+export function PlaygroundTemplate(props: PlaygroundTemplateProps) {
+  const id = props.playgroundInstanceId;
+  const instances = usePlaygroundContext((state) => state.instances);
+  const playground = instances.find((instance) => instance.id === id);
+  if (!playground) {
+    throw new Error(`Playground instance ${id} not found`);
+  }
+  const { template } = playground;
+
   return (
-    <Card title="Template" collapsible variant="compact">
-      {operationType === "chat" ? (
-        <div>Chat Template goes here</div>
+    <Card
+      title="Template"
+      collapsible
+      variant="compact"
+      extra={
+        instances.length >= NUM_MAX_PLAYGROUND_INSTANCES ? (
+          <DeleteButton {...props} />
+        ) : (
+          <CompareButton />
+        )
+      }
+    >
+      {template.__type === "chat" ? (
+        <PlaygroundChatTemplate {...props} />
       ) : (
-        <div>Completion Template goes here</div>
+        "Completion Template"
       )}
     </Card>
   );
 }
+
+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 (
+    <Button
+      variant="default"
+      size="compact"
+      icon={<Icon svg={<Icons.TrashOutline />} />}
+      onClick={() => {
+        deleteInstance(props.playgroundInstanceId);
+      }}
+    />
+  );
+}

app/src/pages/playground/SpanPlaygroundPage.tsx

@@ -0,0 +1,10 @@
+import React from "react";
+
+// import { useLoaderData } from "react-router";
+import { Playground } from "./Playground";
+
+export function SpanPlaygroundPage() {
+  // const data = useLoaderData();
+
+  return <Playground />;
+}