Don’t serialize large samples in vscode

UKGovernmentBEIS · Jan 10, 2025 · e62cd77 · e62cd77
1 parent d84a9af
commit e62cd77
Show file tree

Hide file tree

Showing 2 changed files with 132 additions and 5 deletions.
diff --git a/src/inspect_ai/_view/www/dist/assets/index.js b/src/inspect_ai/_view/www/dist/assets/index.js
@@ -15923,7 +15923,9 @@ var require_assets = __commonJS({
         const scrollProps = scrollRef ? {} : { onscroll: handleScroll };
         return m$1`
       <div ref=${baseRef} ...${props} ...${scrollProps}>
-        <div style=${{ ...style_inner, height: `${listMetrics.totalHeight}px` }}>
+        <div
+          style=${{ ...style_inner, height: `${listMetrics.totalHeight}px` }}
+        >
           <div
             style=${{ ...style_content, top: `${top2}px` }}
             ref=${containerRef}
@@ -32637,18 +32639,72 @@ self.onmessage = function (e) {
     const vscode = getVscodeApi();
     let initialState = void 0;
     if (vscode) {
-      initialState = vscode.getState();
+      initialState = filterState(vscode.getState());
     }
     D$1(m$1`<${App}
     api=${api}
     initialState=${initialState}
     saveInitialState=${throttle$1((state) => {
       const vscode2 = getVscodeApi();
       if (vscode2) {
-        vscode2.setState(state);
+        vscode2.setState(filterState(state));
       }
     }, 1e3)}
   />`, document.getElementById("app"));
+    function filterState(state) {
+      if (!state) {
+        return state;
+      }
+      const filters = [filterLargeSample, filterLargeSelectedLog];
+      return filters.reduce((filteredState, filter) => filter(filteredState), state);
+    }
+    function filterLargeSample(state) {
+      if (!state || !state.selectedSample) {
+        return state;
+      }
+      const estimatedTotalSize = estimateSize(state.selectedSample.messages);
+      if (estimatedTotalSize > 4e5) {
+        const {
+          selectedSample,
+          ...filteredState
+        } = state;
+        return filteredState;
+      } else {
+        return state;
+      }
+    }
+    function filterLargeSelectedLog(state) {
+      var _a2;
+      if (!state || !((_a2 = state.selectedLog) == null ? void 0 : _a2.contents)) {
+        return state;
+      }
+      const estimatedSize = estimateSize(state.selectedLog.contents.sampleSummaries);
+      if (estimatedSize > 4e5) {
+        const {
+          selectedLog,
+          ...filteredState
+        } = state;
+        return filteredState;
+      } else {
+        return state;
+      }
+    }
+    function estimateSize(list2, frequency = 0.2) {
+      if (!list2 || list2.len === 0) {
+        return 0;
+      }
+      const sampleSize = Math.ceil(list2.length * frequency);
+      const messageIndices = /* @__PURE__ */ new Set();
+      while (messageIndices.size < sampleSize && messageIndices.size < list2.length) {
+        const randomIndex = Math.floor(Math.random() * list2.length);
+        messageIndices.add(randomIndex);
+      }
+      const totalSize = Array.from(messageIndices).reduce((size, index) => {
+        return size + JSON.stringify(list2[index]).length;
+      }, 0);
+      const estimatedTotalSize = totalSize / sampleSize * list2.length;
+      return estimatedTotalSize;
+    }
   }
 });
 export default require_assets();

diff --git a/src/inspect_ai/_view/www/src/index.js b/src/inspect_ai/_view/www/src/index.js
@@ -10,7 +10,7 @@ import { throttle } from "./utils/sync.mjs";
 const vscode = getVscodeApi();
 let initialState = undefined;
 if (vscode) {
-  initialState = vscode.getState();
+  initialState = filterState(vscode.getState());
 }
 
 render(
@@ -20,9 +20,80 @@ render(
     saveInitialState=${throttle((state) => {
       const vscode = getVscodeApi();
       if (vscode) {
-        vscode.setState(state);
+        vscode.setState(filterState(state));
       }
     }, 1000)}
   />`,
   document.getElementById("app"),
 );
+
+function filterState(state) {
+  if (!state) {
+    return state;
+  }
+
+  // When saving state, we can't store vast amounts of data (like a large sample)
+  const filters = [filterLargeSample, filterLargeSelectedLog];
+  return filters.reduce(
+    (filteredState, filter) => filter(filteredState),
+    state,
+  );
+}
+
+// Filters the selected Sample if it is large
+function filterLargeSample(state) {
+  if (!state || !state.selectedSample) {
+    return state;
+  }
+
+  const estimatedTotalSize = estimateSize(state.selectedSample.messages);
+  if (estimatedTotalSize > 400000) {
+    const { selectedSample, ...filteredState } = state; // eslint-disable-line
+    return filteredState;
+  } else {
+    return state;
+  }
+}
+
+// Filters the selectedlog if it is too large
+function filterLargeSelectedLog(state) {
+  if (!state || !state.selectedLog?.contents) {
+    return state;
+  }
+
+  const estimatedSize = estimateSize(state.selectedLog.contents.sampleSummaries);
+  if (estimatedSize > 400000) {
+    const { selectedLog, ...filteredState } = state; // eslint-disable-line
+    return filteredState;
+  } else {
+    return state;
+  }
+}
+
+function estimateSize(list, frequency=0.2) {
+  if (!list || list.len === 0) {
+    return 0;
+  }
+
+  // Total number of samples
+  const sampleSize = Math.ceil(list.length * frequency);
+
+  // Get a proper random sample without duplicates
+  const messageIndices = new Set();
+  while (
+    messageIndices.size < sampleSize &&
+    messageIndices.size < list.length
+  ) {
+    const randomIndex = Math.floor(Math.random() * list.length);
+    messageIndices.add(randomIndex);
+  }
+
+  // Calculate size from sampled messages
+  const totalSize = Array.from(messageIndices).reduce((size, index) => {
+    return size + JSON.stringify(list[index]).length;
+  }, 0);
+
+  // Estimate total size based on sample
+  const estimatedTotalSize = (totalSize / sampleSize) * list.length;
+  return estimatedTotalSize;
+}