[OPIK-398] Project statistics - FE implementation (#735)

* [OPIK-398] Project statistics - FE implementation * - invalidate queries in case of annotate
comet-ml · Nov 27, 2024 · d117259 · d117259
1 parent ec91598
commit d117259
Show file tree

Hide file tree

Showing 16 changed files with 446 additions and 23 deletions.
diff --git a/apps/opik-frontend/src/api/traces/useSpansStatistic.ts b/apps/opik-frontend/src/api/traces/useSpansStatistic.ts
@@ -0,0 +1,46 @@
+import { QueryFunctionContext, useQuery } from "@tanstack/react-query";
+import api, { QueryConfig, SPANS_REST_ENDPOINT } from "@/api/api";
+import { SPAN_TYPE } from "@/types/traces";
+import { ColumnsStatistic } from "@/types/shared";
+import { Filters } from "@/types/filters";
+import { generateSearchByIDFilters, processFilters } from "@/lib/filters";
+
+type UseSpansStatisticParams = {
+  projectId: string;
+  traceId?: string;
+  type?: SPAN_TYPE;
+  filters?: Filters;
+  search?: string;
+};
+
+export type UseSpansStatisticResponse = {
+  stats: ColumnsStatistic;
+};
+
+const getSpansStatistic = async (
+  { signal }: QueryFunctionContext,
+  { projectId, traceId, type, filters, search }: UseSpansStatisticParams,
+) => {
+  const { data } = await api.get(`${SPANS_REST_ENDPOINT}stats`, {
+    signal,
+    params: {
+      project_id: projectId,
+      ...(traceId && { trace_id: traceId }),
+      ...(type && { type }),
+      ...processFilters(filters, generateSearchByIDFilters(search)),
+    },
+  });
+
+  return data;
+};
+
+export default function useSpansStatistic(
+  params: UseSpansStatisticParams,
+  options?: QueryConfig<UseSpansStatisticResponse>,
+) {
+  return useQuery({
+    queryKey: ["spans-statistic", params],
+    queryFn: (context) => getSpansStatistic(context, params),
+    ...options,
+  });
+}
diff --git a/apps/opik-frontend/src/api/traces/useTraceFeedbackScoreDeleteMutation.ts b/apps/opik-frontend/src/api/traces/useTraceFeedbackScoreDeleteMutation.ts
@@ -91,10 +91,12 @@ const useTraceFeedbackScoreDeleteMutation = () => {
       if (variables.spanId) {
         await queryClient.invalidateQueries({ queryKey: [SPANS_KEY] });
         await queryClient.invalidateQueries({ queryKey: ["spans-columns"] });
+        await queryClient.invalidateQueries({ queryKey: ["spans-statistic"] });
       }
 
       await queryClient.invalidateQueries({ queryKey: [TRACES_KEY] });
       await queryClient.invalidateQueries({ queryKey: ["traces-columns"] });
+      await queryClient.invalidateQueries({ queryKey: ["traces-statistic"] });
 
       await queryClient.invalidateQueries({
         queryKey: [TRACE_KEY, { traceId: variables.traceId }],

diff --git a/apps/opik-frontend/src/api/traces/useTraceFeedbackScoreSetMutation.ts b/apps/opik-frontend/src/api/traces/useTraceFeedbackScoreSetMutation.ts
@@ -106,10 +106,12 @@ const useTraceFeedbackScoreSetMutation = () => {
       if (variables.spanId) {
         await queryClient.invalidateQueries({ queryKey: [SPANS_KEY] });
         await queryClient.invalidateQueries({ queryKey: ["spans-columns"] });
+        await queryClient.invalidateQueries({ queryKey: ["spans-statistic"] });
       }
 
       await queryClient.invalidateQueries({ queryKey: [TRACES_KEY] });
       await queryClient.invalidateQueries({ queryKey: ["traces-columns"] });
+      await queryClient.invalidateQueries({ queryKey: ["traces-statistic"] });
 
       await queryClient.invalidateQueries({
         queryKey: [TRACE_KEY, { traceId: variables.traceId }],

diff --git a/apps/opik-frontend/src/api/traces/useTracesStatistic.ts b/apps/opik-frontend/src/api/traces/useTracesStatistic.ts
@@ -0,0 +1,44 @@
+import { QueryFunctionContext, useQuery } from "@tanstack/react-query";
+import api, { QueryConfig, TRACES_REST_ENDPOINT } from "@/api/api";
+import { ColumnsStatistic } from "@/types/shared";
+import { Filters } from "@/types/filters";
+import { generateSearchByIDFilters, processFilters } from "@/lib/filters";
+
+type UseTracesStatisticParams = {
+  projectId: string;
+  filters?: Filters;
+  search?: string;
+};
+
+export type UseTracesStatisticResponse = {
+  stats: ColumnsStatistic;
+};
+
+const getTracesStatistic = async (
+  { signal }: QueryFunctionContext,
+  { projectId, filters, search }: UseTracesStatisticParams,
+) => {
+  const { data } = await api.get<UseTracesStatisticResponse>(
+    `${TRACES_REST_ENDPOINT}stats`,
+    {
+      signal,
+      params: {
+        project_id: projectId,
+        ...processFilters(filters, generateSearchByIDFilters(search)),
+      },
+    },
+  );
+
+  return data;
+};
+
+export default function useTracesStatistic(
+  params: UseTracesStatisticParams,
+  options?: QueryConfig<UseTracesStatisticResponse>,
+) {
+  return useQuery({
+    queryKey: ["traces-statistic", params],
+    queryFn: (context) => getTracesStatistic(context, params),
+    ...options,
+  });
+}
diff --git a/apps/opik-frontend/src/components/pages/TracesPage/TracesSpansTab.tsx b/apps/opik-frontend/src/components/pages/TracesPage/TracesSpansTab.tsx
@@ -22,6 +22,7 @@ import {
   COLUMN_SELECT_ID,
   COLUMN_TYPE,
   ColumnData,
+  ColumnsStatistic,
   DynamicColumn,
   ROW_HEIGHT,
 } from "@/types/shared";
@@ -48,9 +49,12 @@ import FeedbackScoreHeader from "@/components/shared/DataTableHeaders/FeedbackSc
 import TraceDetailsPanel from "@/components/shared/TraceDetailsPanel/TraceDetailsPanel";
 import TooltipWrapper from "@/components/shared/TooltipWrapper/TooltipWrapper";
 import { formatDate } from "@/lib/date";
+import useTracesOrSpansStatistic from "@/hooks/useTracesOrSpansStatistic";
 
 const getRowId = (d: Trace | Span) => d.id;
 
+const REFETCH_INTERVAL = 30000;
+
 export const TRACES_PAGE_COLUMNS: ColumnData<BaseTraceData>[] = [
   {
     id: "name",
@@ -131,6 +135,12 @@ export const TRACES_PAGE_FILTERS_COLUMNS = [
     label: "ID",
     type: COLUMN_TYPE.string,
   },
+
+  {
+    id: "feedback_scores",
+    label: "Feedback scores",
+    type: COLUMN_TYPE.numberDictionary,
+  },
   ...TRACES_PAGE_COLUMNS,
 ];
 
@@ -207,18 +217,31 @@ export const TracesSpansTab: React.FC<TracesSpansTabProps> = ({
       truncate: true,
     },
     {
-      refetchInterval: 30000,
+      refetchInterval: REFETCH_INTERVAL,
     },
   );
 
+  const { data: statisticData, refetch: refetchStatistic } =
+    useTracesOrSpansStatistic(
+      {
+        projectId,
+        type: type as TRACE_DATA_TYPE,
+        filters,
+        search: search as string,
+      },
+      {
+        refetchInterval: REFETCH_INTERVAL,
+      },
+    );
+
   const { data: feedbackScoresData, isPending: isFeedbackScoresPending } =
     useTracesOrSpansScoresColumns(
       {
         projectId,
         type: type as TRACE_DATA_TYPE,
       },
       {
-        refetchInterval: 30000,
+        refetchInterval: REFETCH_INTERVAL,
       },
     );
 
@@ -230,6 +253,10 @@ export const TracesSpansTab: React.FC<TracesSpansTabProps> = ({
     : "No search results";
 
   const rows: Array<Span | Trace> = useMemo(() => data?.content ?? [], [data]);
+  const columnsStatistic: ColumnsStatistic = useMemo(
+    () => statisticData?.stats ?? [],
+    [statisticData],
+  );
 
   const dynamicColumns = useMemo(() => {
     return (feedbackScoresData?.scores ?? []).map<DynamicColumn>((c) => ({
@@ -297,6 +324,7 @@ export const TracesSpansTab: React.FC<TracesSpansTabProps> = ({
             cell: FeedbackScoreCell as never,
             accessorFn: (row) =>
               row.feedback_scores?.find((f) => f.name === label),
+            statisticKey: `feedback_score.${label}`,
           }) as ColumnData<BaseTraceData>,
       ),
     ];
@@ -426,7 +454,10 @@ export const TracesSpansTab: React.FC<TracesSpansTabProps> = ({
               variant="outline"
               size="icon"
               className="shrink-0"
-              onClick={() => refetch()}
+              onClick={() => {
+                refetch();
+                refetchStatistic();
+              }}
             >
               <RotateCw className="size-4" />
             </Button>
@@ -452,6 +483,7 @@ export const TracesSpansTab: React.FC<TracesSpansTabProps> = ({
       </div>
       <DataTable
         columns={columns}
+        columnsStatistic={columnsStatistic}
         data={rows}
         activeRowId={activeRowId ?? ""}
         resizeConfig={resizeConfig}

diff --git a/apps/opik-frontend/src/components/shared/DataTable/DataTable.tsx b/apps/opik-frontend/src/components/shared/DataTable/DataTable.tsx
@@ -31,6 +31,7 @@ import DataTableColumnResizer from "@/components/shared/DataTable/DataTableColum
 import {
   CELL_VERTICAL_ALIGNMENT,
   COLUMN_TYPE,
+  ColumnsStatistic,
   OnChangeFn,
   ROW_HEIGHT,
 } from "@/types/shared";
@@ -44,6 +45,7 @@ import { TABLE_HEADER_Z_INDEX } from "@/constants/shared";
 declare module "@tanstack/react-table" {
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
   interface TableMeta<TData extends RowData> {
+    columnsStatistic?: ColumnsStatistic;
     rowHeight: ROW_HEIGHT;
     rowHeightClass: string;
   }
@@ -55,6 +57,7 @@ declare module "@tanstack/react-table" {
     iconType?: COLUMN_TYPE;
     verticalAlignment?: CELL_VERTICAL_ALIGNMENT;
     overrideRowHeight?: ROW_HEIGHT;
+    statisticKey?: string;
     custom?: object;
   }
 }
@@ -90,6 +93,7 @@ interface ExpandingConfig {
 
 interface DataTableProps<TData, TValue> {
   columns: ColumnDef<TData, TValue>[];
+  columnsStatistic?: ColumnsStatistic;
   data: TData[];
   renderCustomRow?: (row: Row<TData>) => ReactNode | null;
   getIsCustomRow?: (row: Row<TData>) => boolean;
@@ -109,6 +113,7 @@ interface DataTableProps<TData, TValue> {
 
 const DataTable = <TData, TValue>({
   columns,
+  columnsStatistic,
   data,
   renderCustomRow,
   getIsCustomRow = () => false,
@@ -165,6 +170,7 @@ const DataTable = <TData, TValue>({
       ...(columnPinning && { columnPinning }),
     },
     meta: {
+      columnsStatistic,
       rowHeight,
       rowHeightClass: getRowHeightClass(rowHeight),
     },

diff --git a/apps/opik-frontend/src/components/shared/DataTable/DataTableColumnResizer.tsx b/apps/opik-frontend/src/components/shared/DataTable/DataTableColumnResizer.tsx
@@ -10,6 +10,9 @@ const DataTableColumnResizer = <TData,>({
   header,
 }: DataTableColumnResizerProps<TData>) => {
   if (!header.column.getCanResize()) return <></>;
+  const hasStatistic = Boolean(
+    header.getContext().table.options.meta?.columnsStatistic,
+  );
 
   return (
     <div
@@ -26,7 +29,12 @@ const DataTableColumnResizer = <TData,>({
         header.column.getIsLastColumn() ? "right-0" : "-right-1",
       )}
     >
-      <div className="absolute top-3 h-5 w-px bg-border"></div>
+      <div
+        className={cn(
+          "absolute top-2 h-7 w-px bg-border",
+          hasStatistic && "h-10",
+        )}
+      ></div>
       <div className="absolute inset-y-0 w-px bg-transparent transition-colors group-hover:bg-gray-600 group-active:bg-blue-600"></div>
     </div>
   );

diff --git a/apps/opik-frontend/src/components/shared/DataTableCells/ListCell.tsx b/apps/opik-frontend/src/components/shared/DataTableCells/ListCell.tsx
@@ -12,19 +12,20 @@ const ListCell = (context: CellContext<unknown, unknown>) => {
     return null;
   }
 
-  const rowHeight = context.table.options.meta?.rowHeight ?? ROW_HEIGHT.small;
+  const isSmall =
+    (context.table.options.meta?.rowHeight ?? ROW_HEIGHT.small) ===
+    ROW_HEIGHT.small;
 
   return (
     <CellWrapper
       metadata={context.column.columnDef.meta}
       tableMetadata={context.table.options.meta}
+      className={cn(isSmall && "py-0")}
     >
       <div
         className={cn(
           "flex max-h-full flex-row gap-2",
-          rowHeight === ROW_HEIGHT.small
-            ? "overflow-x-auto"
-            : "flex-wrap overflow-auto",
+          isSmall ? "overflow-x-auto" : "flex-wrap overflow-auto",
         )}
       >
         {items.sort().map((item) => {

diff --git a/apps/opik-frontend/src/components/shared/DataTableHeaders/FeedbackScoreHeader.tsx b/apps/opik-frontend/src/components/shared/DataTableHeaders/FeedbackScoreHeader.tsx
@@ -2,21 +2,26 @@ import React from "react";
 import { HeaderContext } from "@tanstack/react-table";
 import { TAG_VARIANTS_COLOR_MAP } from "@/components/ui/tag";
 import { generateTagVariant } from "@/lib/traces";
+import HeaderWrapper from "@/components/shared/DataTableHeaders/HeaderWrapper";
 
-const FeedbackScoreHeader = <TData,>({
-  column,
-}: HeaderContext<TData, unknown>) => {
+const FeedbackScoreHeader = <TData,>(
+  context: HeaderContext<TData, unknown>,
+) => {
+  const { column } = context;
   const { header } = column.columnDef.meta ?? {};
   const color = TAG_VARIANTS_COLOR_MAP[generateTagVariant(header!)!];
 
   return (
-    <div className="flex size-full items-center justify-end gap-1.5 px-3">
+    <HeaderWrapper
+      metadata={context.column.columnDef.meta}
+      tableMetadata={context.table.options.meta}
+    >
       <div
-        className="size-2 shrink-0 rounded-[2px] bg-[--color-bg]"
+        className="mr-0.5 size-2 shrink-0 rounded-[2px] bg-[--color-bg]"
         style={{ "--color-bg": color } as React.CSSProperties}
       ></div>
       <span className="truncate">{header}</span>
-    </div>
+    </HeaderWrapper>
   );
 };