feat(sesions): add session details page (#5225)

* feat(sesions): add session details page * remove unused import, add error boundary * pull user from root span
Arize-ai · Oct 31, 2024 · ad6caf2 · ad6caf2
1 parent b8f38fc
commit ad6caf2
Show file tree

Hide file tree

Showing 11 changed files with 744 additions and 58 deletions.
diff --git a/app/src/Routes.tsx b/app/src/Routes.tsx
@@ -7,6 +7,7 @@ import { embeddingLoaderQuery$data } from "./pages/embedding/__generated__/embed
 import { Layout } from "./pages/Layout";
 import { spanPlaygroundPageLoaderQuery$data } from "./pages/playground/__generated__/spanPlaygroundPageLoaderQuery.graphql";
 import { projectLoaderQuery$data } from "./pages/project/__generated__/projectLoaderQuery.graphql";
+import { sessionLoader } from "./pages/trace/sessionLoader";
 import { SessionPage } from "./pages/trace/SessionPage";
 import {
   APIsPage,
@@ -118,7 +119,11 @@ const router = createBrowserRouter(
               <Route index element={<ProjectPage />} />
               <Route element={<ProjectPage />}>
                 <Route path="traces/:traceId" element={<TracePage />} />
-                <Route path="sessions/:sessionId" element={<SessionPage />} />
+                <Route
+                  path="sessions/:sessionId"
+                  element={<SessionPage />}
+                  loader={sessionLoader}
+                />
               </Route>
             </Route>
           </Route>

diff --git a/app/src/components/table/TimestampCell.tsx b/app/src/components/table/TimestampCell.tsx
@@ -11,7 +11,7 @@ export function TimestampCell<TData extends object, TValue>({
 }: CellContext<TData, TValue>) {
   const value = getValue();
   if (!isStringOrNull(value)) {
-    throw new Error("TimestampCell only supports number or null values.");
+    throw new Error("TimestampCell only supports string or null values.");
   }
   const timestamp =
     value != null

diff --git a/app/src/pages/trace/EditSpanAnnotationsButton.tsx b/app/src/pages/trace/EditSpanAnnotationsButton.tsx
@@ -1,19 +1,34 @@
 import React, { ReactNode, useState } from "react";
 
-import { Button, DialogContainer, Icon, Icons } from "@arizeai/components";
+import {
+  Button,
+  ButtonProps,
+  DialogContainer,
+  Icon,
+  Icons,
+} from "@arizeai/components";
 
 import { EditSpanAnnotationsDialog } from "@phoenix/components/trace/EditSpanAnnotationsDialog";
 
-export function EditSpanAnnotationsButton(props: {
+export function EditSpanAnnotationsButton({
+  spanNodeId,
+  projectId,
+  size = "default",
+}: {
   spanNodeId: string;
   projectId: string;
+  /**
+   * The size of the button
+   * @default default
+   */
+  size?: ButtonProps["size"];
 }) {
-  const { spanNodeId, projectId } = props;
   const [dialog, setDialog] = useState<ReactNode>(null);
   return (
     <>
       <Button
         variant="default"
+        size={size}
         icon={<Icon svg={<Icons.EditOutline />} />}
         onClick={() =>
           setDialog(

diff --git a/app/src/pages/trace/SessionDetails.tsx b/app/src/pages/trace/SessionDetails.tsx
@@ -1,8 +1,54 @@
-import React, { useMemo } from "react";
+import React from "react";
 import { graphql, useLazyLoadQuery } from "react-relay";
 import { css } from "@emotion/react";
 
-import { SessionDetailsQuery } from "./__generated__/SessionDetailsQuery.graphql";
+import { Flex, Text, View } from "@arizeai/components";
+
+import { TokenCount } from "@phoenix/components/trace/TokenCount";
+
+import {
+  SessionDetailsQuery,
+  SessionDetailsQuery$data,
+} from "./__generated__/SessionDetailsQuery.graphql";
+import { SessionDetailsTraceList } from "./SessionDetailsTraceList";
+
+function SessionDetailsHeader({
+  traceCount,
+  tokenUsage,
+}: {
+  traceCount: number;
+  tokenUsage?: NonNullable<SessionDetailsQuery$data["session"]>["tokenUsage"];
+}) {
+  return (
+    <View
+      padding={"size-200"}
+      borderBottomWidth={"thin"}
+      borderBottomColor={"dark"}
+    >
+      <Flex direction={"row"} gap={"size-400"}>
+        <Flex direction={"column"}>
+          <Text elementType={"h3"} textSize={"medium"} color={"text-700"}>
+            Traces Count
+          </Text>
+          <Text textSize={"xlarge"}>{traceCount}</Text>
+        </Flex>
+        {tokenUsage != null ? (
+          <Flex direction={"column"}>
+            <Text elementType={"h3"} textSize={"medium"} color={"text-700"}>
+              Total Tokens
+            </Text>
+            <TokenCount
+              tokenCountTotal={tokenUsage.total}
+              tokenCountCompletion={tokenUsage.completion}
+              tokenCountPrompt={tokenUsage.prompt}
+              textSize={"xlarge"}
+            />
+          </Flex>
+        ) : null}
+      </Flex>
+    </View>
+  );
+}
 
 export type SessionDetailsProps = {
   sessionId: string;
@@ -18,16 +64,45 @@ export function SessionDetails(props: SessionDetailsProps) {
       query SessionDetailsQuery($id: GlobalID!) {
         session: node(id: $id) {
           ... on ProjectSession {
+            numTraces
+            durationS
+            tokenUsage {
+              total
+              completion
+              prompt
+            }
+            sessionId
             traces {
               edges {
                 trace: node {
                   rootSpan {
+                    id
+                    attributes
+                    project {
+                      id
+                    }
                     input {
                       value
                     }
                     output {
                       value
                     }
+                    cumulativeTokenCountTotal
+                    cumulativeTokenCountCompletion
+                    cumulativeTokenCountPrompt
+                    latencyMs
+                    startTime
+                    spanAnnotations {
+                      name
+                      label
+                      score
+                      explanation
+                      annotatorKind
+                    }
+                    context {
+                      traceId
+                      spanId
+                    }
                   }
                 }
               }
@@ -43,37 +118,25 @@ export function SessionDetails(props: SessionDetailsProps) {
       fetchPolicy: "store-and-network",
     }
   );
-  const spansList = useMemo(() => {
-    const gqlSpans = data.session?.traces?.edges || [];
-    return gqlSpans.map(({ trace }) => trace);
-  }, [data]);
+
+  if (data.session == null) {
+    throw new Error("Session not found");
+  }
   return (
     <main
       css={css`
         flex: 1 1 auto;
-        overflow: hidden;
         display: flex;
         flex-direction: column;
+        height: 100%;
+        overflow: hidden;
       `}
     >
-      <table border={1} cellPadding={20}>
-        <thead>
-          <tr>
-            <th>#</th>
-            <th>User</th>
-            <th>Assistant</th>
-          </tr>
-        </thead>
-        <tbody>
-          {spansList.map((trace, index) => (
-            <tr key={index}>
-              <td>{index + 1}</td>
-              <td>{trace.rootSpan?.input?.value}</td>
-              <td>{trace.rootSpan?.output?.value}</td>
-            </tr>
-          ))}
-        </tbody>
-      </table>
+      <SessionDetailsHeader
+        traceCount={data.session.numTraces ?? 0}
+        tokenUsage={data.session.tokenUsage}
+      />
+      <SessionDetailsTraceList traces={data.session.traces} />
     </main>
   );
 }