fix: select columns for child call table preview

wandb · Oct 2, 2024 · c47e99e · c47e99e
1 parent 4e22e1e
commit c47e99e
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 4 deletions.
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallDetails.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallPage/CallDetails.tsx
@@ -71,6 +71,8 @@ export const CallSchemaLink = ({call}: {call: CallSchema}) => {
   );
 };
 
+const SHOWN_COLS = ['op_name', 'status', 'inputs.*', 'output', 'output.*'];
+
 export const CallDetails: FC<{
   call: CallSchema;
 }> = ({call}) => {
@@ -176,6 +178,7 @@ export const CallDetails: FC<{
               }}
               entity={call.entity}
               project={call.project}
+              shownCols={SHOWN_COLS}
             />
           );
           if (isPeeking) {

diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/CallsTable.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/CallsTable.tsx
@@ -140,6 +140,9 @@ export const CallsTable: FC<{
 
   paginationModel?: GridPaginationModel;
   setPaginationModel?: (newModel: GridPaginationModel) => void;
+
+  // Can include glob for prefix match, e.g. "inputs.*"
+  shownCols?: string[];
 }> = ({
   entity,
   project,
@@ -158,6 +161,7 @@ export const CallsTable: FC<{
   setSortModel,
   paginationModel,
   setPaginationModel,
+  shownCols,
 }) => {
   const {loading: loadingUserInfo, userInfo} = useViewerInfo();
 
@@ -303,6 +307,7 @@ export const CallsTable: FC<{
     onCollapse,
     onExpand,
     columnIsRefExpanded,
+    shownCols,
     onAddFilter
   );
 

diff --git a/...-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/callsTableColumns.tsx b/...-js/src/components/PagePanelComponents/Home/Browse3/pages/CallsPage/callsTableColumns.tsx
@@ -54,6 +54,7 @@ export const useCallsTableColumns = (
   onCollapse: (col: string) => void,
   onExpand: (col: string) => void,
   columnIsRefExpanded: (col: string) => boolean,
+  shownCols?: string[],
   onAddFilter?: (field: string, operator: string | null, value: any) => void
 ) => {
   const [userDefinedColumnWidths, setUserDefinedColumnWidths] = useState<
@@ -129,6 +130,7 @@ export const useCallsTableColumns = (
         onExpand,
         columnIsRefExpanded,
         userDefinedColumnWidths,
+        shownCols,
         onAddFilter
       ),
     [
@@ -144,6 +146,7 @@ export const useCallsTableColumns = (
       onExpand,
       columnIsRefExpanded,
       userDefinedColumnWidths,
+      shownCols,
       onAddFilter,
     ]
   );
@@ -168,15 +171,29 @@ function buildCallsTableColumns(
   onExpand: (col: string) => void,
   columnIsRefExpanded: (col: string) => boolean,
   userDefinedColumnWidths: Record<string, number>,
+  shownCols?: string[],
   onAddFilter?: (field: string, operator: string | null, value: any) => void
 ): {
   cols: Array<GridColDef<TraceCallSchema>>;
   colGroupingModel: GridColumnGroupingModel;
 } {
   // Filters summary.usage. because we add a derived column for tokens and cost
-  const filteredDynamicColumnNames = allDynamicColumnNames.filter(
-    c => !HIDDEN_DYNAMIC_COLUMN_PREFIXES.some(p => c.startsWith(p + '.'))
-  );
+  // Sort attributes after inputs and outputs.
+  const filteredDynamicColumnNames = allDynamicColumnNames
+    .filter(
+      c => !HIDDEN_DYNAMIC_COLUMN_PREFIXES.some(p => c.startsWith(p + '.'))
+    )
+    .sort((a, b) => {
+      const prefixes = ['inputs.', 'output.', 'attributes.'];
+      const aPrefix =
+        a === 'output' ? 'output.' : prefixes.find(p => a.startsWith(p)) ?? '';
+      const bPrefix =
+        b === 'output' ? 'output.' : prefixes.find(p => b.startsWith(p)) ?? '';
+      if (aPrefix !== bPrefix) {
+        return prefixes.indexOf(aPrefix) - prefixes.indexOf(bPrefix);
+      }
+      return a.localeCompare(b);
+    });
 
   const cols: Array<GridColDef<TraceCallSchema>> = [
     {
@@ -435,7 +452,21 @@ function buildCallsTableColumns(
     }
   });
 
-  return {cols, colGroupingModel: groupingModel};
+  // TODO: It would be better to build up the cols rather than throwing away
+  //       some at the end, but making simpler change for now.
+  let orderedCols = cols;
+  if (shownCols !== undefined) {
+    orderedCols = shownCols.flatMap(shownCol => {
+      if (shownCol.includes('*')) {
+        const regex = new RegExp('^' + shownCol.replace('*', '.*') + '$');
+        return cols.filter(col => regex.test(col.field));
+      } else {
+        return cols.filter(col => col.field === shownCol);
+      }
+    });
+  }
+
+  return {cols: orderedCols, colGroupingModel: groupingModel};
 }
 /**
  * This function maintains an ever-growing list of dynamic column names. It is used to