feat(playground): show span details when playground run ends (#5108)

* feat: show span details

* cleanup

* fix

app/src/pages/playground/PlaygroundOutput.tsx

@@ -1,9 +1,9 @@
-import React, { useMemo, useState } from "react";
+import React, { Suspense, useMemo, useState } from "react";
 import { useSubscription } from "react-relay";
 import { graphql, GraphQLSubscriptionConfig } from "relay-runtime";
 import { css } from "@emotion/react";
 
-import { Card, Flex, Icon, Icons } from "@arizeai/components";
+import { Card, Flex, Icon, Icons, View } from "@arizeai/components";
 
 import { useNotifyError } from "@phoenix/contexts";
 import { useCredentialsContext } from "@phoenix/contexts/CredentialsContext";
@@ -22,6 +22,7 @@ import {
   PlaygroundOutputSubscription$variables,
 } from "./__generated__/PlaygroundOutputSubscription.graphql";
 import { isChatMessages } from "./playgroundUtils";
+import { RunMetadataFooter } from "./RunMetadataFooter";
 import { TitleWithAlphabeticIndex } from "./TitleWithAlphabeticIndex";
 import { PlaygroundInstanceProps } from "./types";
 import { useDerivedPlaygroundVariables } from "./useDerivedPlaygroundVariables";
@@ -111,8 +112,14 @@ export function PlaygroundOutput(props: PlaygroundOutputProps) {
       title={<TitleWithAlphabeticIndex index={index} title="Output" />}
       collapsible
       variant="compact"
+      bodyStyle={{ padding: 0 }}
     >
-      {OutputEl}
+      <View padding="size-200">{OutputEl}</View>
+      <Suspense>
+        {instance.spanId ? (
+          <RunMetadataFooter spanId={instance.spanId} />
+        ) : null}
+      </Suspense>
     </Card>
   );
 }
@@ -164,6 +171,11 @@ function useChatCompletionSubscription({
                 arguments
               }
             }
+            ... on FinishedChatCompletion {
+              span {
+                id
+              }
+            }
           }
         }
       `,
@@ -311,6 +323,13 @@ function PlaygroundOutputText(props: PlaygroundInstanceProps) {
           }
           return updated;
         });
+      } else if (chatCompletion.__typename === "FinishedChatCompletion") {
+        updateInstance({
+          instanceId: props.playgroundInstanceId,
+          patch: {
+            spanId: chatCompletion.span.id,
+          },
+        });
       }
     },
     onCompleted: () => {

app/src/pages/playground/RunMetadataFooter.tsx

@@ -0,0 +1,68 @@
+import React from "react";
+import { useLazyLoadQuery } from "react-relay";
+import { graphql } from "relay-runtime";
+
+import { Flex, View } from "@arizeai/components";
+
+import { ExternalLink } from "@phoenix/components";
+import { LatencyText } from "@phoenix/components/trace/LatencyText";
+import { TokenCount } from "@phoenix/components/trace/TokenCount";
+
+import { RunMetadataFooterQuery } from "./__generated__/RunMetadataFooterQuery.graphql";
+
+export function RunMetadataFooter({ spanId }: { spanId: string }) {
+  const data = useLazyLoadQuery<RunMetadataFooterQuery>(
+    graphql`
+      query RunMetadataFooterQuery($spanId: GlobalID!) {
+        span: node(id: $spanId) {
+          id
+          ... on Span {
+            project {
+              id
+            }
+            context {
+              traceId
+              spanId
+            }
+            tokenCountCompletion
+            tokenCountPrompt
+            tokenCountTotal
+            latencyMs
+          }
+        }
+      }
+    `,
+    { spanId }
+  );
+
+  if (!data.span) {
+    return null;
+  }
+
+  return (
+    <View
+      borderTopColor="light"
+      borderTopWidth="thin"
+      paddingStart="size-200"
+      paddingEnd="size-200"
+      paddingTop="size-100"
+      paddingBottom="size-100"
+    >
+      <Flex direction="row" gap="size-200" justifyContent="space-between">
+        <Flex direction="row" gap="size-100">
+          <TokenCount
+            tokenCountTotal={data.span.tokenCountTotal || 0}
+            tokenCountPrompt={data.span.tokenCountPrompt || 0}
+            tokenCountCompletion={data.span.tokenCountCompletion || 0}
+          />
+          <LatencyText latencyMs={data.span.latencyMs || 0} />
+        </Flex>
+        <ExternalLink
+          href={`/projects/${data.span?.project?.id}/traces/${data.span?.context?.traceId}?selectedSpanId=${data.span?.context?.spanId}`}
+        >
+          View Trace
+        </ExternalLink>
+      </Flex>
+    </View>
+  );
+}