feat(playground): provide a back to trace button from the span playgr… (

#4954) * feat(playground): provide a back to trace button from the span playground * fix * fix types in tests
Arize-ai · Oct 11, 2024 · 01227e2 · 01227e2
1 parent 2b6529f
commit 01227e2
Show file tree

Hide file tree

Showing 4 changed files with 139 additions and 37 deletions.
diff --git a/app/src/pages/playground/SpanPlaygroundPage.tsx b/app/src/pages/playground/SpanPlaygroundPage.tsx
@@ -1,5 +1,7 @@
-import React, { useMemo } from "react";
-import { useLoaderData } from "react-router";
+import React, { useMemo, useState } from "react";
+import { useLoaderData, useNavigate } from "react-router";
+
+import { Alert, Button, Flex, Icon, Icons } from "@arizeai/components";
 
 import { createPlaygroundInstance } from "@phoenix/store";
 
@@ -26,12 +28,55 @@ export function SpanPlaygroundPage() {
   );
 
   return (
-    <Playground
-      instances={
-        playgroundInstance != null
-          ? [playgroundInstance]
-          : [createPlaygroundInstance()]
-      }
-    />
+    <Flex direction="column" height="100%">
+      <SpanPlaygroundBanners span={span} />
+      <Playground
+        instances={
+          playgroundInstance != null
+            ? [playgroundInstance]
+            : [createPlaygroundInstance()]
+        }
+      />
+    </Flex>
+  );
+}
+
+function SpanPlaygroundBanners({
+  span,
+}: {
+  span: Extract<
+    NonNullable<spanPlaygroundPageLoaderQuery$data["span"]>,
+    { __typename: "Span" }
+  >;
+}) {
+  const navigate = useNavigate();
+  const [showBackBanner, setShowBackBanner] = useState(true);
+  return (
+    <div>
+      {showBackBanner && (
+        <Alert
+          variant="info"
+          title="LLM Span Replay"
+          banner
+          dismissable
+          onDismissClick={() => {
+            setShowBackBanner(false);
+          }}
+          extra={
+            <Button
+              variant="default"
+              icon={<Icon svg={<Icons.ArrowBack />} />}
+              onClick={() => {
+                navigate(
+                  `/projects/${span.project.id}/traces/${span.context.traceId}?selectedSpanNodeId=${span.id}`
+                );
+              }}
+            >
+              Back to Trace
+            </Button>
+          }
+        >{`Replay and iterate on your LLM call from your ${span.project.name} project`}</Alert>
+      )}
+    </div>
   );
 }
diff --git a/app/src/pages/playground/__generated__/spanPlaygroundPageLoaderQuery.graphql.ts b/app/src/pages/playground/__generated__/spanPlaygroundPageLoaderQuery.graphql.ts
diff --git a/app/src/pages/playground/__tests__/fixtures.ts b/app/src/pages/playground/__tests__/fixtures.ts
@@ -2,9 +2,15 @@ import { PlaygroundSpan } from "../spanPlaygroundPageLoader";
 
 export const basePlaygroundSpan: PlaygroundSpan = {
   __typename: "Span",
+  id: "fake-id",
   context: {
+    traceId: "test",
     spanId: "test",
   },
+  project: {
+    id: "test",
+    name: "test",
+  },
   attributes: "",
 };
 export const spanAttributesWithInputMessages = {

diff --git a/app/src/pages/playground/spanPlaygroundPageLoader.ts b/app/src/pages/playground/spanPlaygroundPageLoader.ts
@@ -29,8 +29,14 @@ export async function spanPlaygroundPageLoader(args: LoaderFunctionArgs) {
         span: node(id: $spanId) {
           __typename
           ... on Span {
+            id
+            project {
+              id
+              name
+            }
             context {
               spanId
+              traceId
             }
             attributes
           }