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

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

@@ -0,0 +1,181 @@
+import { EvalLogHeader, LogFiles } from "../api/types";
+import { LogsState } from "../types";
+import { createLogger } from "../utils/logger";
+import { createLogsPolling } from "./logsPolling";
+import { StoreState } from "./store";
+
+const log = createLogger("Log Slice");
+
+export interface LogsSlice {
+  logs: LogsState;
+  logsActions: {
+    // Update State
+    setLogs: (logs: LogFiles) => void;
+    setLogHeaders: (headers: Record<string, EvalLogHeader>) => void;
+    setHeadersLoading: (loading: boolean) => void;
+    setSelectedLogIndex: (index: number) => void;
+    setSelectedLogFile: (logUrl: string) => void;
+    updateLogHeaders: (headers: Record<string, EvalLogHeader>) => void;
+
+    // Fetch or update logs
+    refreshLogs: () => Promise<void>;
+    selectLogFile: (logUrl: string) => Promise<void>;
+    loadLogs: () => Promise<LogFiles>;
+
+    // Computed values
+    getSelectedLogFile: () => string | undefined;
+  };
+}
+
+const initialState: LogsState = {
+  logs: { log_dir: "", files: [] },
+  logHeaders: {},
+  headersLoading: false,
+  selectedLogIndex: -1,
+};
+
+export const createLogsSlice = (
+  set: (fn: (state: StoreState) => void) => void,
+  get: () => StoreState,
+  _store: any,
+): [LogsSlice, () => void] => {
+  const logsPolling = createLogsPolling(get, set);
+
+  const slice = {
+    // State
+    logs: initialState,
+
+    // Actions
+    logsActions: {
+      setLogs: (logs: LogFiles) => {
+        set((state) => {
+          state.logs.logs = logs;
+        });
+
+        // If we have files in the logs, load the headers
+        if (logs.files.length > 0) {
+          // ensure state is updated first
+          setTimeout(() => {
+            const currentState = get();
+            if (!currentState.logs.headersLoading) {
+              logsPolling.startPolling(logs);
+            }
+          }, 100);
+        }
+      },
+      setLogHeaders: (headers: Record<string, EvalLogHeader>) =>
+        set((state) => {
+          state.logs.logHeaders = headers;
+        }),
+      setHeadersLoading: (loading: boolean) =>
+        set((state) => {
+          state.logs.headersLoading = loading;
+        }),
+      setSelectedLogIndex: (selectedLogIndex: number) => {
+        set((state) => {
+          state.logs.selectedLogIndex = selectedLogIndex;
+        });
+      },
+      updateLogHeaders: (headers: Record<string, EvalLogHeader>) =>
+        set((state) => {
+          state.logs.logHeaders = { ...get().logs.logHeaders, ...headers };
+        }),
+
+      setSelectedLogFile: (logUrl: string) => {
+        const state = get();
+        const index = state.logs.logs.files.findIndex((val) =>
+          logUrl.endsWith(val.name),
+        );
+
+        if (index > -1) {
+          state.logsActions.setSelectedLogIndex(index);
+        }
+      },
+
+      // Helper function to load logs
+      loadLogs: async () => {
+        const api = get().api;
+        if (!api) {
+          console.error("API not initialized in LogsStore");
+          return { log_dir: "", files: [] };
+        }
+
+        try {
+          log.debug("LOADING LOG FILES");
+          return await api.get_log_paths();
+        } catch (e) {
+          console.log(e);
+          get().appActions.setStatus({ loading: false, error: e as Error });
+          return { log_dir: "", files: [] };
+        }
+      },
+      refreshLogs: async () => {
+        log.debug("REFRESH LOGS");
+        const state = get();
+        const refreshedLogs = await state.logsActions.loadLogs();
+
+        // Set the logs first
+        state.logsActions.setLogs(refreshedLogs || { log_dir: "", files: [] });
+
+        // Preserve the selected log even if new logs appear
+        const currentLog =
+          refreshedLogs.files[
+            state.logs.selectedLogIndex > -1 ? state.logs.selectedLogIndex : 0
+          ];
+
+        if (currentLog) {
+          const newIndex = refreshedLogs?.files.findIndex((file) =>
+            currentLog.name.endsWith(file.name),
+          );
+
+          if (newIndex !== undefined && newIndex !== -1) {
+            state.logsActions.setSelectedLogIndex(newIndex);
+          }
+        }
+      },
+      // Select a specific log file
+      selectLogFile: async (logUrl: string) => {
+        const state = get();
+        const index = state.logs.logs.files.findIndex((val) =>
+          val.name.endsWith(logUrl),
+        );
+
+        // It is already loaded
+        if (index > -1) {
+          state.logsActions.setSelectedLogIndex(index);
+        } else {
+          // It isn't yet loaded, so refresh the logs and try to load it from there
+          const result = await state.logsActions.loadLogs();
+          const idx = result?.files.findIndex((file) =>
+            logUrl.endsWith(file.name),
+          );
+
+          state.logsActions.setLogs(result || { log_dir: "", files: [] });
+          state.logsActions.setSelectedLogIndex(
+            idx !== undefined && idx > -1 ? idx : 0,
+          );
+        }
+      },
+
+      getSelectedLogFile: () => {
+        const state = get();
+        const file = state.logs.logs.files[state.logs.selectedLogIndex];
+        return file !== undefined ? file.name : undefined;
+      },
+    },
+  } as const;
+
+  const cleanup = () => {};
+
+  return [slice, cleanup];
+};
+
+export const initializeLogsSlice = <T extends LogsSlice>(
+  set: (fn: (state: T) => void) => void,
+) => {
+  set((state) => {
+    if (!state.logs) {
+      state.logs = initialState;
+    }
+  });
+};

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

@@ -0,0 +1,314 @@
+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 ||
+          state.sample.sampleStatus === "streaming"
+        ) {
+          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,140 @@
+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>;
+
+    pollSample: (
+      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;
+        }),
+      pollSample: async (logFile: string, sampleSummary: SampleSummary) => {
+        // Poll running sample
+        const state = get();
+        if (state.log.loadedLog && state.sample.selectedSample) {
+          samplePolling.startPolling(logFile, sampleSummary);
+          state.sampleActions.setSampleStatus("streaming");
+        }
+      },
+      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;
+};