feat: lazy load spans

Arize-ai · Jul 26, 2024 · ee74148 · ee74148
1 parent 9c15cf7
commit ee74148
Show file tree

Hide file tree

Showing 10 changed files with 816 additions and 553 deletions.
diff --git a/app/src/components/trace/TraceTree.tsx b/app/src/components/trace/TraceTree.tsx
@@ -13,12 +13,12 @@ import { createSpanTree, SpanTreeNode } from "./utils";
 
 type TraceTreeProps = {
   spans: ISpanItem[];
-  onSpanClick?: (spanId: string) => void;
-  selectedSpanId: string;
+  onSpanClick?: (span: ISpanItem) => void;
+  selectedSpanNodeId: string;
 };
 
 export function TraceTree(props: TraceTreeProps) {
-  const { spans, onSpanClick, selectedSpanId } = props;
+  const { spans, onSpanClick, selectedSpanNodeId } = props;
   const spanTree = createSpanTree(spans);
   return (
     <ul
@@ -31,10 +31,10 @@ export function TraceTree(props: TraceTreeProps) {
     >
       {spanTree.map((spanNode) => (
         <SpanTreeItem
-          key={spanNode.span.context.spanId}
+          key={spanNode.span.id}
           node={spanNode}
           onSpanClick={onSpanClick}
-          selectedSpanId={selectedSpanId}
+          selectedSpanNodeId={selectedSpanNodeId}
         />
       ))}
     </ul>
@@ -43,10 +43,10 @@ export function TraceTree(props: TraceTreeProps) {
 
 function SpanTreeItem<TSpan extends ISpanItem>(props: {
   node: SpanTreeNode<TSpan>;
-  selectedSpanId: string;
-  onSpanClick?: (spanId: string) => void;
+  selectedSpanNodeId: string;
+  onSpanClick?: (span: ISpanItem) => void;
 }) {
-  const { node, selectedSpanId, onSpanClick } = props;
+  const { node, selectedSpanNodeId, onSpanClick } = props;
   const childNodes = node.children;
   return (
     <div>
@@ -58,11 +58,11 @@ function SpanTreeItem<TSpan extends ISpanItem>(props: {
         `}
         onClick={() => {
           startTransition(() => {
-            onSpanClick && onSpanClick(node.span.context.spanId);
+            onSpanClick && onSpanClick(node.span);
           });
         }}
       >
-        <SpanNodeWrap isSelected={selectedSpanId === node.span.context.spanId}>
+        <SpanNodeWrap isSelected={selectedSpanNodeId === node.span.id}>
           <Flex
             direction="row"
             gap="size-100"
@@ -102,7 +102,7 @@ function SpanTreeItem<TSpan extends ISpanItem>(props: {
                 <SpanTreeItem
                   node={leafNode}
                   onSpanClick={onSpanClick}
-                  selectedSpanId={selectedSpanId}
+                  selectedSpanNodeId={selectedSpanNodeId}
                 />
               </div>
             );

diff --git a/app/src/components/trace/types.ts b/app/src/components/trace/types.ts
@@ -2,6 +2,7 @@
  * A generic interface for a span to be re-used as a constraint
  */
 export interface ISpanItem {
+  id: string;
   name: string;
   spanKind: string;
   statusCode: SpanStatusCodeType;

diff --git a/app/src/pages/example/ExampleDetailsDialog.tsx b/app/src/pages/example/ExampleDetailsDialog.tsx
@@ -47,8 +47,8 @@ export function ExampleDetailsDialog({
               metadata
             }
             span {
+              id
               context {
-                spanId
                 traceId
               }
               project {
@@ -77,7 +77,7 @@ export function ExampleDetailsDialog({
       return null;
     }
     return {
-      spanId: sourceSpan.context.spanId,
+      id: sourceSpan.id,
       traceId: sourceSpan.context.traceId,
       projectId: sourceSpan.project.id,
     };
@@ -98,7 +98,7 @@ export function ExampleDetailsDialog({
                 size="compact"
                 onClick={() => {
                   navigate(
-                    `/projects/${sourceSpanInfo.projectId}/traces/${sourceSpanInfo.traceId}?selectedSpanId=${sourceSpanInfo.spanId}`
+                    `/projects/${sourceSpanInfo.projectId}/traces/${sourceSpanInfo.traceId}?selectedSpanNodeId=${sourceSpanInfo.id}`
                   );
                 }}
               >

diff --git a/app/src/pages/example/__generated__/ExampleDetailsDialogQuery.graphql.ts b/app/src/pages/example/__generated__/ExampleDetailsDialogQuery.graphql.ts
diff --git a/app/src/pages/project/SpansTable.tsx b/app/src/pages/project/SpansTable.tsx
@@ -264,9 +264,10 @@ export function SpansTable(props: SpansTableProps) {
       accessorKey: "name",
       enableSorting: false,
       cell: ({ getValue, row }) => {
-        const { spanId, traceId } = row.original.context;
+        const span = row.original;
+        const { traceId } = span.context;
         return (
-          <Link to={`traces/${traceId}?selectedSpanId=${spanId}`}>
+          <Link to={`traces/${traceId}?selectedSpanNodeId=${span.id}`}>
             {getValue() as string}
           </Link>
         );
@@ -470,7 +471,7 @@ export function SpansTable(props: SpansTableProps) {
                     key={row.id}
                     onClick={() =>
                       navigate(
-                        `traces/${row.original.context.traceId}?selectedSpanId=${row.original.context.spanId}`
+                        `traces/${row.original.context.traceId}?selectedSpanNodeId=${row.original.id}`
                       )
                     }
                   >

diff --git a/app/src/pages/project/TracesTable.tsx b/app/src/pages/project/TracesTable.tsx
@@ -372,9 +372,9 @@ export function TracesTable(props: TracesTableProps) {
       accessorKey: "name",
       enableSorting: false,
       cell: ({ getValue, row }) => {
-        const { spanId, traceId } = row.original.context;
+        const { traceId } = row.original.context;
         return (
-          <Link to={`traces/${traceId}?selectedSpanId=${spanId}`}>
+          <Link to={`traces/${traceId}?selectedSpanNodeId=${row.original.id}`}>
             {getValue() as string}
           </Link>
         );

diff --git a/app/src/pages/trace/SpanDetails.tsx b/app/src/pages/trace/SpanDetails.tsx
@@ -6,6 +6,7 @@ import React, {
   useMemo,
   useState,
 } from "react";
+import { graphql, useLazyLoadQuery } from "react-relay";
 import { useNavigate } from "react-router";
 import { json } from "@codemirror/lang-json";
 import { nord } from "@uiw/codemirror-theme-nord";
@@ -80,8 +81,9 @@ import { RetrievalEvaluationLabel } from "../project/RetrievalEvaluationLabel";
 
 import {
   MimeType,
-  TraceDetailsQuery$data,
-} from "./__generated__/TraceDetailsQuery.graphql";
+  SpanDetailsQuery,
+  SpanDetailsQuery$data,
+} from "./__generated__/SpanDetailsQuery.graphql";
 import { EditSpanAnnotationsButton } from "./EditSpanAnnotationsButton";
 import { SpanCodeDropdown } from "./SpanCodeDropdown";
 import { SpanEvaluationsTable } from "./SpanEvaluationsTable";
@@ -98,11 +100,9 @@ type AttributeObject = {
   [SemanticAttributePrefixes.llm]?: AttributeLlm;
 };
 
-type Span = NonNullable<
-  TraceDetailsQuery$data["project"]["trace"]
->["spans"]["edges"][number]["span"];
+type Span = Extract<SpanDetailsQuery$data["span"], { __typename: "Span" }>;
 
-type DocumentEvaluation = Span["documentEvaluations"][number];
+type DocumentEvaluation = NonNullable<Span["documentEvaluations"]>[number];
 
 /**
  * Hook that safely parses a JSON string.
@@ -134,15 +134,87 @@ const defaultCardProps: Partial<CardProps> = {
 };
 
 export function SpanDetails({
-  selectedSpan,
+  spanNodeId,
   projectId,
 }: {
-  selectedSpan: Span;
+  /**
+   * The Global ID of the span
+   */
+  spanNodeId: string;
   projectId: string;
 }) {
+  const { span } = useLazyLoadQuery<SpanDetailsQuery>(
+    graphql`
+      query SpanDetailsQuery($spanId: GlobalID!) {
+        span: node(id: $spanId) {
+          __typename
+          ... on Span {
+            id
+            context {
+              spanId
+              traceId
+            }
+            name
+            spanKind
+            statusCode: propagatedStatusCode
+            statusMessage
+            startTime
+            parentId
+            latencyMs
+            tokenCountTotal
+            tokenCountPrompt
+            tokenCountCompletion
+            input {
+              value
+              mimeType
+            }
+            output {
+              value
+              mimeType
+            }
+            attributes
+            events @required(action: THROW) {
+              name
+              message
+              timestamp
+            }
+            spanEvaluations {
+              name
+              label
+              score
+            }
+            documentRetrievalMetrics {
+              evaluationName
+              ndcg
+              precision
+              hit
+            }
+            documentEvaluations {
+              documentPosition
+              name
+              label
+              score
+              explanation
+            }
+            ...SpanEvaluationsTable_evals
+          }
+        }
+      }
+    `,
+    {
+      spanId: spanNodeId,
+    }
+  );
+
+  if (span.__typename !== "Span") {
+    throw new Error(
+      "Expected a span, but got a different type" + span.__typename
+    );
+  }
+
   const hasExceptions = useMemo<boolean>(() => {
-    return spanHasException(selectedSpan);
-  }, [selectedSpan]);
+    return spanHasException(span);
+  }, [span]);
   const showAnnotations = useFeatureFlag("annotations");
   return (
     <Flex direction="column" flex="1 1 auto" height="100%">
@@ -159,16 +231,16 @@ export function SpanDetails({
           justifyContent="space-between"
           alignItems="center"
         >
-          <SpanItem {...selectedSpan} />
+          <SpanItem {...span} />
           <Flex flex="none" direction="row" alignItems="center" gap="size-100">
             <SpanCodeDropdown
-              traceId={selectedSpan.context.traceId}
-              spanId={selectedSpan.context.spanId}
+              traceId={span.context.traceId}
+              spanId={span.context.spanId}
             />
-            <AddSpanToDatasetButton span={selectedSpan} />
+            <AddSpanToDatasetButton span={span} />
             {showAnnotations ? (
               <EditSpanAnnotationsButton
-                spanNodeId={selectedSpan.id}
+                spanNodeId={span.id}
                 projectId={projectId}
               />
             ) : null}
@@ -177,18 +249,16 @@ export function SpanDetails({
       </View>
       <Tabs>
         <TabPane name={"Info"}>
-          <SpanInfo span={selectedSpan} />
+          <SpanInfo span={span} />
         </TabPane>
         <TabPane
           name={"Evaluations"}
           extra={
-            <Counter variant={"light"}>
-              {selectedSpan.spanEvaluations.length}
-            </Counter>
+            <Counter variant={"light"}>{span.spanEvaluations.length}</Counter>
           }
         >
           {(selected) => {
-            return selected ? <SpanEvaluations span={selectedSpan} /> : null;
+            return selected ? <SpanEvaluations span={span} /> : null;
           }}
         </TabPane>
         <TabPane name={"Attributes"} title="Attributes">
@@ -197,22 +267,22 @@ export function SpanDetails({
               title="All Attributes"
               {...defaultCardProps}
               titleExtra={attributesContextualHelp}
-              extra={<CopyToClipboardButton text={selectedSpan.attributes} />}
+              extra={<CopyToClipboardButton text={span.attributes} />}
               bodyStyle={{ padding: 0 }}
             >
-              <JSONBlock>{selectedSpan.attributes}</JSONBlock>
+              <JSONBlock>{span.attributes}</JSONBlock>
             </Card>
           </View>
         </TabPane>
         <TabPane
           name={"Events"}
           extra={
             <Counter variant={hasExceptions ? "danger" : "light"}>
-              {selectedSpan.events.length}
+              {span.events.length}
             </Counter>
           }
         >
-          <SpanEventsList events={selectedSpan.events} />
+          <SpanEventsList events={span.events} />
         </TabPane>
       </Tabs>
     </Flex>