Implement global pivot configuration management in GlobalState and update PivotTab to utilize it, including loading, saving, and resetting functionality.

Author: Dylan Huang

vite-app/src/GlobalState.tsx

@@ -2,17 +2,75 @@ import { makeAutoObservable } from "mobx";
 import type { EvaluationRow } from "./types/eval-protocol";
 import flattenJson from "./util/flatten-json";
 
+// Pivot configuration interface
+export interface PivotConfig {
+  selectedRowFields: string[];
+  selectedColumnFields: string[];
+  selectedValueField: string;
+  selectedAggregator: string;
+  filters: Array<{ field: string; operator: string; value: string }>;
+}
+
+// Default pivot configuration
+const DEFAULT_PIVOT_CONFIG: PivotConfig = {
+  selectedRowFields: ["$.eval_metadata.name"],
+  selectedColumnFields: ["$.input_metadata.completion_params.model"],
+  selectedValueField: "$.evaluation_result.score",
+  selectedAggregator: "avg",
+  filters: [],
+};
+
 export class GlobalState {
   isConnected: boolean = false;
   // rollout_id -> EvaluationRow
   dataset: Record<string, EvaluationRow> = {};
   // rollout_id -> expanded
   expandedRows: Record<string, boolean> = {};
+  // Pivot configuration
+  pivotConfig: PivotConfig;
 
   constructor() {
+    // Load pivot config from localStorage or use defaults
+    this.pivotConfig = this.loadPivotConfig();
     makeAutoObservable(this);
   }
 
+  // Load pivot configuration from localStorage
+  private loadPivotConfig(): PivotConfig {
+    try {
+      const stored = localStorage.getItem("pivotConfig");
+      if (stored) {
+        const parsed = JSON.parse(stored);
+        // Merge with defaults to handle any missing properties
+        return { ...DEFAULT_PIVOT_CONFIG, ...parsed };
+      }
+    } catch (error) {
+      console.warn("Failed to load pivot config from localStorage:", error);
+    }
+    return { ...DEFAULT_PIVOT_CONFIG };
+  }
+
+  // Save pivot configuration to localStorage
+  private savePivotConfig() {
+    try {
+      localStorage.setItem("pivotConfig", JSON.stringify(this.pivotConfig));
+    } catch (error) {
+      console.warn("Failed to save pivot config to localStorage:", error);
+    }
+  }
+
+  // Update pivot configuration and save to localStorage
+  updatePivotConfig(updates: Partial<PivotConfig>) {
+    Object.assign(this.pivotConfig, updates);
+    this.savePivotConfig();
+  }
+
+  // Reset pivot configuration to defaults
+  resetPivotConfig() {
+    this.pivotConfig = { ...DEFAULT_PIVOT_CONFIG };
+    this.savePivotConfig();
+  }
+
   upsertRows(dataset: EvaluationRow[]) {
     dataset.forEach((row) => {
       if (!row.execution_metadata?.rollout_id) {

vite-app/src/components/PivotTab.tsx

@@ -1,8 +1,9 @@
 import { observer } from "mobx-react";
 import PivotTable from "./PivotTable";
 import Select from "./Select";
+import Button from "./Button";
 import { state } from "../App";
-import { useState } from "react";
+import { useEffect } from "react";
 
 interface FieldSelectorProps {
   title: string;
@@ -243,45 +244,61 @@ const FilterSelector = ({
 };
 
 const PivotTab = observer(() => {
-  const [selectedRowFields, setSelectedRowFields] = useState<string[]>([
-    "$.eval_metadata.name",
-  ]);
-  const [selectedColumnFields, setSelectedColumnFields] = useState<string[]>([
-    "$.input_metadata.completion_params.model",
-  ]);
-  const [selectedValueField, setSelectedValueField] = useState<string>(
-    "$.evaluation_result.score"
-  );
-  const [selectedAggregator, setSelectedAggregator] = useState<string>("avg");
-  const [filters, setFilters] = useState<
-    Array<{ field: string; operator: string; value: string }>
-  >([]);
+  // Use global state instead of local state
+  const { pivotConfig } = state;
+
+  // Update global state when configuration changes
+  const updateRowFields = (index: number, value: string) => {
+    const newRowFields = [...pivotConfig.selectedRowFields];
+    newRowFields[index] = value;
+    state.updatePivotConfig({ selectedRowFields: newRowFields });
+  };
+
+  const updateColumnFields = (index: number, value: string) => {
+    const newColumnFields = [...pivotConfig.selectedColumnFields];
+    newColumnFields[index] = value;
+    state.updatePivotConfig({ selectedColumnFields: newColumnFields });
+  };
+
+  const updateValueField = (value: string) => {
+    state.updatePivotConfig({ selectedValueField: value });
+  };
+
+  const updateAggregator = (value: string) => {
+    state.updatePivotConfig({ selectedAggregator: value });
+  };
+
+  const updateFilters = (
+    filters: Array<{ field: string; operator: string; value: string }>
+  ) => {
+    state.updatePivotConfig({ filters });
+  };
 
   const createFieldHandler = (
-    setter: React.Dispatch<React.SetStateAction<string[]>>
+    updater: (index: number, value: string) => void
   ) => {
     return (index: number, value: string) => {
-      setter((prev) => {
-        const newFields = [...prev];
-        newFields[index] = value;
-        return newFields;
-      });
+      updater(index, value);
     };
   };
 
   const createAddHandler = (
-    setter: React.Dispatch<React.SetStateAction<string[]>>
+    fields: string[],
+    updater: (fields: string[]) => void
   ) => {
     return () => {
-      setter((prev) => (prev.length < 3 ? [...prev, ""] : prev));
+      if (fields.length < 3) {
+        updater([...fields, ""]);
+      }
     };
   };
 
   const createRemoveHandler = (
-    setter: React.Dispatch<React.SetStateAction<string[]>>
+    fields: string[],
+    updater: (fields: string[]) => void
   ) => {
     return (index: number) => {
-      setter((prev) => prev.filter((_, i) => i !== index));
+      updater(fields.filter((_, i) => i !== index));
     };
   };
 
@@ -338,63 +355,85 @@ const PivotTab = observer(() => {
         specific subsets of your data.
       </div>
 
+      {/* Controls Section with Reset Button */}
+      <div className="mb-4 flex justify-between items-center">
+        <Button
+          onClick={() => state.resetPivotConfig()}
+          variant="secondary"
+          size="sm"
+        >
+          Reset to Defaults
+        </Button>
+      </div>
+
       <FieldSelector
         title="Row Fields"
-        fields={selectedRowFields}
-        onFieldChange={createFieldHandler(setSelectedRowFields)}
-        onAddField={createAddHandler(setSelectedRowFields)}
-        onRemoveField={createRemoveHandler(setSelectedRowFields)}
+        fields={pivotConfig.selectedRowFields}
+        onFieldChange={createFieldHandler(updateRowFields)}
+        onAddField={createAddHandler(pivotConfig.selectedRowFields, (fields) =>
+          state.updatePivotConfig({ selectedRowFields: fields })
+        )}
+        onRemoveField={createRemoveHandler(
+          pivotConfig.selectedRowFields,
+          (fields) => state.updatePivotConfig({ selectedRowFields: fields })
+        )}
         availableKeys={availableKeys}
         variant="row"
       />
 
       <FieldSelector
         title="Column Fields"
-        fields={selectedColumnFields}
-        onFieldChange={createFieldHandler(setSelectedColumnFields)}
-        onAddField={createAddHandler(setSelectedColumnFields)}
-        onRemoveField={createRemoveHandler(setSelectedColumnFields)}
+        fields={pivotConfig.selectedColumnFields}
+        onFieldChange={createFieldHandler(updateColumnFields)}
+        onAddField={createAddHandler(
+          pivotConfig.selectedColumnFields,
+          (fields) => state.updatePivotConfig({ selectedColumnFields: fields })
+        )}
+        onRemoveField={createRemoveHandler(
+          pivotConfig.selectedColumnFields,
+          (fields) => state.updatePivotConfig({ selectedColumnFields: fields })
+        )}
         availableKeys={availableKeys}
         variant="column"
       />
 
       <SingleFieldSelector
         title="Value Field"
-        field={selectedValueField}
-        onFieldChange={setSelectedValueField}
+        field={pivotConfig.selectedValueField}
+        onFieldChange={updateValueField}
         availableKeys={availableKeys}
       />
 
       <AggregatorSelector
-        aggregator={selectedAggregator}
-        onAggregatorChange={setSelectedAggregator}
+        aggregator={pivotConfig.selectedAggregator}
+        onAggregatorChange={updateAggregator}
       />
 
       <FilterSelector
-        filters={filters}
-        onFiltersChange={setFilters}
+        filters={pivotConfig.filters}
+        onFiltersChange={updateFilters}
         availableKeys={availableKeys}
       />
 
       <PivotTable
         data={state.flattenedDataset}
         rowFields={
-          selectedRowFields.filter(
+          pivotConfig.selectedRowFields.filter(
             (field) => field !== ""
           ) as (keyof (typeof state.flattenedDataset)[number])[]
         }
         columnFields={
-          selectedColumnFields.filter(
+          pivotConfig.selectedColumnFields.filter(
             (field) => field !== ""
           ) as (keyof (typeof state.flattenedDataset)[number])[]
         }
         valueField={
-          selectedValueField as keyof (typeof state.flattenedDataset)[number]
+          pivotConfig.selectedValueField as keyof (typeof state.flattenedDataset)[number]
         }
-        aggregator={selectedAggregator as any}
+        aggregator={pivotConfig.selectedAggregator as any}
         showRowTotals
         showColumnTotals
-        filter={createFilterFunction(filters)}
+        filter={createFilterFunction(pivotConfig.filters)}
       />
     </div>
   );