inspect-ai 0.3.81__py3-none-any.whl → 0.3.82__py3-none-any.whl

inspect_ai/_view/www/src/state/samplePolling.ts

@@ -0,0 +1,311 @@
+import {
+  AttachmentData,
+  EventData,
+  SampleData,
+  SampleSummary,
+} from "../api/types";
+import { Event } from "../types";
+import { resolveAttachments } from "../utils/attachments";
+import { createLogger } from "../utils/logger";
+import { createPolling } from "../utils/polling";
+import { resolveSample } from "./sampleUtils"; // Import the shared utility
+import { StoreState } from "./store";
+
+const log = createLogger("samplePolling");
+
+const kNoId = -1;
+const kPollingInterval = 2;
+const kPollingMaxRetries = 10;
+
+// Keeps the state for polling (the last ids for events
+// and attachments, the attachments and events, and
+// a mapping from eventIds to event indexes to enable
+// replacing events)
+interface PollingState {
+  eventId: number;
+  attachmentId: number;
+
+  attachments: Record<string, string>;
+
+  eventMapping: Record<string, number>;
+  events: Event[];
+}
+
+export function createSamplePolling(
+  get: () => StoreState,
+  set: (fn: (state: StoreState) => void) => void,
+) {
+  // The polling function that will be returned
+  let currentPolling: ReturnType<typeof createPolling> | null = null;
+
+  // handle aborts
+  let abortController: AbortController;
+
+  // The inintial polling state
+  const pollingState: PollingState = {
+    eventId: kNoId,
+    attachmentId: kNoId,
+
+    eventMapping: {},
+    attachments: {},
+    events: [],
+  };
+
+  // Function to start polling for a specific log file
+  const startPolling = (logFile: string, summary: SampleSummary) => {
+    // Create a unique identifier for this polling session
+    const pollingId = `${logFile}:${summary.id}-${summary.epoch}`;
+    log.debug(`Start Polling ${pollingId}`);
+
+    // If we're already polling this resource, don't restart
+    if (currentPolling && currentPolling.name === pollingId) {
+      log.debug(`Aleady polling, ignoring start`);
+      return;
+    }
+
+    // Stop any existing polling first
+    if (currentPolling) {
+      log.debug(`Resetting existing polling`);
+      currentPolling.stop();
+
+      // Clear any current running events
+      set((state) => {
+        state.sample.runningEvents = [];
+      });
+
+      // Reset the current polling state
+      resetPollingState(pollingState);
+    }
+    abortController = new AbortController();
+
+    // Create the polling callback
+    log.debug(`Polling sample: ${summary.id}-${summary.epoch}`);
+    const pollCallback = async () => {
+      const state = get();
+
+      // Get the api
+      const api = state.api;
+      if (!api) {
+        throw new Error("Required API is missing");
+      }
+
+      if (!api.get_log_sample_data) {
+        throw new Error("Required API get_log_sample_data is undefined.");
+      }
+
+      if (abortController.signal.aborted) {
+        return false;
+      }
+
+      // Fetch sample data
+      const eventId = pollingState.eventId;
+      const attachmentId = pollingState.attachmentId;
+      const sampleDataResponse = await api.get_log_sample_data(
+        logFile,
+        summary.id,
+        summary.epoch,
+        eventId,
+        attachmentId,
+      );
+
+      if (abortController.signal.aborted) {
+        return false;
+      }
+
+      if (sampleDataResponse?.status === "NotFound") {
+        // A 404 from the server means that this sample
+        // has been flushed to the main eval file, no events
+        // are available and we should retrieve the data from the
+        // sample file itself.
+
+        // Stop polling since we now have the complete sample
+        stopPolling();
+
+        // Also fetch a fresh sample and clear the runnning Events
+        // (if there were ever running events)
+        if (state.sample.runningEvents.length > 0) {
+          try {
+            log.debug(
+              `LOADING COMPLETED SAMPLE AFTER FLUSH: ${summary.id}-${summary.epoch}`,
+            );
+            const sample = await api.get_log_sample(
+              logFile,
+              summary.id,
+              summary.epoch,
+            );
+
+            if (sample) {
+              const migratedSample = resolveSample(sample);
+
+              // Update the store with the completed sample
+              set((state) => {
+                state.sample.selectedSample = migratedSample;
+                state.sample.sampleStatus = "ok";
+                state.sample.runningEvents = [];
+              });
+            } else {
+              set((state) => {
+                state.sample.sampleStatus = "error";
+                state.sample.sampleError = new Error(
+                  "Unable to load sample - an unknown error occurred",
+                );
+                state.sample.runningEvents = [];
+              });
+            }
+          } catch (e) {
+            set((state) => {
+              state.sample.sampleError = e as Error;
+              state.sample.sampleStatus = "error";
+              state.sample.runningEvents = [];
+            });
+          }
+        }
+
+        return false;
+      }
+
+      if (
+        sampleDataResponse?.status === "OK" &&
+        sampleDataResponse.sampleData
+      ) {
+        if (abortController.signal.aborted) {
+          return false;
+        }
+
+        if (sampleDataResponse.sampleData) {
+          // Process attachments
+          processAttachments(sampleDataResponse.sampleData, pollingState);
+
+          // Process events
+          const processedEvents = processEvents(
+            sampleDataResponse.sampleData,
+            pollingState,
+          );
+
+          // update max attachment id
+          if (sampleDataResponse.sampleData.attachments.length > 0) {
+            const maxAttachment = findMaxId(
+              sampleDataResponse.sampleData.attachments,
+              pollingState.attachmentId,
+            );
+            log.debug(`New max attachment ${maxAttachment}`);
+            pollingState.attachmentId = maxAttachment;
+          }
+
+          // update max event id
+          if (sampleDataResponse.sampleData.events.length > 0) {
+            const maxEvent = findMaxId(
+              sampleDataResponse.sampleData.events,
+              pollingState.eventId,
+            );
+            log.debug(`New max event ${maxEvent}`);
+            pollingState.eventId = maxEvent;
+          }
+
+          // Update the running events (ensure identity of runningEvents fails equality)
+          if (processedEvents) {
+            set((state) => {
+              state.sample.runningEvents = [...pollingState.events];
+            });
+          }
+        }
+      }
+
+      // Continue polling
+      return true;
+    };
+
+    // Create the polling instance
+    const polling = createPolling(pollingId, pollCallback, {
+      maxRetries: kPollingMaxRetries,
+      interval: kPollingInterval,
+    });
+
+    // Store the polling instance and start it
+    currentPolling = polling;
+    polling.start();
+  };
+
+  // Stop polling
+  const stopPolling = () => {
+    if (currentPolling) {
+      currentPolling.stop();
+      currentPolling = null;
+    }
+  };
+
+  const cleanup = () => {
+    log.debug(`CLEANUP`);
+    abortController.abort();
+    stopPolling();
+  };
+
+  return {
+    startPolling,
+    stopPolling,
+    cleanup,
+  };
+}
+
+const resetPollingState = (state: PollingState) => {
+  state.eventId = -1;
+  state.attachmentId = -1;
+  state.eventMapping = {};
+  state.attachments = {};
+  state.events = [];
+};
+
+function processAttachments(
+  sampleData: SampleData,
+  pollingState: PollingState,
+) {
+  log.debug(`Processing ${sampleData.attachments.length} attachments`);
+  Object.values(sampleData.attachments).forEach((v) => {
+    pollingState.attachments[v.hash] = v.content;
+  });
+}
+
+function processEvents(sampleData: SampleData, pollingState: PollingState) {
+  // Go through each event and resolve it, either appending or replacing
+  log.debug(`Processing ${sampleData.events.length} events`);
+  if (sampleData.events.length === 0) {
+    return false;
+  }
+
+  for (const eventData of sampleData.events) {
+    // Identify if this event id already has an event in the event list
+    const existingIndex = pollingState.eventMapping[eventData.event_id];
+
+    // Resolve attachments within this event
+    const resolvedEvent = resolveAttachments<Event>(
+      eventData.event,
+      pollingState.attachments,
+    );
+
+    if (existingIndex) {
+      // There is an existing event in the stream, replace it
+      log.debug(`Replace event ${existingIndex}`);
+      pollingState.events[existingIndex] = resolvedEvent;
+    } else {
+      // This is a new event, add to the event list and note
+      // its position
+      log.debug(`New event ${pollingState.events.length}`);
+
+      const currentIndex = pollingState.events.length;
+      pollingState.eventMapping[eventData.event_id] = currentIndex;
+      pollingState.events.push(resolvedEvent);
+    }
+  }
+  return true;
+}
+
+const findMaxId = (
+  items: EventData[] | AttachmentData[],
+  currentMax: number,
+) => {
+  if (items.length > 0) {
+    const newMax = Math.max(...items.map((i) => i.id), currentMax);
+    return newMax;
+  }
+  return currentMax;
+};

inspect_ai/_view/www/src/state/sampleSlice.ts

@@ -0,0 +1,127 @@
+import { SampleSummary } from "../api/types";
+import { kSampleMessagesTabId } from "../constants";
+import { SampleState, SampleStatus } from "../types";
+import { EvalSample } from "../types/log";
+import { createLogger } from "../utils/logger";
+import { createSamplePolling } from "./samplePolling";
+import { resolveSample } from "./sampleUtils"; // Import the shared utility
+import { StoreState } from "./store";
+
+const log = createLogger("sampleSlice");
+
+export interface SampleSlice {
+  sample: SampleState;
+  sampleActions: {
+    // The actual sample data
+    setSelectedSample: (sample: EvalSample) => void;
+    clearSelectedSample: () => void;
+    setSampleStatus: (status: SampleStatus) => void;
+    setSampleError: (error: Error | undefined) => void;
+
+    // Loading
+    loadSample: (
+      logFile: string,
+      sampleSummary: SampleSummary,
+    ) => Promise<void>;
+  };
+}
+
+const initialState: SampleState = {
+  selectedSample: undefined,
+  sampleStatus: "ok",
+  sampleError: undefined,
+
+  // The resolved events
+  runningEvents: [],
+};
+
+export const createSampleSlice = (
+  set: (fn: (state: StoreState) => void) => void,
+  get: () => StoreState,
+  _store: any,
+): [SampleSlice, () => void] => {
+  // The sample poller
+  const samplePolling = createSamplePolling(get, set);
+
+  const slice = {
+    // Actions
+    sample: initialState,
+    sampleActions: {
+      setSelectedSample: (sample: EvalSample) => {
+        set((state) => {
+          state.sample.selectedSample = sample;
+        });
+        if (sample.events.length < 1) {
+          // If there are no events, use the messages tab as the default
+          get().appActions.setSampleTab(kSampleMessagesTabId);
+        }
+      },
+      clearSelectedSample: () =>
+        set((state) => {
+          state.sample.selectedSample = undefined;
+        }),
+      setSampleStatus: (status: SampleStatus) =>
+        set((state) => {
+          state.sample.sampleStatus = status;
+        }),
+      setSampleError: (error: Error | undefined) =>
+        set((state) => {
+          state.sample.sampleError = error;
+        }),
+      loadSample: async (logFile: string, sampleSummary: SampleSummary) => {
+        const sampleActions = get().sampleActions;
+
+        sampleActions.setSampleError(undefined);
+        sampleActions.setSampleStatus("loading");
+        try {
+          if (sampleSummary.completed !== false) {
+            log.debug(
+              `LOADING COMPLETED SAMPLE: ${sampleSummary.id}-${sampleSummary.epoch}`,
+            );
+            const sample = await get().api?.get_log_sample(
+              logFile,
+              sampleSummary.id,
+              sampleSummary.epoch,
+            );
+            if (sample) {
+              const migratedSample = resolveSample(sample);
+              sampleActions.setSelectedSample(migratedSample);
+              sampleActions.setSampleStatus("ok");
+            } else {
+              sampleActions.setSampleStatus("error");
+              throw new Error(
+                "Unable to load sample - an unknown error occurred",
+              );
+            }
+          } else {
+            log.debug(
+              `POLLING RUNNING SAMPLE: ${sampleSummary.id}-${sampleSummary.epoch}`,
+            );
+
+            // Poll running sample
+            samplePolling.startPolling(logFile, sampleSummary);
+            sampleActions.setSampleStatus("streaming");
+          }
+        } catch (e) {
+          sampleActions.setSampleError(e as Error);
+          sampleActions.setSampleStatus("error");
+        }
+      },
+    },
+  } as const;
+
+  const cleanup = () => {
+    samplePolling.cleanup();
+  };
+  return [slice, cleanup];
+};
+
+export const initializeSampleSlice = (
+  set: (fn: (state: StoreState) => void) => void,
+) => {
+  set((state) => {
+    if (!state.sample) {
+      state.sample = initialState;
+    }
+  });
+};

inspect_ai/_view/www/src/state/sampleUtils.ts

@@ -0,0 +1,21 @@
+import { EvalSample } from "../types/log";
+import { resolveAttachments } from "../utils/attachments";
+
+/**
+ * Migrates and resolves attachments for a sample
+ */
+export const resolveSample = (sample: any): EvalSample => {
+  sample = { ...sample };
+
+  // Migrates old versions of samples to the new structure
+  if (sample.transcript) {
+    sample.events = sample.transcript.events;
+    sample.attachments = sample.transcript.content;
+  }
+  sample.attachments = sample.attachments || {};
+  sample.input = resolveAttachments(sample.input, sample.attachments);
+  sample.messages = resolveAttachments(sample.messages, sample.attachments);
+  sample.events = resolveAttachments(sample.events, sample.attachments);
+  sample.attachments = {};
+  return sample;
+};

inspect_ai/_view/www/src/state/scrolling.ts

@@ -0,0 +1,206 @@
+import { RefObject, useCallback, useEffect, useRef } from "react";
+import { StateCallback, StateSnapshot, VirtuosoHandle } from "react-virtuoso";
+import { createLogger } from "../utils/logger";
+import { debounce } from "../utils/sync";
+import { useStore } from "./store";
+
+const log = createLogger("scrolling");
+
+export function useStatefulScrollPosition<
+  T extends HTMLElement = HTMLDivElement,
+>(
+  elementRef: RefObject<T | null>,
+  elementKey: string,
+  delay = 500,
+  scrollable = true,
+) {
+  const getScrollPosition = useStore(
+    (state) => state.appActions.getScrollPosition,
+  );
+  const setScrollPosition = useStore(
+    (state) => state.appActions.setScrollPosition,
+  );
+
+  // Create debounced scroll handler
+  const handleScroll = useCallback(
+    debounce((e: Event) => {
+      const target = e.target as HTMLElement;
+      const position = target.scrollTop;
+      log.debug(`Storing scroll position`, elementKey, position);
+      setScrollPosition(elementKey, position);
+    }, delay),
+    [elementKey, setScrollPosition, delay],
+  );
+
+  // Function to manually restore scroll position
+  const restoreScrollPosition = useCallback(() => {
+    const element = elementRef.current;
+    const savedPosition = getScrollPosition(elementKey);
+
+    if (element && savedPosition !== undefined) {
+      requestAnimationFrame(() => {
+        element.scrollTop = savedPosition;
+
+        requestAnimationFrame(() => {
+          if (element.scrollTop !== savedPosition) {
+            element.scrollTop = savedPosition;
+          }
+        });
+      });
+    }
+  }, [elementKey, getScrollPosition, elementRef]);
+
+  // Set up scroll listener and restore position on mount
+  useEffect(() => {
+    const element = elementRef.current;
+    if (!element || !scrollable) {
+      return;
+    }
+    log.debug(`Restore Scroll Hook`, elementKey);
+
+    // Restore scroll position on mount
+    const savedPosition = getScrollPosition(elementKey);
+    if (savedPosition !== undefined) {
+      log.debug(`Restoring scroll position`, savedPosition);
+      // Ensure the element has fully rendered
+      requestAnimationFrame(() => {
+        if (element.scrollTop !== savedPosition) {
+          element.scrollTop = savedPosition;
+        }
+      });
+    }
+
+    // Set up scroll listener
+    if (element.addEventListener) {
+      element.addEventListener("scroll", handleScroll);
+    } else {
+      log.warn("Element has no way to add event listener", element);
+    }
+
+    // Clean up
+    return () => {
+      if (element.removeEventListener) {
+        element.removeEventListener("scroll", handleScroll);
+      } else {
+        log.warn("Element has no way to remove event listener", element);
+      }
+    };
+  }, [elementKey, elementRef, handleScroll]);
+
+  return { restoreScrollPosition };
+}
+
+// Define a type for the debounced function that includes the cancel method
+type DebouncedFunction<T extends (...args: any[]) => any> = T & {
+  cancel: () => void;
+  flush: () => void;
+};
+
+export const useVirtuosoState = (
+  virtuosoRef: RefObject<VirtuosoHandle | null>,
+  elementKey: string,
+  delay = 1000,
+) => {
+  // Use useCallback to stabilize the selectors
+  const restoreState = useStore(
+    useCallback((state) => state.app.listPositions[elementKey], [elementKey]),
+  );
+
+  const setListPosition = useStore(
+    useCallback((state) => state.appActions.setListPosition, []),
+  );
+
+  const clearListPosition = useStore(
+    useCallback((state) => state.appActions.clearListPosition, []),
+  );
+
+  // Properly type the debounced function ref
+  const debouncedFnRef = useRef<DebouncedFunction<
+    (isScrolling: boolean) => void
+  > | null>(null);
+
+  // Create the state change handler
+  const handleStateChange: StateCallback = useCallback(
+    (state: StateSnapshot) => {
+      log.debug(`Storing list state: [${elementKey}]`, state);
+      setListPosition(elementKey, state);
+    },
+    [elementKey, setListPosition],
+  );
+
+  // Setup the debounced function once
+  useEffect(() => {
+    debouncedFnRef.current = debounce((isScrolling: boolean) => {
+      log.debug("List scroll", isScrolling);
+      const element = virtuosoRef.current;
+      if (!element) {
+        return;
+      }
+      element.getState(handleStateChange);
+    }, delay) as DebouncedFunction<(isScrolling: boolean) => void>;
+
+    return () => {
+      // Clear the stored position when component unmounts
+      clearListPosition(elementKey);
+    };
+  }, [delay, elementKey, handleStateChange, clearListPosition, virtuosoRef]);
+
+  // Return a stable function reference that uses the ref internally
+  const isScrolling = useCallback((scrolling: boolean) => {
+    if (!scrolling) {
+      return;
+    }
+
+    if (debouncedFnRef.current) {
+      debouncedFnRef.current(scrolling);
+    }
+  }, []);
+
+  // Use a state to prevent re-rendering just because the restore
+  // state changes
+  const stateRef = useRef(restoreState);
+  useEffect(() => {
+    stateRef.current = restoreState;
+  }, [restoreState]);
+
+  const getRestoreState = useCallback(() => stateRef.current, []);
+
+  return { getRestoreState, isScrolling };
+};
+
+export function useRafThrottle<T extends (...args: any[]) => any>(
+  callback: T,
+  dependencies: any[] = [],
+): (...args: Parameters<T>) => void {
+  const rafRef = useRef<number | null>(null);
+  const callbackRef = useRef<T>(callback);
+
+  // Update the callback ref when the callback changes
+  useEffect(() => {
+    callbackRef.current = callback;
+  }, [callback, ...dependencies]);
+
+  const throttledCallback = useCallback((...args: Parameters<T>) => {
+    // Skip if we already have a frame queued
+    if (rafRef.current) {
+      return;
+    }
+
+    rafRef.current = requestAnimationFrame(() => {
+      callbackRef.current(...args);
+      rafRef.current = null;
+    });
+  }, []);
+
+  // Clean up any pending animation frame on unmount
+  useEffect(() => {
+    return () => {
+      if (rafRef.current) {
+        cancelAnimationFrame(rafRef.current);
+        rafRef.current = null;
+      }
+    };
+  }, []);
+
+  return throttledCallback;
+}