inspect-ai 0.3.99__py3-none-any.whl → 0.3.101__py3-none-any.whl

inspect_ai/_view/www/src/app/samples/transcript/state/StateEventView.tsx

@@ -18,7 +18,7 @@ import {
 
 import { FC, useEffect, useMemo } from "react";
 import { useStore } from "../../../../state/store";
-import { EventNode } from "../types";
+import { EventNode, kTranscriptCollapseScope } from "../types";
 import styles from "./StateEventView.module.css";
 
 interface StateEventViewProps {
@@ -35,7 +35,6 @@ export const StateEventView: FC<StateEventViewProps> = ({
   className,
 }) => {
   const event = eventNode.event;
-  const id = eventNode.id;
 
   const summary = useMemo(() => {
     return summarizeChanges(event.changes);
@@ -67,13 +66,13 @@ export const StateEventView: FC<StateEventViewProps> = ({
   const collapseEvent = useStore((state) => state.sampleActions.collapseEvent);
   useEffect(() => {
     if (changePreview === undefined) {
-      collapseEvent(id, true);
+      collapseEvent(kTranscriptCollapseScope, eventNode.id, true);
     }
   }, [changePreview, collapseEvent]);
 
   return (
     <EventPanel
-      id={id}
+      eventNodeId={eventNode.id}
       depth={eventNode.depth}
       title={title}
       className={className}

inspect_ai/_view/www/src/app/samples/transcript/transform/hooks.ts

@@ -0,0 +1,78 @@
+import { useMemo } from "react";
+import {
+  Events,
+  SpanBeginEvent,
+  StepEvent,
+  SubtaskEvent,
+  ToolEvent,
+} from "../../../../@types/log";
+import { EventNode } from "../types";
+import { fixupEventStream, kSandboxSignalName } from "./fixups";
+import { treeifyEvents } from "./treeify";
+
+export const useEventNodes = (events: Events, running: boolean) => {
+  // Normalize Events in a flattened filtered list
+  const { eventTree, defaultCollapsedIds } = useMemo((): {
+    eventTree: EventNode[];
+    defaultCollapsedIds: Record<string, true>;
+  } => {
+    // Apply fixups to the event string
+    const resolvedEvents = fixupEventStream(events, !running);
+
+    // Build the event tree
+    const eventTree = treeifyEvents(resolvedEvents, 0);
+
+    // Apply collapse filters to the event tree
+    const defaultCollapsedIds: Record<string, true> = {};
+    const findCollapsibleEvents = (nodes: EventNode[]) => {
+      for (const node of nodes) {
+        if (
+          (node.event.event === "step" ||
+            node.event.event === "span_begin" ||
+            node.event.event === "tool" ||
+            node.event.event === "subtask") &&
+          collapseFilters.some((filter) =>
+            filter(
+              node.event as
+                | StepEvent
+                | SpanBeginEvent
+                | ToolEvent
+                | SubtaskEvent,
+            ),
+          )
+        ) {
+          defaultCollapsedIds[node.id] = true;
+        }
+
+        // Recursively check children
+        findCollapsibleEvents(node.children);
+      }
+    };
+    findCollapsibleEvents(eventTree);
+
+    return { eventTree, defaultCollapsedIds };
+  }, [events, running]);
+
+  return { eventNodes: eventTree, defaultCollapsedIds };
+};
+
+const collapseFilters: Array<
+  (event: StepEvent | SpanBeginEvent | ToolEvent | SubtaskEvent) => boolean
+> = [
+  (event: StepEvent | SpanBeginEvent | ToolEvent | SubtaskEvent) =>
+    event.type === "solver" && event.name === "system_message",
+  (event: StepEvent | SpanBeginEvent | ToolEvent | SubtaskEvent) => {
+    if (event.event === "step" || event.event === "span_begin") {
+      return (
+        event.name === kSandboxSignalName ||
+        event.name === "init" ||
+        event.name === "sample_init"
+      );
+    }
+    return false;
+  },
+  (event: StepEvent | SpanBeginEvent | ToolEvent | SubtaskEvent) =>
+    event.event === "tool" && !event.agent && !event.failed,
+  (event: StepEvent | SpanBeginEvent | ToolEvent | SubtaskEvent) =>
+    event.event === "subtask",
+];

inspect_ai/_view/www/src/app/samples/transcript/transform/treeify.ts

@@ -1,4 +1,4 @@
-import { Events } from "../../../../@types/log";
+import { Events, SpanBeginEvent, SpanEndEvent } from "../../../../@types/log";
 import { EventNode, EventType } from "../types";
 import {
   ACTION_BEGIN,
@@ -11,7 +11,10 @@ import {
   TOOL,
   TYPE_AGENT,
   TYPE_HANDOFF,
+  TYPE_SCORER,
+  TYPE_SCORERS,
   TYPE_SOLVER,
+  TYPE_SOLVERS,
   TYPE_SUBTASK,
   TYPE_TOOL,
   hasSpans,
@@ -31,6 +34,9 @@ export function treeifyEvents(events: Events, depth: number): EventNode[] {
   const rootNodes: EventNode[] = [];
   const stack: EventNode[] = [];
 
+  // The function used to build the tree
+  const treeifyFn = getTreeifyFunction();
+
   const addNode = (event: EventType): EventNode => {
     const currentDepth = stack.length;
 
@@ -45,7 +51,11 @@ export function treeifyEvents(events: Events, depth: number): EventNode[] {
 
     // Create a new node
     const idPath = pathIndices.slice(0, currentDepth + 1).join(".");
-    const node = new EventNode(idPath, event, currentDepth + depth);
+    const node = new EventNode(
+      `event_node_${idPath}`,
+      event,
+      currentDepth + depth,
+    );
     if (stack.length > 0) {
       const parentNode = stack[stack.length - 1];
       parentNode.children.push(node);
@@ -65,6 +75,10 @@ export function treeifyEvents(events: Events, depth: number): EventNode[] {
     pathIndices.pop();
   };
 
+  // First inject spans that may be needed
+  events = injectScorersSpan(events);
+
+  // Now treeify the list
   events.forEach((event) => {
     treeifyFn(event, addNode, pushStack, popStack);
   });
@@ -76,156 +90,294 @@ export function treeifyEvents(events: Events, depth: number): EventNode[] {
   }
 }
 
-const treeifyFn: TreeifyFunction = (
-  event: EventType,
-  addNode: (event: EventType) => EventNode,
-  pushStack: (node: EventNode) => void,
-  popStack: () => void,
-): void => {
-  switch (event.event) {
-    case STEP:
-      if (event.action === ACTION_BEGIN) {
-        // Starting a new step
-        const node = addNode(event);
-        pushStack(node);
+// This injects a scorer span around top level scorer events if one
+// isn't already present
+const kBeginScorerId = "E617087FA405";
+const kEndScorerId = "C39922B09481";
+const kScorersSpanId = "C5A831026F2C";
+const injectScorersSpan = (events: Events): Events => {
+  const results: Events = [];
+  const collectedScorerEvents: Events = [];
+  let hasCollectedScorers = false;
+  let collecting: string | null = null;
+
+  const flushCollected = (): Events => {
+    if (collectedScorerEvents.length > 0) {
+      const beginSpan: SpanBeginEvent = {
+        name: "scorers",
+        id: kBeginScorerId,
+        span_id: kScorersSpanId,
+        event: SPAN_BEGIN,
+        type: TYPE_SCORERS,
+        timestamp: collectedScorerEvents[0].timestamp,
+        working_start: collectedScorerEvents[0].working_start,
+        pending: false,
+        parent_id: null,
+      };
+
+      const scoreEvents: Events = collectedScorerEvents.map((event) => {
+        return {
+          ...event,
+          parent_id:
+            event.event === "span_begin"
+              ? event.parent_id || kScorersSpanId
+              : null,
+        };
+      });
+
+      const endSpan: SpanEndEvent = {
+        id: kEndScorerId,
+        span_id: kScorersSpanId,
+        event: SPAN_END,
+        pending: false,
+        working_start:
+          collectedScorerEvents[collectedScorerEvents.length - 1].working_start,
+        timestamp:
+          collectedScorerEvents[collectedScorerEvents.length - 1].timestamp,
+      };
+
+      collectedScorerEvents.length = 0;
+      hasCollectedScorers = true;
+      return [beginSpan, ...scoreEvents, endSpan];
+    }
+    return [];
+  };
+
+  for (const event of events) {
+    // Return events immediately if the scorers span is present
+    if (event.event === SPAN_BEGIN && event.type === TYPE_SCORERS) {
+      return events;
+    }
+
+    if (
+      event.event === SPAN_BEGIN &&
+      event.type === TYPE_SCORER &&
+      !hasCollectedScorers
+    ) {
+      collecting = event.span_id;
+    }
+
+    // Look for the first scorer event and then begin
+    if (collecting) {
+      if (event.event === SPAN_END && event.span_id === collecting) {
+        collecting = null;
+        results.push(...flushCollected());
+        results.push(event);
       } else {
-        // An ending step
-        popStack();
+        collectedScorerEvents.push(event);
       }
-      break;
-    case SPAN_BEGIN: {
-      const node = addNode(event);
-      pushStack(node);
-      break;
-    }
-    case SPAN_END: {
-      popStack();
-      break;
+    } else {
+      results.push(event);
     }
-    case TOOL:
-      {
-        const node = addNode(event);
+  }
 
-        // In the span world, the first child will be a span of type tool
-        if (
-          event.events.length > 0 &&
-          (event.events[0].event !== SPAN_BEGIN ||
-            event.events[0].type !== TYPE_TOOL)
-        ) {
-          // Expand the children
+  return results;
+};
+
+const getTreeifyFunction = () => {
+  const treeifyFn: TreeifyFunction = (
+    event: EventType,
+    addNode: (event: EventType) => EventNode,
+    pushStack: (node: EventNode) => void,
+    popStack: () => void,
+  ): void => {
+    switch (event.event) {
+      case STEP:
+        if (event.action === ACTION_BEGIN) {
+          // Starting a new step
+          const node = addNode(event);
           pushStack(node);
-          for (const child of event.events) {
-            treeifyFn(child, addNode, pushStack, popStack);
-          }
+        } else {
+          // An ending step
           popStack();
         }
-      }
-
-      break;
-    case SUBTASK:
-      {
+        break;
+      case SPAN_BEGIN: {
         const node = addNode(event);
+        pushStack(node);
+        break;
+      }
+      case SPAN_END: {
+        popStack();
+        break;
+      }
+      case TOOL:
+        {
+          const node = addNode(event);
+
+          // In the span world, the first child will be a span of type tool
+          if (
+            event.events.length > 0 &&
+            (event.events[0].event !== SPAN_BEGIN ||
+              event.events[0].type !== TYPE_TOOL)
+          ) {
+            // Expand the children
+            pushStack(node);
+            for (const child of event.events) {
+              treeifyFn(child, addNode, pushStack, popStack);
+            }
+            popStack();
+          }
+        }
 
-        // In the span world, the first child will be a span of type tool
-        if (
-          event.events.length > 0 &&
-          (event.events[0].event !== SPAN_BEGIN ||
-            event.events[0].type !== TYPE_SUBTASK)
-        ) {
-          // Expand the children
-          pushStack(node);
-          for (const child of event.events) {
-            treeifyFn(child, addNode, pushStack, popStack);
+        break;
+      case SUBTASK:
+        {
+          const node = addNode(event);
+
+          // In the span world, the first child will be a span of type tool
+          if (
+            event.events.length > 0 &&
+            (event.events[0].event !== SPAN_BEGIN ||
+              event.events[0].type !== TYPE_SUBTASK)
+          ) {
+            // Expand the children
+            pushStack(node);
+            for (const child of event.events) {
+              treeifyFn(child, addNode, pushStack, popStack);
+            }
+            popStack();
           }
-          popStack();
         }
-      }
 
-      break;
-    default:
-      // An event
-      addNode(event);
-      break;
-  }
-};
+        break;
 
-type TreeNodeTransformer = {
-  name: string;
-  matches: (node: EventNode) => boolean;
-  process: (node: EventNode) => EventNode;
+      default:
+        // An event
+        addNode(event);
+        break;
+    }
+  };
+  return treeifyFn;
 };
 
-const treeNodeTransformers: TreeNodeTransformer[] = [
-  {
-    name: "unwrap_tools",
-    matches: (node) =>
-      node.event.event === SPAN_BEGIN && node.event.type === TYPE_TOOL,
-    process: (node) => elevateChildNode(node, TYPE_TOOL) || node,
-  },
-  {
-    name: "unwrap_subtasks",
-    matches: (node) =>
-      node.event.event === SPAN_BEGIN && node.event.type === TYPE_SUBTASK,
-    process: (node) => elevateChildNode(node, TYPE_SUBTASK) || node,
-  },
-  {
-    name: "unwrap_agent_solver",
-    matches: (node) =>
-      node.event.event === SPAN_BEGIN &&
-      node.event["type"] === TYPE_SOLVER &&
-      node.children.length === 2 &&
-      node.children[0].event.event === SPAN_BEGIN &&
-      node.children[0].event.type === TYPE_AGENT &&
-      node.children[1].event.event === STATE,
-
-    process: (node) => skipFirstChildNode(node),
-  },
-  {
-    name: "unwrap_agent_solver w/store",
-    matches: (node) =>
-      node.event.event === SPAN_BEGIN &&
-      node.event["type"] === TYPE_SOLVER &&
-      node.children.length === 3 &&
-      node.children[0].event.event === SPAN_BEGIN &&
-      node.children[0].event.type === TYPE_AGENT &&
-      node.children[1].event.event === STATE &&
-      node.children[2].event.event === STORE,
-    process: (node) => skipFirstChildNode(node),
-  },
-  {
-    name: "unwrap_handoff",
-    matches: (node) =>
-      node.event.event === SPAN_BEGIN &&
-      node.event["type"] === TYPE_HANDOFF &&
-      node.children.length === 2 &&
-      node.children[0].event.event === TOOL &&
-      node.children[1].event.event === STORE &&
-      node.children[0].children.length === 2 &&
-      node.children[0].children[0].event.event === SPAN_BEGIN &&
-      node.children[0].children[0].event.type === TYPE_AGENT,
-    process: (node) => skipThisNode(node),
-  },
-];
-
 const transformTree = (roots: EventNode[]): EventNode[] => {
-  const visitNode = (node: EventNode): EventNode => {
-    let processedNode = node;
+  // Gather the transformers that we'll use
+  const treeNodeTransformers: TreeNodeTransformer[] = transformers();
+
+  const visitNode = (node: EventNode): EventNode | EventNode[] => {
+    // Start with the original node
+    let currentNodes: EventNode[] = [node];
 
-    // Visit children (depth first)
-    processedNode.children = processedNode.children.map(visitNode);
+    // Process children of all nodes first (depth-first)
+    currentNodes = currentNodes.map((n) => {
+      n.children = n.children.flatMap(visitNode);
+      return n;
+    });
 
-    // Apply any visitors to this node
+    // Apply each transformer to all nodes that match
     for (const transformer of treeNodeTransformers) {
-      if (transformer.matches(processedNode)) {
-        processedNode = transformer.process(processedNode);
-        // Only apply the first matching transformer
-        break;
+      const nextNodes: EventNode[] = [];
+
+      // Process each current node with this transformer
+      for (const currentNode of currentNodes) {
+        if (transformer.matches(currentNode)) {
+          const result = transformer.process(currentNode);
+          if (Array.isArray(result)) {
+            nextNodes.push(...result);
+          } else {
+            nextNodes.push(result);
+          }
+        } else {
+          // Node doesn't match this transformer, keep it unchanged
+          nextNodes.push(currentNode);
+        }
       }
+
+      // Update current nodes for next transformer
+      currentNodes = nextNodes;
     }
-    return processedNode;
+
+    // Return all processed nodes
+    return currentNodes.length === 1 ? currentNodes[0] : currentNodes;
   };
 
-  return roots.map(visitNode);
+  // Process all nodes first
+  const processedRoots = roots.flatMap(visitNode);
+
+  // Call flush on any transformers that have it
+  const flushedNodes: EventNode[] = [];
+  for (const transformer of treeNodeTransformers) {
+    if (transformer.flush) {
+      const flushResults = transformer.flush();
+      if (flushResults && flushResults.length > 0) {
+        flushedNodes.push(...flushResults);
+      }
+    }
+  }
+
+  return [...processedRoots, ...flushedNodes];
+};
+
+const transformers = () => {
+  const treeNodeTransformers: TreeNodeTransformer[] = [
+    {
+      name: "unwrap_tools",
+      matches: (node) =>
+        node.event.event === SPAN_BEGIN && node.event.type === TYPE_TOOL,
+      process: (node) => elevateChildNode(node, TYPE_TOOL) || node,
+    },
+    {
+      name: "unwrap_subtasks",
+      matches: (node) =>
+        node.event.event === SPAN_BEGIN && node.event.type === TYPE_SUBTASK,
+      process: (node) => elevateChildNode(node, TYPE_SUBTASK) || node,
+    },
+    {
+      name: "unwrap_agent_solver",
+      matches: (node) =>
+        node.event.event === SPAN_BEGIN &&
+        node.event["type"] === TYPE_SOLVER &&
+        node.children.length === 2 &&
+        node.children[0].event.event === SPAN_BEGIN &&
+        node.children[0].event.type === TYPE_AGENT &&
+        node.children[1].event.event === STATE,
+
+      process: (node) => skipFirstChildNode(node),
+    },
+    {
+      name: "unwrap_agent_solver w/store",
+      matches: (node) =>
+        node.event.event === SPAN_BEGIN &&
+        node.event["type"] === TYPE_SOLVER &&
+        node.children.length === 3 &&
+        node.children[0].event.event === SPAN_BEGIN &&
+        node.children[0].event.type === TYPE_AGENT &&
+        node.children[1].event.event === STATE &&
+        node.children[2].event.event === STORE,
+      process: (node) => skipFirstChildNode(node),
+    },
+    {
+      name: "unwrap_handoff",
+      matches: (node) =>
+        node.event.event === SPAN_BEGIN &&
+        node.event["type"] === TYPE_HANDOFF &&
+        node.children.length === 2 &&
+        node.children[0].event.event === TOOL &&
+        node.children[1].event.event === STORE &&
+        node.children[0].children.length === 2 &&
+        node.children[0].children[0].event.event === SPAN_BEGIN &&
+        node.children[0].children[0].event.type === TYPE_AGENT,
+      process: (node) => skipThisNode(node),
+    },
+    {
+      name: "discard_solvers_span",
+      matches: (Node) =>
+        Node.event.event === SPAN_BEGIN && Node.event.type === TYPE_SOLVERS,
+      process: (node) => {
+        const nodes = discardNode(node);
+        return nodes;
+      },
+    },
+  ];
+  return treeNodeTransformers;
+};
+
+type TreeNodeTransformer = {
+  name: string;
+  matches: (node: EventNode) => boolean;
+  process: (node: EventNode) => EventNode | EventNode[];
+  flush?: () => EventNode[];
 };
 
 /**
@@ -275,6 +427,11 @@ const skipThisNode = (node: EventNode): EventNode => {
   return newNode;
 };
 
+const discardNode = (node: EventNode): EventNode[] => {
+  const nodes = reduceDepth(node.children, 1);
+  return nodes;
+};
+
 // Reduce the depth of the children by 1
 // This is used when we hoist a child node to the parent
 const reduceDepth = (nodes: EventNode[], depth: number = 1): EventNode[] => {
@@ -297,23 +454,71 @@ const setDepth = (nodes: EventNode[], depth: number): EventNode[] => {
   });
 };
 
+export interface TreeNodeVisitor {
+  visit: (node: EventNode, parent?: EventNode) => EventNode[];
+  flush?: () => EventNode[];
+}
+
 /**
  * Flatten the tree structure into a flat array of EventNode objects
  * Each node in the result will have its children set properly
- * @param events - The events to flatten
- * @param depth - The current depth in the tree
+ * @param eventNodes - The event nodes to flatten
+ * @param collapsed - Record indicating which nodes are collapsed
+ * @param visitors - Array of visitors to apply to each node
+ * @param parentNode - The parent node of the current nodes being processed
  * @returns An array of EventNode objects
  */
 export const flatTree = (
   eventNodes: EventNode[],
-  collapsed: Record<string, true> | null,
+  collapsed: Record<string, boolean> | null,
+  visitors?: TreeNodeVisitor[],
+  parentNode?: EventNode,
 ): EventNode[] => {
   const result: EventNode[] = [];
   for (const node of eventNodes) {
-    result.push(node);
-    if (collapsed === null || collapsed[node.id] !== true) {
-      result.push(...flatTree(node.children, collapsed));
+    if (visitors && visitors.length > 0) {
+      let pendingNodes: EventNode[] = [{ ...node }];
+
+      for (const visitor of visitors) {
+        const allResults: EventNode[] = [];
+        for (const pendingNode of pendingNodes) {
+          const visitorResult = visitor.visit(pendingNode);
+          if (parentNode) {
+            parentNode.children = visitorResult;
+          }
+          allResults.push(...visitorResult);
+        }
+        pendingNodes = allResults;
+      }
+
+      for (const pendingNode of pendingNodes) {
+        const children = flatTree(
+          pendingNode.children,
+          collapsed,
+          visitors,
+          pendingNode,
+        );
+        pendingNode.children = children;
+        result.push(pendingNode);
+        if (collapsed === null || collapsed[pendingNode.id] !== true) {
+          result.push(...children);
+        }
+      }
+
+      for (const visitor of visitors) {
+        if (visitor.flush) {
+          const finalNodes = visitor.flush();
+          result.push(...finalNodes);
+        }
+      }
+    } else {
+      result.push(node);
+      const children = flatTree(node.children, collapsed, visitors, node);
+      if (collapsed === null || collapsed[node.id] !== true) {
+        result.push(...children);
+      }
     }
   }
+
   return result;
 };

inspect_ai/_view/www/src/app/samples/transcript/transform/utils.ts

@@ -13,8 +13,11 @@ export const STATE = "state";
 export const TYPE_TOOL = "tool";
 export const TYPE_SUBTASK = "subtask";
 export const TYPE_SOLVER = "solver";
+export const TYPE_SOLVERS = "solvers";
 export const TYPE_AGENT = "agent";
 export const TYPE_HANDOFF = "handoff";
+export const TYPE_SCORERS = "scorers";
+export const TYPE_SCORER = "scorer";
 
 export const hasSpans = (events: Events): boolean => {
   return events.some((event) => event.event === SPAN_BEGIN);

inspect_ai/_view/www/src/app/samples/transcript/types.ts

@@ -26,6 +26,8 @@ export interface StateManager {
   applyChanges(changes: Changes): object;
 }
 
+export const kTranscriptCollapseScope = "transcript-collapse";
+
 export type EventType =
   | SampleInitEvent
   | SampleLimitEvent