fix(traces): make column selector sync'd between tabs (#1816)

* fix(traces): make column selector sync'd between tabs * lint * reverse the aliasing so columns show/hide together
Arize-ai · Dec 1, 2023 · 631f3c8 · 631f3c8
1 parent 63ac7f9
commit 631f3c8
Show file tree

Hide file tree

Showing 11 changed files with 192 additions and 151 deletions.
diff --git a/app/src/contexts/TracingContext.tsx b/app/src/contexts/TracingContext.tsx
@@ -0,0 +1,35 @@
+import React, { createContext, PropsWithChildren, useRef } from "react";
+import { useZustand } from "use-zustand";
+
+import {
+  createTracingStore,
+  TracingProps,
+  TracingState,
+  TracingStore,
+} from "@phoenix/store/tracingStore";
+
+export const TracingContext = createContext<TracingStore | null>(null);
+
+export function TracingProvider({
+  children,
+  ...props
+}: PropsWithChildren<Partial<TracingProps>>) {
+  const storeRef = useRef<TracingStore>();
+  if (!storeRef.current) {
+    storeRef.current = createTracingStore(props);
+  }
+  return (
+    <TracingContext.Provider value={storeRef.current}>
+      {children}
+    </TracingContext.Provider>
+  );
+}
+
+export function useTracingContext<T>(
+  selector: (state: TracingState) => T,
+  equalityFn?: (left: T, right: T) => boolean
+): T {
+  const store = React.useContext(TracingContext);
+  if (!store) throw new Error("Missing TracingContext.Provider in the tree");
+  return useZustand(store, selector, equalityFn);
+}
diff --git a/app/src/pages/TracingRoot.tsx b/app/src/pages/TracingRoot.tsx
@@ -2,11 +2,14 @@ import React from "react";
 import { Outlet } from "react-router";
 
 import { StreamStateProvider } from "@phoenix/contexts/StreamStateContext";
+import { TracingProvider } from "@phoenix/contexts/TracingContext";
 
 export function TracingRoot() {
   return (
-    <StreamStateProvider>
-      <Outlet />
-    </StreamStateProvider>
+    <TracingProvider>
+      <StreamStateProvider>
+        <Outlet />
+      </StreamStateProvider>
+    </TracingProvider>
   );
 }
diff --git a/app/src/pages/tracing/SpanColumnSelector.tsx b/app/src/pages/tracing/SpanColumnSelector.tsx
@@ -1,9 +1,11 @@
 import React, { ChangeEvent, useCallback, useMemo } from "react";
-import { Column, VisibilityState } from "@tanstack/react-table";
+import { Column } from "@tanstack/react-table";
 import { css } from "@emotion/react";
 
 import { Dropdown, Flex, Icon, Icons, View } from "@arizeai/components";
 
+import { useTracingContext } from "@phoenix/contexts/TracingContext";
+
 const UN_HIDABLE_COLUMN_IDS = ["spanKind", "name"];
 
 type SpanColumnSelectorProps = {
@@ -14,14 +16,6 @@ type SpanColumnSelectorProps = {
    */
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   columns: Column<any>[];
-  /**
-   * Map of the column id to the visibility state
-   */
-  columnVisibility: VisibilityState;
-  /*
-   * Callback to set the visibility state of a column
-   */
-  onColumnVisibilityChange: (visibilityState: VisibilityState) => void;
 };
 
 export function SpanColumnSelector(props: SpanColumnSelectorProps) {
@@ -51,12 +45,12 @@ const columCheckboxItemCSS = css`
 `;
 
 function ColumnSelectorMenu(props: SpanColumnSelectorProps) {
-  const {
-    columns: propsColumns,
-    columnVisibility,
-    onColumnVisibilityChange,
-  } = props;
+  const { columns: propsColumns } = props;
 
+  const columnVisibility = useTracingContext((state) => state.columnVisibility);
+  const setColumnVisibility = useTracingContext(
+    (state) => state.setColumnVisibility
+  );
   const columns = useMemo(() => {
     return propsColumns.filter((column) => {
       return !UN_HIDABLE_COLUMN_IDS.includes(column.id);
@@ -74,9 +68,9 @@ function ColumnSelectorMenu(props: SpanColumnSelectorProps) {
   const onCheckboxChange = useCallback(
     (event: ChangeEvent<HTMLInputElement>) => {
       const { name, checked } = event.target;
-      onColumnVisibilityChange({ ...columnVisibility, [name]: checked });
+      setColumnVisibility({ ...columnVisibility, [name]: checked });
     },
-    [columnVisibility, onColumnVisibilityChange]
+    [columnVisibility, setColumnVisibility]
   );
 
   const onToggleAll = useCallback(
@@ -85,9 +79,9 @@ function ColumnSelectorMenu(props: SpanColumnSelectorProps) {
       const newVisibilityState = columns.reduce((acc, column) => {
         return { ...acc, [column.id]: checked };
       }, {});
-      onColumnVisibilityChange(newVisibilityState);
+      setColumnVisibility(newVisibilityState);
     },
-    [columns, onColumnVisibilityChange]
+    [columns, setColumnVisibility]
   );
 
   return (

diff --git a/app/src/pages/tracing/SpansTable.tsx b/app/src/pages/tracing/SpansTable.tsx
@@ -14,7 +14,6 @@ import {
   getSortedRowModel,
   SortingState,
   useReactTable,
-  VisibilityState,
 } from "@tanstack/react-table";
 import { css } from "@emotion/react";
 
@@ -29,6 +28,7 @@ import { LatencyText } from "@phoenix/components/trace/LatencyText";
 import { SpanKindLabel } from "@phoenix/components/trace/SpanKindLabel";
 import { SpanStatusCodeIcon } from "@phoenix/components/trace/SpanStatusCodeIcon";
 import { useStreamState } from "@phoenix/contexts/StreamStateContext";
+import { useTracingContext } from "@phoenix/contexts/TracingContext";
 
 import {
   SpansTable_spans$key,
@@ -57,7 +57,7 @@ export function SpansTable(props: SpansTableProps) {
   const tableContainerRef = useRef<HTMLDivElement>(null);
   const [sorting, setSorting] = useState<SortingState>([]);
   const [filterCondition, setFilterCondition] = useState<string>("");
-  const [columnVisibility, setColumnVisibility] = useState<VisibilityState>({});
+  const columnVisibility = useTracingContext((state) => state.columnVisibility);
   const navigate = useNavigate();
   const { data, loadNext, hasNext, isLoadingNext, refetch } =
     usePaginationFragment<SpansTableSpansQuery, SpansTable_spans$key>(
@@ -235,6 +235,7 @@ export function SpansTable(props: SpansTableProps) {
     data: tableData,
     state: {
       sorting,
+      columnVisibility,
     },
     onSortingChange: setSorting,
     getCoreRowModel: getCoreRowModel(),
@@ -255,11 +256,7 @@ export function SpansTable(props: SpansTableProps) {
       >
         <Flex direction="row" gap="size-100" width="100%" alignItems="center">
           <SpanFilterConditionField onValidCondition={setFilterCondition} />
-          <SpanColumnSelector
-            columns={computedColumns}
-            columnVisibility={columnVisibility}
-            onColumnVisibilityChange={setColumnVisibility}
-          />
+          <SpanColumnSelector columns={computedColumns} />
         </Flex>
       </View>
       <div

diff --git a/app/src/pages/tracing/TracesTable.tsx b/app/src/pages/tracing/TracesTable.tsx
@@ -17,7 +17,6 @@ import {
   getSortedRowModel,
   SortingState,
   useReactTable,
-  VisibilityState,
 } from "@tanstack/react-table";
 import { css } from "@emotion/react";
 
@@ -35,6 +34,7 @@ import { SpanStatusCodeIcon } from "@phoenix/components/trace/SpanStatusCodeIcon
 import { ISpanItem } from "@phoenix/components/trace/types";
 import { createSpanTree, SpanTreeNode } from "@phoenix/components/trace/utils";
 import { useStreamState } from "@phoenix/contexts/StreamStateContext";
+import { useTracingContext } from "@phoenix/contexts/TracingContext";
 
 import {
   SpanStatusCode,
@@ -120,9 +120,9 @@ export function TracesTable(props: TracesTableProps) {
                 statusCode
                 startTime
                 latencyMs
-                cumulativeTokenCountTotal
-                cumulativeTokenCountPrompt
-                cumulativeTokenCountCompletion
+                tokenCountTotal: cumulativeTokenCountTotal
+                tokenCountPrompt: cumulativeTokenCountPrompt
+                tokenCountCompletion: cumulativeTokenCountCompletion
                 parentId
                 input {
                   value
@@ -141,9 +141,9 @@ export function TracesTable(props: TracesTableProps) {
                   startTime
                   latencyMs
                   parentId
-                  cumulativeTokenCountTotal: tokenCountTotal # this switcheroo is to allow descendant Rows to unfurl in the same Table while displaying a different value under the same Column
-                  cumulativeTokenCountPrompt: tokenCountPrompt
-                  cumulativeTokenCountCompletion: tokenCountCompletion
+                  tokenCountTotal
+                  tokenCountPrompt
+                  tokenCountCompletion
                   input {
                     value
                   }
@@ -261,7 +261,7 @@ export function TracesTable(props: TracesTableProps) {
     },
     {
       header: "total tokens",
-      accessorKey: "cumulativeTokenCountTotal",
+      accessorKey: "tokenCountTotal",
       cell: ({ row, getValue }) => {
         const value = getValue();
         if (value === null) {
@@ -270,10 +270,8 @@ export function TracesTable(props: TracesTableProps) {
         return (
           <TokenCount
             tokenCountTotal={value as number}
-            tokenCountPrompt={row.original.cumulativeTokenCountPrompt || 0}
-            tokenCountCompletion={
-              row.original.cumulativeTokenCountCompletion || 0
-            }
+            tokenCountPrompt={row.original.tokenCountPrompt || 0}
+            tokenCountCompletion={row.original.tokenCountCompletion || 0}
           />
         );
       },
@@ -327,7 +325,7 @@ export function TracesTable(props: TracesTableProps) {
     [hasNext, isLoadingNext, loadNext]
   );
   const [expanded, setExpanded] = useState<ExpandedState>({});
-  const [columnVisibility, setColumnVisibility] = useState<VisibilityState>({});
+  const columnVisibility = useTracingContext((state) => state.columnVisibility);
   const table = useReactTable<TableRow>({
     columns,
     data: tableData,
@@ -359,11 +357,7 @@ export function TracesTable(props: TracesTableProps) {
       >
         <Flex direction="row" gap="size-100" width="100%" alignItems="center">
           <SpanFilterConditionField onValidCondition={setFilterCondition} />
-          <SpanColumnSelector
-            columns={computedColumns}
-            columnVisibility={columnVisibility}
-            onColumnVisibilityChange={setColumnVisibility}
-          />
+          <SpanColumnSelector columns={computedColumns} />
         </Flex>
       </View>
       <div

diff --git a/app/src/pages/tracing/__generated__/TracesTableQuery.graphql.ts b/app/src/pages/tracing/__generated__/TracesTableQuery.graphql.ts