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

inspect_ai/_view/www/src/samples/transcript/TranscriptView.tsx

@@ -1,4 +1,4 @@
-import React, { RefObject, useCallback, useState } from "react";
+import { FC, memo, RefObject, useMemo } from "react";
 import { Events } from "../../types/log";
 import { ApprovalEventView } from "./ApprovalEventView";
 import { ErrorEventView } from "./ErrorEventView";
@@ -8,17 +8,17 @@ import { LoggerEventView } from "./LoggerEventView";
 import { ModelEventView } from "./ModelEventView";
 import { SampleInitEventView } from "./SampleInitEventView";
 import { SampleLimitEventView } from "./SampleLimitEventView";
+import { SandboxEventView } from "./SandboxEventView";
 import { ScoreEventView } from "./ScoreEventView";
 import { StateEventView } from "./state/StateEventView";
 import { StepEventView } from "./StepEventView";
 import { SubtaskEventView } from "./SubtaskEventView";
 import { ToolEventView } from "./ToolEventView";
-import { EventNode, EventType, TranscriptEventState } from "./types";
+import { EventNode, EventType } from "./types";
 
 import clsx from "clsx";
-import { Virtuoso } from "react-virtuoso";
-import { SandboxEventView } from "./SandboxEventView";
 import styles from "./TranscriptView.module.css";
+import { TranscriptVirtualListComponent } from "./TranscriptVirtualListComponent";
 
 interface TranscriptViewProps {
   id: string;
@@ -26,38 +26,21 @@ interface TranscriptViewProps {
   depth: number;
 }
 
-type TranscriptState = Record<string, TranscriptEventState>;
-
 /**
  * Renders the TranscriptView component.
  */
-export const TranscriptView: React.FC<TranscriptViewProps> = ({
+export const TranscriptView: FC<TranscriptViewProps> = ({
   id,
   events,
   depth,
 }) => {
-  const [transcriptState, setTranscriptState] = useState<TranscriptState>({});
-  const onTranscriptState = useCallback(
-    (state: TranscriptState) => {
-      setTranscriptState(state);
-    },
-    [setTranscriptState],
-  );
-
   // Normalize Events themselves
   const resolvedEvents = fixupEventStream(events);
   const eventNodes = treeifyEvents(
     resolvedEvents,
     depth !== undefined ? depth : 0,
   );
-  return (
-    <TranscriptComponent
-      id={id}
-      eventNodes={eventNodes}
-      transcriptState={transcriptState}
-      setTranscriptState={onTranscriptState}
-    />
-  );
+  return <TranscriptComponent id={id} eventNodes={eventNodes} />;
 };
 
 interface TranscriptVirtualListProps {
@@ -65,367 +48,243 @@ interface TranscriptVirtualListProps {
   events: Events;
   depth?: number;
   scrollRef: RefObject<HTMLDivElement | null>;
+  running?: boolean;
 }
 
 /**
  * Renders the Transcript Virtual List.
  */
-export const TranscriptVirtualList: React.FC<TranscriptVirtualListProps> = (
-  props,
-) => {
-  let { id, scrollRef, events, depth } = props;
-
-  // Normalize Events themselves
-  const resolvedEvents = fixupEventStream(events);
-  const eventNodes = treeifyEvents(resolvedEvents, depth || 0);
-
-  const [transcriptState, setTranscriptState] = useState({});
-  const onTranscriptState = useCallback(
-    (state: TranscriptEventState) => {
-      setTranscriptState(state);
-    },
-    [setTranscriptState],
-  );
-
-  return (
-    <TranscriptVirtualListComponent
-      id={id}
-      eventNodes={eventNodes}
-      scrollRef={scrollRef}
-      transcriptState={transcriptState}
-      setTranscriptState={onTranscriptState}
-    />
-  );
-};
-
-interface TranscriptVirtualListComponentProps {
-  id: string;
-  eventNodes: EventNode[];
-  transcriptState: TranscriptState;
-  setTranscriptState: (state: TranscriptState) => void;
-  scrollRef?: RefObject<HTMLDivElement | null>;
-}
-
-/**
- * Renders the Transcript component.
- */
-export const TranscriptVirtualListComponent: React.FC<
-  TranscriptVirtualListComponentProps
-> = ({ id, eventNodes, scrollRef, transcriptState, setTranscriptState }) => {
-  const setEventState = useCallback(
-    (eventId: string, state: TranscriptEventState) => {
-      setTranscriptState({ ...transcriptState, [eventId]: state });
-    },
-    [transcriptState, setTranscriptState],
-  );
+export const TranscriptVirtualList: FC<TranscriptVirtualListProps> = memo(
+  (props) => {
+    let { id, scrollRef, events, depth, running } = props;
 
-  const [followOutput, setFollowOutput] = useState(false);
+    // Normalize Events themselves
+    const eventNodes = useMemo(() => {
+      const resolvedEvents = fixupEventStream(events, !running);
+      const eventNodes = treeifyEvents(resolvedEvents, depth || 0);
 
-  const renderRow = (item: EventNode, index: number) => {
-    const bgClass = item.depth % 2 == 0 ? styles.darkenedBg : styles.normalBg;
-    const paddingClass = index === 0 ? styles.first : undefined;
-
-    const eventId = `${id}-event${index}`;
+      return eventNodes;
+    }, [events, depth]);
 
     return (
-      <div key={eventId} className={clsx(styles.node, paddingClass)}>
-        <RenderedEventNode
-          id={eventId}
-          node={item}
-          className={clsx(bgClass)}
-          scrollRef={scrollRef}
-          eventState={transcriptState[eventId] || {}}
-          setEventState={(state) => setEventState(eventId, state)}
-        />
-      </div>
+      <TranscriptVirtualListComponent
+        id={id}
+        eventNodes={eventNodes}
+        scrollRef={scrollRef}
+        running={running}
+      />
     );
-  };
-
-  return (
-    <Virtuoso
-      customScrollParent={scrollRef?.current ? scrollRef.current : undefined}
-      style={{ height: "100%", width: "100%" }}
-      data={eventNodes}
-      itemContent={(index: number, data: EventNode) => {
-        return renderRow(data, index);
-      }}
-      increaseViewportBy={{ top: 1000, bottom: 1000 }}
-      overscan={{
-        main: 10,
-        reverse: 10,
-      }}
-      followOutput={followOutput}
-      atBottomStateChange={(atBottom: boolean) => {
-        setFollowOutput(atBottom);
-      }}
-      skipAnimationFrameInResizeObserver={true}
-      className={clsx("transcript")}
-    />
-  );
-};
+  },
+);
 
 interface TranscriptComponentProps {
   id: string;
-  transcriptState: TranscriptState;
-  setTranscriptState: (state: TranscriptState) => void;
   eventNodes: EventNode[];
 }
 /**
  * Renders the Transcript component.
  */
-export const TranscriptComponent: React.FC<TranscriptComponentProps> = ({
-  id,
-  transcriptState,
-  setTranscriptState,
-  eventNodes,
-}) => {
-  const setEventState = useCallback(
-    (state: TranscriptEventState, eventId: string) => {
-      setTranscriptState({ ...transcriptState, [eventId]: state });
-    },
-    [setTranscriptState, transcriptState],
-  );
-
-  const rows = eventNodes.map((eventNode, index) => {
-    const clz = [styles.eventNode];
-    if (eventNode.depth % 2 == 0) {
-      clz.push(styles.darkenBg);
-    }
-    if (index === eventNodes.length - 1) {
-      clz.push(styles.lastNode);
-    }
+export const TranscriptComponent: FC<TranscriptComponentProps> = memo(
+  ({ id, eventNodes }) => {
+    const rows = eventNodes.map((eventNode, index) => {
+      const clz = [styles.eventNode];
+      if (eventNode.depth % 2 == 0) {
+        clz.push(styles.darkenBg);
+      }
+      if (index === eventNodes.length - 1) {
+        clz.push(styles.lastNode);
+      }
 
-    const eventId = `${id}-event${index}`;
+      const eventId = `${id}|event|${index}`;
+
+      const row = (
+        <div
+          key={eventId}
+          className={clsx(
+            styles.eventNodeContainer,
+            index === eventNodes.length - 1 ? styles.noBottom : undefined,
+          )}
+        >
+          <RenderedEventNode
+            id={eventId}
+            node={eventNode}
+            className={clsx(clz)}
+          />
+        </div>
+      );
+      return row;
+    });
 
-    const row = (
+    return (
       <div
-        key={eventId}
-        className={clsx(
-          styles.eventNodeContainer,
-          index === eventNodes.length - 1 ? styles.noBottom : undefined,
-        )}
+        id={id}
+        className={clsx("text-size-small", styles.transcriptComponent)}
       >
-        <RenderedEventNode
-          id={eventId}
-          node={eventNode}
-          className={clsx(clz)}
-          eventState={transcriptState[eventId] || {}}
-          setEventState={(state: TranscriptEventState) => {
-            setEventState(state, eventId);
-          }}
-        />
+        {rows}
       </div>
     );
-    return row;
-  });
-
-  return (
-    <div
-      id={id}
-      className={clsx("text-size-small", styles.transcriptComponent)}
-    >
-      {rows}
-    </div>
-  );
-};
+  },
+);
 
 interface RenderedEventNodeProps {
   id: string;
   node: EventNode;
-  scrollRef?: RefObject<HTMLDivElement | null>;
-  eventState: TranscriptEventState;
-  setEventState: (state: TranscriptEventState) => void;
   className?: string | string[];
 }
 /**
  * Renders the event based on its type.
  */
-export const RenderedEventNode: React.FC<RenderedEventNodeProps> = ({
-  id,
-  node,
-  scrollRef,
-  eventState,
-  setEventState,
-  className,
-}) => {
-  switch (node.event.event) {
-    case "sample_init":
-      return (
-        <SampleInitEventView
-          id={id}
-          event={node.event}
-          eventState={eventState}
-          setEventState={setEventState}
-          className={className}
-        />
-      );
-
-    case "sample_limit":
-      return (
-        <SampleLimitEventView
-          id={id}
-          event={node.event}
-          eventState={eventState}
-          setEventState={setEventState}
-          className={className}
-        />
-      );
-
-    case "info":
-      return (
-        <InfoEventView
-          id={id}
-          event={node.event}
-          eventState={eventState}
-          setEventState={setEventState}
-          className={className}
-        />
-      );
-
-    case "logger":
-      return <LoggerEventView event={node.event} className={className} />;
-
-    case "model":
-      return (
-        <ModelEventView
-          id={id}
-          event={node.event}
-          eventState={eventState}
-          setEventState={setEventState}
-          className={className}
-        />
-      );
-
-    case "score":
-      return (
-        <ScoreEventView
-          id={id}
-          event={node.event}
-          eventState={eventState}
-          setEventState={setEventState}
-          className={className}
-        />
-      );
-
-    case "state":
-      return (
-        <StateEventView
-          id={id}
-          event={node.event}
-          eventState={eventState}
-          setEventState={setEventState}
-          className={className}
-        />
-      );
-
-    case "step":
-      return (
-        <StepEventView
-          event={node.event}
-          eventState={eventState}
-          setEventState={setEventState}
-          children={node.children}
-          scrollRef={scrollRef}
-          className={className}
-        />
-      );
-
-    case "store":
-      return (
-        <StateEventView
-          id={id}
-          event={node.event}
-          eventState={eventState}
-          setEventState={setEventState}
-          className={className}
-          isStore={true}
-        />
-      );
-
-    case "subtask":
-      return (
-        <SubtaskEventView
-          id={id}
-          event={node.event}
-          eventState={eventState}
-          setEventState={setEventState}
-          className={className}
-          depth={node.depth}
-        />
-      );
-
-    case "tool":
-      return (
-        <ToolEventView
-          id={id}
-          event={node.event}
-          eventState={eventState}
-          setEventState={setEventState}
-          className={className}
-          depth={node.depth}
-        />
-      );
-
-    case "input":
-      return (
-        <InputEventView
-          id={id}
-          event={node.event}
-          eventState={eventState}
-          setEventState={setEventState}
-          className={className}
-        />
-      );
-
-    case "error":
-      return (
-        <ErrorEventView
-          id={id}
-          event={node.event}
-          eventState={eventState}
-          setEventState={setEventState}
-          className={className}
-        />
-      );
-
-    case "approval":
-      return <ApprovalEventView event={node.event} className={className} />;
-
-    case "sandbox":
-      return (
-        <SandboxEventView
-          id={id}
-          event={node.event}
-          className={className}
-          eventState={eventState}
-          setEventState={setEventState}
-        />
-      );
-
-    default:
-      return null;
-  }
-};
+export const RenderedEventNode: FC<RenderedEventNodeProps> = memo(
+  ({ id, node, className }) => {
+    switch (node.event.event) {
+      case "sample_init":
+        return (
+          <SampleInitEventView
+            id={id}
+            event={node.event}
+            className={className}
+          />
+        );
+
+      case "sample_limit":
+        return (
+          <SampleLimitEventView
+            id={id}
+            event={node.event}
+            className={className}
+          />
+        );
+
+      case "info":
+        return (
+          <InfoEventView id={id} event={node.event} className={className} />
+        );
+
+      case "logger":
+        return <LoggerEventView event={node.event} className={className} />;
+
+      case "model":
+        return (
+          <ModelEventView id={id} event={node.event} className={className} />
+        );
+
+      case "score":
+        return (
+          <ScoreEventView id={id} event={node.event} className={className} />
+        );
+
+      case "state":
+        return (
+          <StateEventView id={id} event={node.event} className={className} />
+        );
+
+      case "step":
+        return (
+          <StepEventView
+            id={id}
+            event={node.event}
+            children={node.children}
+            className={className}
+          />
+        );
+
+      case "store":
+        return (
+          <StateEventView
+            id={id}
+            event={node.event}
+            className={className}
+            isStore={true}
+          />
+        );
+
+      case "subtask":
+        return (
+          <SubtaskEventView
+            id={id}
+            event={node.event}
+            className={className}
+            depth={node.depth}
+          />
+        );
+
+      case "tool":
+        return (
+          <ToolEventView
+            id={id}
+            event={node.event}
+            className={className}
+            depth={node.depth}
+          />
+        );
+
+      case "input":
+        return (
+          <InputEventView id={id} event={node.event} className={className} />
+        );
+
+      case "error":
+        return (
+          <ErrorEventView id={id} event={node.event} className={className} />
+        );
+
+      case "approval":
+        return <ApprovalEventView event={node.event} className={className} />;
+
+      case "sandbox":
+        return (
+          <SandboxEventView id={id} event={node.event} className={className} />
+        );
+
+      default:
+        return null;
+    }
+  },
+);
 
 /**
  * Normalizes event content
  */
-const fixupEventStream = (events: Events) => {
+const fixupEventStream = (events: Events, filterPending: boolean = true) => {
   const initEventIndex = events.findIndex((e) => {
     return e.event === "sample_init";
   });
   const initEvent = events[initEventIndex];
 
-  // Filter pending events
-  const finalEvents = events.filter((e) => !e.pending);
-
-  // See if the find an init step
+  // If filtering pending, just remove all pending events
+  // otherise, collapse sequential pending events of the same
+  // type
+  const collapsed = filterPending
+    ? events.filter((e) => !e.pending)
+    : events.reduce<Events>((acc, event) => {
+        // Collapse sequential pending events of the same type
+        if (!event.pending) {
+          // Not a pending event
+          acc.push(event);
+        } else {
+          // For pending events, replace previous pending or add new
+          const lastIndex = acc.length - 1;
+          if (
+            lastIndex >= 0 &&
+            acc[lastIndex].pending &&
+            acc[lastIndex].event === event.event
+          ) {
+            // Replace previous pending with current one (if they're of the same type)
+            acc[lastIndex] = event;
+          } else {
+            // First event or follows non-pending
+            acc.push(event);
+          }
+        }
+        return acc;
+      }, []);
+  // See if the events have an init step
   const hasInitStep =
-    events.findIndex((e) => {
+    collapsed.findIndex((e) => {
       return e.event === "step" && e.name === "init";
     }) !== -1;
 
-  const fixedUp = [...finalEvents];
+  const fixedUp = [...collapsed];
   if (!hasInitStep && initEvent) {
     fixedUp.splice(initEventIndex, 0, {
       timestamp: initEvent.timestamp,

inspect_ai/_view/www/src/samples/transcript/TranscriptVirtualListComponent.module.css

@@ -0,0 +1,16 @@
+.darkenedBg {
+  background-color: var(--bs-light-bg-subtle);
+}
+
+.normalBg {
+  background-color: var(--bs-body-bg);
+}
+
+.node.first {
+  padding-top: 0.5em;
+}
+
+.node {
+  padding-top: 0;
+  padding-bottom: 0.5em;
+}

inspect_ai/_view/www/src/samples/transcript/TranscriptVirtualListComponent.tsx

@@ -0,0 +1,44 @@
+import clsx from "clsx";
+import { FC, RefObject, useCallback } from "react";
+import { RenderedEventNode } from "./TranscriptView";
+import { EventNode } from "./types";
+
+import { LiveVirtualList } from "../../components/LiveVirtualList";
+import styles from "./TranscriptVirtualListComponent.module.css";
+
+interface TranscriptVirtualListComponentProps {
+  id: string;
+  eventNodes: EventNode[];
+  scrollRef?: RefObject<HTMLDivElement | null>;
+  running?: boolean;
+}
+
+/**
+ * Renders the Transcript component.
+ */
+export const TranscriptVirtualListComponent: FC<
+  TranscriptVirtualListComponentProps
+> = ({ id, eventNodes, scrollRef, running }) => {
+  const renderRow = useCallback((index: number, item: EventNode) => {
+    const bgClass = item.depth % 2 == 0 ? styles.darkenedBg : styles.normalBg;
+    const paddingClass = index === 0 ? styles.first : undefined;
+
+    const eventId = `${id}-event-${index}`;
+
+    return (
+      <div key={eventId} className={clsx(styles.node, paddingClass)}>
+        <RenderedEventNode id={eventId} node={item} className={clsx(bgClass)} />
+      </div>
+    );
+  }, []);
+
+  return (
+    <LiveVirtualList<EventNode>
+      id={id}
+      scrollRef={scrollRef}
+      data={eventNodes}
+      renderRow={renderRow}
+      live={running}
+    />
+  );
+};