feat(trace): add live updating UI, stream on/off switch

app/package.json

@@ -6,7 +6,7 @@
   "license": "None",
   "private": true,
   "dependencies": {
-    "@arizeai/components": "^0.17.0",
+    "@arizeai/components": "^0.17.3",
     "@arizeai/point-cloud": "^3.0.4",
     "@codemirror/lang-json": "^6.0.1",
     "@codemirror/view": "^6.16.0",
@@ -74,7 +74,7 @@
     "build:relay": "relay-compiler",
     "watch": "./esbuild.config.mjs dev",
     "test": "jest --config ./jest.config.js",
-    "dev": "npm run dev:server:image & npm run build:static && npm run watch",
+    "dev": "npm run dev:server:traces:llama_index_rag & npm run build:static && npm run watch",
     "dev:server:mnist": "python3 -m phoenix.server.main fixture fashion_mnist",
     "dev:server:mnist:single": "python3 -m phoenix.server.main fixture fashion_mnist --primary-only true",
     "dev:server:sentiment": "python3 -m phoenix.server.main fixture sentiment_classification_language_drift",

app/src/components/table/TableEmpty.tsx

@@ -12,7 +12,7 @@ export function TableEmpty() {
             padding: ${theme.spacing.margin24}px ${theme.spacing.margin24}px !important;
           `}
         >
-          No data
+          No Data
         </td>
       </tr>
     </tbody>

app/src/hooks/useInterval.ts

@@ -0,0 +1,22 @@
+import { useEffect, useRef } from "react";
+
+type Callback = () => void;
+export function useInterval(callback: Callback, delay: number) {
+  const savedCallback = useRef<Callback | null>(null);
+
+  // Remember the latest callback.
+  useEffect(() => {
+    savedCallback.current = callback;
+  }, [callback]);
+
+  // Set up the interval.
+  useEffect(() => {
+    function tick() {
+      savedCallback.current && savedCallback.current();
+    }
+    if (delay !== null) {
+      const id = setInterval(tick, delay);
+      return () => clearInterval(id);
+    }
+  }, [delay]);
+}

app/src/pages/tracing/SpansTable.tsx

@@ -15,6 +15,7 @@ import { Icon, Icons } from "@arizeai/components";
 
 import { Link } from "@phoenix/components/Link";
 import { selectableTableCSS } from "@phoenix/components/table/styles";
+import { TableEmpty } from "@phoenix/components/table/TableEmpty";
 import { TextCell } from "@phoenix/components/table/TextCell";
 import { TimestampCell } from "@phoenix/components/table/TimestampCell";
 import { LatencyText } from "@phoenix/components/trace/LatencyText";
@@ -214,6 +215,7 @@ export function SpansTable(props: SpansTableProps) {
     getSortedRowModel: getSortedRowModel(),
   });
   const rows = table.getRowModel().rows;
+  const isEmpty = rows.length === 0;
   return (
     <div
       css={css`
@@ -261,31 +263,35 @@ export function SpansTable(props: SpansTableProps) {
             </tr>
           ))}
         </thead>
-        <tbody>
-          {rows.map((row) => {
-            return (
-              <tr
-                key={row.id}
-                onClick={() =>
-                  navigate(
-                    `traces/${row.original.context.traceId}?selectedSpanId=${row.original.context.spanId}`
-                  )
-                }
-              >
-                {row.getVisibleCells().map((cell) => {
-                  return (
-                    <td key={cell.id}>
-                      {flexRender(
-                        cell.column.columnDef.cell,
-                        cell.getContext()
-                      )}
-                    </td>
-                  );
-                })}
-              </tr>
-            );
-          })}
-        </tbody>
+        {isEmpty ? (
+          <TableEmpty />
+        ) : (
+          <tbody>
+            {rows.map((row) => {
+              return (
+                <tr
+                  key={row.id}
+                  onClick={() =>
+                    navigate(
+                      `traces/${row.original.context.traceId}?selectedSpanId=${row.original.context.spanId}`
+                    )
+                  }
+                >
+                  {row.getVisibleCells().map((cell) => {
+                    return (
+                      <td key={cell.id}>
+                        {flexRender(
+                          cell.column.columnDef.cell,
+                          cell.getContext()
+                        )}
+                      </td>
+                    );
+                  })}
+                </tr>
+              );
+            })}
+          </tbody>
+        )}
       </table>
     </div>
   );

app/src/pages/tracing/StreamToggle.tsx

@@ -0,0 +1,79 @@
+import React, {
+  startTransition,
+  useCallback,
+  useEffect,
+  useRef,
+  useState,
+} from "react";
+import { graphql, useRefetchableFragment } from "react-relay";
+
+import { Switch } from "@arizeai/components";
+
+import { useInterval } from "@phoenix/hooks/useInterval";
+
+import { StreamToggle_data$key } from "./__generated__/StreamToggle_data.graphql";
+
+/**
+ * Check every few seconds for new data
+ */
+const REFRESH_INTERVAL_MS = 2000;
+
+export function StreamToggle(props: {
+  query: StreamToggle_data$key;
+  onRefresh: () => void;
+}) {
+  const { onRefresh } = props;
+  const [isStreaming, setIsStreaming] = useState<boolean>(true);
+
+  const [traceCountData, refetchCounts] = useRefetchableFragment(
+    graphql`
+      fragment StreamToggle_data on Query
+      @refetchable(queryName: "StreamToggleRefetchQuery") {
+        traceCount: spans(rootSpansOnly: true) {
+          pageInfo {
+            totalCount
+          }
+        }
+      }
+    `,
+    props.query
+  );
+  // Keep track of the loaded trace count so we can detect when it changes
+  const loadedTraceCountRef = useRef<number>(
+    traceCountData.traceCount.pageInfo.totalCount
+  );
+
+  // Refetch the count of traces if the streaming toggle is on
+  const refetchCountsIfStreaming = useCallback(() => {
+    if (isStreaming) {
+      startTransition(() => {
+        refetchCounts({}, { fetchPolicy: "store-and-network" });
+      });
+    }
+  }, [isStreaming, refetchCounts]);
+
+  // We want to refetch higher up the render tree when the counts change
+  const totalTraceCount = traceCountData.traceCount.pageInfo.totalCount;
+  useEffect(() => {
+    if (loadedTraceCountRef.current !== totalTraceCount) {
+      // Update the loaded trace count so the effect doesn't fire again
+      loadedTraceCountRef.current = totalTraceCount;
+      onRefresh();
+    }
+  }, [onRefresh, totalTraceCount]);
+
+  useInterval(refetchCountsIfStreaming, REFRESH_INTERVAL_MS);
+  return (
+    <Switch
+      labelPlacement="start"
+      defaultSelected
+      onChange={() => {
+        setIsStreaming(!isStreaming);
+        // Perform one last refresh before pausing / resuming
+        onRefresh();
+      }}
+    >
+      Stream
+    </Switch>
+  );
+}

app/src/pages/tracing/TracesTable.tsx

@@ -19,6 +19,7 @@ import { Flex, Icon, Icons } from "@arizeai/components";
 import { Link } from "@phoenix/components/Link";
 import { TextCell } from "@phoenix/components/table";
 import { selectableTableCSS } from "@phoenix/components/table/styles";
+import { TableEmpty } from "@phoenix/components/table/TableEmpty";
 import { TableExpandButton } from "@phoenix/components/table/TableExpandButton";
 import { TimestampCell } from "@phoenix/components/table/TimestampCell";
 import { LatencyText } from "@phoenix/components/trace/LatencyText";
@@ -317,6 +318,7 @@ export function TracesTable(props: TracesTableProps) {
     getExpandedRowModel: getExpandedRowModel(),
   });
   const rows = table.getRowModel().rows;
+  const isEmpty = rows.length === 0;
   return (
     <div
       css={css`
@@ -364,29 +366,33 @@ export function TracesTable(props: TracesTableProps) {
             </tr>
           ))}
         </thead>
-        <tbody>
-          {rows.map((row) => {
-            return (
-              <tr
-                key={row.id}
-                onClick={() =>
-                  navigate(`traces/${row.original.context.traceId}`)
-                }
-              >
-                {row.getVisibleCells().map((cell) => {
-                  return (
-                    <td key={cell.id}>
-                      {flexRender(
-                        cell.column.columnDef.cell,
-                        cell.getContext()
-                      )}
-                    </td>
-                  );
-                })}
-              </tr>
-            );
-          })}
-        </tbody>
+        {isEmpty ? (
+          <TableEmpty />
+        ) : (
+          <tbody>
+            {rows.map((row) => {
+              return (
+                <tr
+                  key={row.id}
+                  onClick={() =>
+                    navigate(`traces/${row.original.context.traceId}`)
+                  }
+                >
+                  {row.getVisibleCells().map((cell) => {
+                    return (
+                      <td key={cell.id}>
+                        {flexRender(
+                          cell.column.columnDef.cell,
+                          cell.getContext()
+                        )}
+                      </td>
+                    );
+                  })}
+                </tr>
+              );
+            })}
+          </tbody>
+        )}
       </table>
     </div>
   );

app/src/pages/tracing/TracingHomePage.tsx

@@ -1,4 +1,4 @@
-import React, { Suspense } from "react";
+import React, { startTransition, Suspense, useCallback, useState } from "react";
 import { graphql, useLazyLoadQuery } from "react-relay";
 import { Outlet } from "react-router";
 import { css } from "@emotion/react";
@@ -7,19 +7,39 @@ import { TabPane, Tabs } from "@arizeai/components";
 
 import { TracingHomePageQuery } from "./__generated__/TracingHomePageQuery.graphql";
 import { SpansTable } from "./SpansTable";
+import { StreamToggle } from "./StreamToggle";
 import { TracesTable } from "./TracesTable";
 import { TracingHomePageHeader } from "./TracingHomePageHeader";
 
+/**
+ * Sets the refetch key to trigger a refetch
+ */
+const useRefetch = (): [number, () => void] => {
+  const [fetchKey, setFetchKey] = useState<number>(0);
+  const refetch = useCallback(() => {
+    startTransition(() => {
+      setFetchKey((prev) => prev + 1);
+    });
+  }, [setFetchKey]);
+  return [fetchKey, refetch];
+};
+
 export function TracingHomePage() {
+  const [fetchKey, refetch] = useRefetch();
   const data = useLazyLoadQuery<TracingHomePageQuery>(
     graphql`
       query TracingHomePageQuery {
         ...SpansTable_spans
         ...TracesTable_spans
         ...TracingHomePageHeader_stats
+        ...StreamToggle_data
       }
     `,
-    {}
+    {},
+    {
+      fetchKey,
+      fetchPolicy: "store-and-network",
+    }
   );
   return (
     <main
@@ -48,7 +68,10 @@ export function TracingHomePage() {
         }
       `}
     >
-      <TracingHomePageHeader query={data} />
+      <TracingHomePageHeader
+        query={data}
+        extra={<StreamToggle query={data} onRefresh={() => refetch()} />}
+      />
       <Tabs>
         <TabPane name="Traces">
           {({ isSelected }) => {