inspect-ai 0.3.93__py3-none-any.whl → 0.3.94__py3-none-any.whl

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

@@ -1,6 +1,12 @@
 // This is a special name that signals a group of sandbox events.
 
-import { Events, StepEvent } from "../../../../@types/log";
+import {
+  Events,
+  SpanBeginEvent,
+  SpanEndEvent,
+  StepEvent,
+} from "../../../../@types/log";
+import { hasSpans } from "./utils";
 
 // It will be caught elsewhere and rendered with a pretty name
 export const kSandboxSignalName = "53787D8A-D3FC-426D-B383-9F880B70E4AA";
@@ -54,39 +60,55 @@ const processPendingEvents = (events: Events, filter: boolean): Events => {
 };
 
 const collapseSampleInit = (events: Events): Events => {
-  // See if the events have an init step
+  // Don't performance sample init logic if spans are present
+  const hasSpans = events.some((e) => {
+    return e.event === "span_begin" || e.event === "span_end";
+  });
+  if (hasSpans) {
+    return events;
+  }
+
+  // Don't synthesize a sample init step if one already exists
   const hasInitStep =
     events.findIndex((e) => {
       return e.event === "step" && e.name === "init";
     }) !== -1;
+  if (hasInitStep) {
+    return events;
+  }
 
+  // Find a sample init event
   const initEventIndex = events.findIndex((e) => {
     return e.event === "sample_init";
   });
   const initEvent = events[initEventIndex];
+  if (!initEvent) {
+    return events;
+  }
 
+  // Splice in sample init step if needed
   const fixedUp = [...events];
-  if (!hasInitStep && initEvent) {
-    fixedUp.splice(initEventIndex, 0, {
-      timestamp: initEvent.timestamp,
-      event: "step",
-      action: "begin",
-      type: null,
-      name: "sample_init",
-      pending: false,
-      working_start: 0,
-    });
-
-    fixedUp.splice(initEventIndex + 2, 0, {
-      timestamp: initEvent.timestamp,
-      event: "step",
-      action: "end",
-      type: null,
-      name: "sample_init",
-      pending: false,
-      working_start: 0,
-    });
-  }
+  fixedUp.splice(initEventIndex, 0, {
+    timestamp: initEvent.timestamp,
+    event: "step",
+    action: "begin",
+    type: null,
+    name: "sample_init",
+    pending: false,
+    working_start: 0,
+    span_id: initEvent.span_id,
+  });
+
+  fixedUp.splice(initEventIndex + 2, 0, {
+    timestamp: initEvent.timestamp,
+    event: "step",
+    action: "end",
+    type: null,
+    name: "sample_init",
+    pending: false,
+    working_start: 0,
+    span_id: initEvent.span_id,
+  });
   return fixedUp;
 };
 
@@ -94,12 +116,22 @@ const groupSandboxEvents = (events: Events): Events => {
   const result: Events = [];
   const pendingSandboxEvents: Events = [];
 
+  const useSpans = hasSpans(events);
+
   const pushPendingSandboxEvents = () => {
     const timestamp =
       pendingSandboxEvents[pendingSandboxEvents.length - 1].timestamp;
-    result.push(createStepEvent(kSandboxSignalName, timestamp, "begin"));
+    if (useSpans) {
+      result.push(createSpanBegin(kSandboxSignalName, timestamp, null));
+    } else {
+      result.push(createStepEvent(kSandboxSignalName, timestamp, "begin"));
+    }
     result.push(...pendingSandboxEvents);
-    result.push(createStepEvent(kSandboxSignalName, timestamp, "end"));
+    if (useSpans) {
+      result.push(createSpanEnd(kSandboxSignalName, timestamp));
+    } else {
+      result.push(createStepEvent(kSandboxSignalName, timestamp, "end"));
+    }
     pendingSandboxEvents.length = 0;
   };
 
@@ -139,4 +171,34 @@ const createStepEvent = (
   name,
   pending: false,
   working_start: 0,
+  span_id: null,
 });
+
+const createSpanBegin = (
+  name: string,
+  timestamp: string,
+  parent_id: string | null,
+): SpanBeginEvent => {
+  return {
+    name,
+    id: `${name}-begin`,
+    span_id: name,
+    parent_id,
+    timestamp,
+    event: "span_begin",
+    type: null,
+    pending: false,
+    working_start: 0,
+  };
+};
+
+const createSpanEnd = (name: string, timestamp: string): SpanEndEvent => {
+  return {
+    id: `${name}-end`,
+    timestamp,
+    event: "span_end",
+    pending: false,
+    working_start: 0,
+    span_id: name,
+  };
+};

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

@@ -1,14 +1,28 @@
 import { Events } from "../../../../@types/log";
 import { EventNode, EventType } from "../types";
+import {
+  ACTION_BEGIN,
+  ET_SPAN_BEGIN,
+  ET_SPAN_END,
+  ET_STEP,
+  hasSpans,
+} from "./utils";
+
+type TreeifyFunction = (
+  event: EventType,
+  addNode: (event: EventType) => EventNode,
+  pushStack: (node: EventNode) => void,
+  popStack: () => void,
+) => void;
 
-/**
- * Gathers events into a hierarchy of EventNodes.
- */
 export function treeifyEvents(events: Events, depth: number): EventNode[] {
+  const useSpans = hasSpans(events);
+  const treeFn = useSpans ? treeifyFnSpan : treeifyFnStep;
+
   const rootNodes: EventNode[] = [];
   const stack: EventNode[] = [];
 
-  const pushNode = (event: EventType): EventNode => {
+  const addNode = (event: EventType): EventNode => {
     const node = new EventNode(event, stack.length + depth);
     if (stack.length > 0) {
       const parentNode = stack[stack.length - 1];
@@ -19,21 +33,219 @@ export function treeifyEvents(events: Events, depth: number): EventNode[] {
     return node;
   };
 
+  const pushStack = (node: EventNode): void => {
+    stack.push(node);
+  };
+
+  const popStack = (): void => {
+    if (stack.length > 0) {
+      stack.pop();
+    }
+  };
+
   events.forEach((event) => {
-    if (event.event === "step" && event.action === "begin") {
-      // Starting a new step
-      const node = pushNode(event);
-      stack.push(node);
-    } else if (event.event === "step" && event.action === "end") {
-      // An ending step
-      if (stack.length > 0) {
-        stack.pop();
+    treeFn(event, addNode, pushStack, popStack);
+  });
+
+  if (useSpans) {
+    return transformTree(rootNodes);
+  } else {
+    return rootNodes;
+  }
+}
+
+const treeifyFnStep: TreeifyFunction = (
+  event: EventType,
+  addNode: (event: EventType) => EventNode,
+  pushStack: (node: EventNode) => void,
+  popStack: () => void,
+): void => {
+  switch (event.event) {
+    case ET_STEP:
+      if (event.action === ACTION_BEGIN) {
+        // Starting a new step
+        const node = addNode(event);
+        pushStack(node);
+      } else {
+        // An ending step
+        popStack();
       }
-    } else {
+      break;
+    case ET_SPAN_BEGIN: {
+      // These shoudn't be here, but throw away
+      break;
+    }
+    case ET_SPAN_END: {
+      // These shoudn't be here, but throw away
+      break;
+    }
+    default:
       // An event
-      pushNode(event);
+      addNode(event);
+      break;
+  }
+};
+
+const treeifyFnSpan: TreeifyFunction = (
+  event: EventType,
+  addNode: (event: EventType) => EventNode,
+  pushStack: (node: EventNode) => void,
+  popStack: () => void,
+): void => {
+  switch (event.event) {
+    case ET_STEP:
+      // strip steps
+      break;
+    case ET_SPAN_BEGIN: {
+      const node = addNode(event);
+      pushStack(node);
+      break;
     }
-  });
+    case ET_SPAN_END: {
+      popStack();
+      break;
+    }
+    default:
+      // An event
+      addNode(event);
+      break;
+  }
+};
 
-  return rootNodes;
-}
+type TreeNodeTransformer = {
+  name: string;
+  matches: (node: EventNode) => boolean;
+  process: (node: EventNode) => EventNode;
+};
+
+const treeNodeTransformers: TreeNodeTransformer[] = [
+  {
+    name: "unwrap_tools",
+    matches: (node) =>
+      node.event.event === "span_begin" && node.event.type === "tool",
+    process: (node) => elevateChildNode(node, "tool") || node,
+  },
+  {
+    name: "unwrap_subtasks",
+    matches: (node) =>
+      node.event.event === "span_begin" && node.event.type === "subtask",
+    process: (node) => elevateChildNode(node, "subtask") || node,
+  },
+  {
+    name: "unwrap_agent_solver",
+    matches: (node) =>
+      node.event.event === "span_begin" &&
+      node.event["type"] === "solver" &&
+      node.children.length === 2 &&
+      node.children[0].event.event === "span_begin" &&
+      node.children[0].event.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"] === "solver" &&
+      node.children.length === 3 &&
+      node.children[0].event.event === "span_begin" &&
+      node.children[0].event.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"] === "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 === "agent",
+    process: (node) => skipThisNode(node),
+  },
+];
+
+const transformTree = (roots: EventNode[]): EventNode[] => {
+  const visitNode = (node: EventNode): EventNode => {
+    let processedNode = node;
+
+    // Visit children (depth first)
+    processedNode.children = processedNode.children.map(visitNode);
+
+    // Apply any visitors to this node
+    for (const transformer of treeNodeTransformers) {
+      if (transformer.matches(processedNode)) {
+        processedNode = transformer.process(processedNode);
+        // Only apply the first matching transformer
+        break;
+      }
+    }
+    return processedNode;
+  };
+
+  return roots.map(visitNode);
+};
+
+/**
+ * Process a span node by elevating a specific child node type and moving its siblings as children
+ * @template T - Type of the event (either ToolEvent or SubtaskEvent)
+ */
+const elevateChildNode = (
+  node: EventNode,
+  childEventType: "tool" | "subtask",
+): EventNode | null => {
+  // Find the specific event child
+  const targetIndex = node.children.findIndex(
+    (child) => child.event.event === childEventType,
+  );
+
+  if (targetIndex === -1) {
+    console.log(
+      `No ${childEventType} event found in a span, this is very unexpected.`,
+    );
+    return null;
+  }
+
+  // Get the target node and set its depth
+  const targetNode = { ...node.children[targetIndex] };
+  const remainingChildren = node.children.filter((_, i) => i !== targetIndex);
+
+  // Process the remaining children
+  targetNode.depth = node.depth;
+  targetNode.children = reduceDepth(remainingChildren);
+
+  // No need to update the event itself (events have been deprecated
+  // and more importantly we drive children / transcripts using the tree structure itself
+  // and notes rather than the event.events itself)
+  return targetNode;
+};
+
+const skipFirstChildNode = (node: EventNode): EventNode => {
+  const agentSpan = node.children.splice(0, 1)[0];
+  node.children.unshift(...reduceDepth(agentSpan.children));
+  return node;
+};
+
+const skipThisNode = (node: EventNode): EventNode => {
+  const newNode = { ...node.children[0] };
+  newNode.depth = node.depth;
+  newNode.children = reduceDepth(newNode.children[0].children, 2);
+  return newNode;
+};
+
+// 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[] => {
+  return nodes.map((node) => {
+    if (node.children.length > 0) {
+      node.children = reduceDepth(node.children, 1);
+    }
+    node.depth = node.depth - depth;
+    return node;
+  });
+};

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

@@ -0,0 +1,11 @@
+import { Events } from "../../../../@types/log";
+
+export const ET_STEP = "step";
+export const ACTION_BEGIN = "begin";
+
+export const ET_SPAN_BEGIN = "span_begin";
+export const ET_SPAN_END = "span_end";
+
+export const hasSpans = (events: Events): boolean => {
+  return events.some((event) => event.event === ET_SPAN_BEGIN);
+};

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

@@ -10,6 +10,8 @@ import {
   SampleLimitEvent,
   SandboxEvent,
   ScoreEvent,
+  SpanBeginEvent,
+  SpanEndEvent,
   StateEvent,
   StepEvent,
   StoreEvent,
@@ -39,7 +41,9 @@ export type EventType =
   | InputEvent
   | ErrorEvent
   | ApprovalEvent
-  | SandboxEvent;
+  | SandboxEvent
+  | SpanBeginEvent
+  | SpanEndEvent;
 
 export class EventNode {
   event: EventType;

inspect_ai/agent/_as_solver.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any
 
 from inspect_ai.util._limit import Limit, apply_limits
+from inspect_ai.util._span import span
 
 if TYPE_CHECKING:
     from inspect_ai.solver._solver import Solver
@@ -61,7 +62,8 @@ def as_solver(agent: Agent, limits: list[Limit] = [], **agent_kwargs: Any) -> So
             try:
                 # run the agent with limits
                 with apply_limits(limits):
-                    agent_state = await agent(agent_state, **agent_kwargs)
+                    async with span(name=agent_name, type="agent"):
+                        agent_state = await agent(agent_state, **agent_kwargs)
             # if an exception occurs, we still want to update the TaskState with the
             # AgentState's messages + output so that it appears in the log and is scored
             finally:

inspect_ai/agent/_as_tool.py

@@ -11,6 +11,7 @@ from inspect_ai.tool._tool_def import ToolDef, validate_tool_parameters
 from inspect_ai.tool._tool_info import ToolInfo, parse_tool_info
 from inspect_ai.tool._tool_params import ToolParam
 from inspect_ai.util._limit import Limit, apply_limits
+from inspect_ai.util._span import span
 
 from ._agent import AGENT_DESCRIPTION, Agent, AgentState
 
@@ -49,13 +50,17 @@ def as_tool(
             "Agent passed to as_tool was not created by an @agent decorated function"
         )
 
+    # get tool_info
+    tool_info = agent_tool_info(agent, description, **agent_kwargs)
+
     async def execute(input: str, *args: Any, **kwargs: Any) -> ToolResult:
         # prepare state
         state = AgentState(messages=[ChatMessageUser(content=input, source="input")])
 
         # run the agent with limits
         with apply_limits(limits):
-            state = await agent(state, *args, **(agent_kwargs | kwargs))
+            async with span(name=tool_info.name, type="agent"):
+                state = await agent(state, *args, **(agent_kwargs | kwargs))
 
         # find assistant message to read content from (prefer output)
         if not state.output.empty:
@@ -67,9 +72,6 @@ def as_tool(
         else:
             return ""
 
-    # get tool_info
-    tool_info = agent_tool_info(agent, description, **agent_kwargs)
-
     # add "input" param
     tool_info.parameters.properties = {
         "input": ToolParam(type="string", description="Input message.")

inspect_ai/agent/_handoff.py

@@ -57,7 +57,9 @@ def handoff(
     tool_info = agent_tool_info(agent, description, **agent_kwargs)
 
     # AgentTool calls will be intercepted by execute_tools
-    agent_tool = AgentTool(agent, input_filter, output_filter, limits, **agent_kwargs)
+    agent_tool = AgentTool(
+        agent, tool_info.name, input_filter, output_filter, limits, **agent_kwargs
+    )
     tool_name = tool_name or f"transfer_to_{tool_info.name}"
     set_registry_info(agent_tool, RegistryInfo(type="tool", name=tool_name))
     set_tool_description(
@@ -75,12 +77,14 @@ class AgentTool(Tool):
     def __init__(
         self,
         agent: Agent,
+        name: str,
         input_filter: MessageFilter | None = None,
         output_filter: MessageFilter | None = None,
         limits: list[Limit] = [],
         **kwargs: Any,
     ):
         self.agent = agent
+        self.name = name
         self.input_filter = input_filter
         self.output_filter = output_filter
         self.limits = limits

inspect_ai/agent/_react.py

@@ -195,9 +195,10 @@ def react(
                     answer = submission(messages)
                     if answer is not None:
                         # set the output to the answer for scoring
-                        state.output.completion = (
-                            f"{state.output.completion}\n\n{answer}".strip()
-                        )
+                        if submit.answer_only:
+                            state.output.completion = answer
+                        else:
+                            state.output.completion = f"{state.output.completion}{submit.answer_delimiter}{answer}".strip()
 
                         # exit if we are at max_attempts
                         attempt_count += 1

inspect_ai/agent/_run.py

@@ -1,8 +1,10 @@
 from copy import copy
 from typing import Any
 
+from inspect_ai._util.registry import registry_unqualified_name
 from inspect_ai.model._chat_message import ChatMessage, ChatMessageUser
 from inspect_ai.util._limit import Limit, apply_limits
+from inspect_ai.util._span import span
 
 from ._agent import Agent, AgentState
 
@@ -52,4 +54,7 @@ async def run(
 
     # run the agent with limits
     with apply_limits(limits):
-        return await agent(state, **agent_kwargs)
+        # run the agent
+        agent_name = registry_unqualified_name(agent)
+        async with span(name=agent_name, type="agent"):
+            return await agent(state, **agent_kwargs)

inspect_ai/agent/_types.py

@@ -96,3 +96,12 @@ class AgentSubmit(NamedTuple):
 
     The tool should return the `answer` provided to it for scoring.
     """
+
+    answer_only: bool = False
+    """Set the completion to only the answer provided by the submit tool.
+
+    By default, the answer is appended (with `answer_delimiter`) to whatever
+    other content the model generated along with the call to `submit()`."""
+
+    answer_delimiter: str = "\n\n"
+    """Delimter used when appending submit tool answer to other content the model generated along with the call to `submit()`."""

inspect_ai/dataset/_dataset.py

@@ -16,6 +16,7 @@ from typing import (
 from pydantic import BaseModel, Field, ValidationError
 from typing_extensions import override
 
+from inspect_ai._util.answer import answer_character, answer_index
 from inspect_ai.model import ChatMessage
 from inspect_ai.util import SandboxEnvironmentSpec, SandboxEnvironmentType
 from inspect_ai.util._sandbox.environment import resolve_sandbox_environment
@@ -328,7 +329,9 @@ class MemoryDataset(Dataset):
             shuffled_choices = [sample.choices[i] for i in positions]
 
             # Map of original position / target letter
-            position_map = {i: chr(65 + new_i) for new_i, i in enumerate(positions)}
+            position_map = {
+                i: answer_character(new_i) for new_i, i in enumerate(positions)
+            }
 
             # Update to the shuffled choices and target
             sample.choices = shuffled_choices
@@ -338,9 +341,9 @@ class MemoryDataset(Dataset):
         self, target: str | list[str], position_map: dict[int, str]
     ) -> str | list[str]:
         if isinstance(target, list):
-            return [position_map[ord(t) - 65] for t in target]
+            return [position_map[answer_index(t)] for t in target]
         else:
-            return position_map[ord(target) - 65]
+            return position_map[answer_index(target)]
 
     @override
     def sort(

inspect_ai/log/__init__.py

@@ -48,6 +48,8 @@ from ._transcript import (
     SampleLimitEvent,
     SandboxEvent,
     ScoreEvent,
+    SpanBeginEvent,
+    SpanEndEvent,
     StateEvent,
     StepEvent,
     StoreEvent,
@@ -56,6 +58,7 @@ from ._transcript import (
     Transcript,
     transcript,
 )
+from ._tree import EventNode, EventTree, SpanNode, event_sequence, event_tree
 
 __all__ = [
     "EvalConfig",
@@ -92,6 +95,8 @@ __all__ = [
     "SampleLimitEvent",
     "SandboxEvent",
     "ScoreEvent",
+    "SpanBeginEvent",
+    "SpanEndEvent",
     "StateEvent",
     "StepEvent",
     "StoreEvent",
@@ -111,4 +116,9 @@ __all__ = [
     "write_log_dir_manifest",
     "retryable_eval_logs",
     "bundle_log_dir",
+    "event_tree",
+    "event_sequence",
+    "EventTree",
+    "EventNode",
+    "SpanNode",
 ]

inspect_ai/log/_convert.py

@@ -2,7 +2,7 @@ import os
 from typing import Literal
 
 from inspect_ai._util.error import PrerequisiteError
-from inspect_ai._util.file import copy_file, exists, filesystem
+from inspect_ai._util.file import exists, filesystem
 from inspect_ai.log._file import (
     log_files_from_ls,
     read_eval_log,
@@ -66,14 +66,9 @@ def convert_eval_logs(
                 "Output file {output_file} already exists (use --overwrite to overwrite existing files)"
             )
 
-        # if the input and output files have the same format just copy
-        if input_file.endswith(f".{to}"):
-            copy_file(input_file, output_file)
-
-        # otherwise do a full read/write
-        else:
-            log = read_eval_log(input_file)
-            write_eval_log(log, output_file)
+        # do a full read/write (normalized deprecated constructs and adds sample summaries)
+        log = read_eval_log(input_file)
+        write_eval_log(log, output_file)
 
     if fs.info(path).type == "file":
         convert_file(path)