From 7333fb79d12240fd8786450bde4a2afa6b853d36 Mon Sep 17 00:00:00 2001 From: Charles Teague Date: Mon, 13 Jul 2026 15:46:43 -0400 Subject: [PATCH 1/4] timeline: require a tool-calling loop before classifying utility agents MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The utility heuristics in timeline_build (foreign-prompt event wrapping and single-turn span classification) assumed a "primary trajectory" defined by the first system prompt seen. In plain workflows of generate() calls — e.g. a monitor making several checker calls with different prompts and no tools — this hid every call after the first prompt group as a utility agent. Both passes now only apply when the span actually contains an agentic (tool-calling) loop: _wrap_utility_events derives the primary prompt solely from the first tool-calling ModelEvent (no fallback to the first event), and _classify_utility_agents only demotes single-turn sub-spans whose parent runs a tool loop. Warmup-call wrapping is unchanged. Adds workflow_generate (mixed-prompt calls, no tools → nothing hidden) and bridge_utility (foreign-prompt call within a tool loop → still utility) scenarios. Mirrors the same fix in ts-mono's core.ts. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 1 + src/inspect_ai/event/_timeline.py | 68 ++++++++------ tests/timeline/generate.py | 149 +++++++++++++++++++++++++++++- tests/timeline/test_timeline.py | 12 +++ 4 files changed, 201 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 217541ab15..831fd17223 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,7 @@ - Bugfix: Errored samples that were scored (e.g. via `score_on_error`) now contribute their scores to metrics and `scored_samples`, matching what the sample log shows. (#4412) - Bugfix: Several raised errors (unexpected Anthropic input block, unsupported image data type, missing sample-buffer database) now interpolate the offending value into the message instead of showing the literal placeholder text. - Bugfix: React agent compaction now works after a checkpoint resume — the restored conversation is no longer treated as an always-preserved prefix, so compaction shrinks the context again. +- Bugfix: Transcript timelines no longer hide workflow `generate()` calls with differing system prompts as utility agents; utility classification now requires a tool-calling agent loop. - Viewer: log listing responses include the log dir's canonical URI (`log_dir_uri`) so the viewer can reliably scope its local cache to the directory. - Inspect View: Reuse one warm async S3 client and connection pool across requests — for both log reads and directory listings — instead of creating one per operation, eliminating the per-request credential/connection cold-start (e.g. `/log-headers` ~3s -> ~0.3s, `/logs` ~1.5s -> ~0.06s). diff --git a/src/inspect_ai/event/_timeline.py b/src/inspect_ai/event/_timeline.py index 071fd68181..5cac7d4970 100644 --- a/src/inspect_ai/event/_timeline.py +++ b/src/inspect_ai/event/_timeline.py @@ -1157,40 +1157,27 @@ def _wrap_utility_events(agent: TimelineSpan) -> None: This function detects them and wraps each one in a ``TimelineSpan`` with ``utility=True`` so downstream code treats them as utility agents. + The primary trajectory is identified by the system prompt of the first + tool-calling ModelEvent. When a span contains no tool-calling model + events there is no agentic loop to distinguish helpers from — plain + workflows of ``generate()`` calls routinely mix system prompts — so + foreign-prompt wrapping is skipped entirely (warmup calls are still + wrapped; they are identified independently of prompts). + Operates recursively on the entire span tree. Args: agent: The span node to process (mutated in place). """ - # --- Determine the primary system prompt for this span --- + # The primary prompt comes only from a tool-calling ModelEvent: without + # an agentic loop there is no primary trajectory (see docstring). primary_prompt: str | None = None - - # Prefer the prompt of the first ModelEvent that has tool calls for item in agent.content: if isinstance(item, TimelineEvent) and isinstance(item.event, ModelEvent): if _has_tool_calls(item.event): primary_prompt = _get_system_prompt_for_event(item.event) break - # Fall back to the first ModelEvent's prompt - if primary_prompt is None: - for item in agent.content: - if isinstance(item, TimelineEvent) and isinstance(item.event, ModelEvent): - primary_prompt = _get_system_prompt_for_event(item.event) - break - - # No ModelEvents at all → nothing to wrap - if primary_prompt is None: - # Still recurse into child spans - for item in agent.content: - if isinstance(item, TimelineSpan): - _wrap_utility_events(item) - for branch in agent.branches: - for item in branch.content: - if isinstance(item, TimelineSpan): - _wrap_utility_events(item) - return - # --- Scan and wrap utility candidates --- new_content: list[TimelineEvent | TimelineSpan] = [] for item in agent.content: @@ -1209,7 +1196,8 @@ def _wrap_utility_events(agent: TimelineSpan) -> None: evt_prompt = _get_system_prompt_for_event(item.event) if ( - evt_prompt is not None + primary_prompt is not None + and evt_prompt is not None and evt_prompt != primary_prompt and not _has_tool_calls(item.event) ): @@ -1316,17 +1304,34 @@ def _is_single_turn(agent: TimelineSpan) -> bool: return False +def _has_agentic_loop(span: TimelineSpan) -> bool: + """Check whether span's direct content contains a tool-calling ModelEvent.""" + return any( + isinstance(item, TimelineEvent) + and isinstance(item.event, ModelEvent) + and _has_tool_calls(item.event) + for item in span.content + ) + + def _classify_utility_agents( - node: TimelineSpan, parent_system_prompt: str | None = None + node: TimelineSpan, + parent_system_prompt: str | None = None, + parent_has_loop: bool = False, ) -> None: """Classify utility agents in the tree via post-processing. An agent is utility if it has a single turn (or single tool-calling turn) - and a different system prompt than its parent. + and a different system prompt than its parent. Classification only applies + when the parent runs an agentic (tool-calling) loop: absent a loop there + is no main trajectory for a helper to be subordinate to — plain workflows + of ``generate()`` calls routinely mix system prompts. Args: node: The span node to classify (and recurse into). parent_system_prompt: The system prompt of the parent agent. + parent_has_loop: Whether the parent has a tool-calling ModelEvent + in its direct content. """ agent_system_prompt = _get_system_prompt(node) @@ -1336,17 +1341,26 @@ def _classify_utility_agents( # helper model calls are handled separately by _wrap_utility_events. if ( parent_system_prompt is not None + and parent_has_loop and agent_system_prompt is not None and not node.tool_invoked ): if agent_system_prompt != parent_system_prompt and _is_single_turn(node): node.utility = True - # Recurse into child spans + # Recurse into child spans. Pure grouping spans (no direct model events) + # inherit from the parent, mirroring the prompt inheritance. effective_prompt = agent_system_prompt or parent_system_prompt + has_model_events = any( + isinstance(item, TimelineEvent) and isinstance(item.event, ModelEvent) + for item in node.content + ) + effective_has_loop = ( + _has_agentic_loop(node) if has_model_events else parent_has_loop + ) for item in node.content: if isinstance(item, TimelineSpan): - _classify_utility_agents(item, effective_prompt) + _classify_utility_agents(item, effective_prompt, effective_has_loop) # ============================================================================= diff --git a/tests/timeline/generate.py b/tests/timeline/generate.py index 2648364c06..2d564a5b83 100644 --- a/tests/timeline/generate.py +++ b/tests/timeline/generate.py @@ -12,6 +12,8 @@ 6. Sequential run – multiple run() calls with inference in between 7. Parallel collect – three parallel sub-agents via collect() 8. Handoff + as_tool – combining handoff() and as_tool() sub-agents + 9. Workflow generate – mixed-prompt generate() calls, no tools → no utility + 10. Bridge utility – foreign-prompt call within a tool loop → utility """ from __future__ import annotations @@ -30,9 +32,21 @@ from inspect_ai.event import Timeline, timeline_build from inspect_ai.event._timeline import TimelineSpan from inspect_ai.log import EvalLog -from inspect_ai.model import ModelOutput, get_model +from inspect_ai.model import ( + ChatMessageSystem, + ChatMessageUser, + ModelOutput, + get_model, +) from inspect_ai.scorer import includes -from inspect_ai.solver import generate, system_message +from inspect_ai.solver import ( + Generate, + Solver, + TaskState, + generate, + solver, + system_message, +) from inspect_ai.tool import tool from inspect_ai.util import collect @@ -601,6 +615,89 @@ def scenario_deep_utility() -> tuple[str, Task, Any]: return "deep_utility", task, model +def scenario_workflow_generate() -> tuple[str, Task, Any]: + """Workflow of generate() calls with mixed system prompts and no tools. + + A monitor-style solver makes five auxiliary model calls — two sharing + one system prompt, three using different prompts. With no tool-calling + loop anywhere there is no primary trajectory, so none of the calls may + be hidden as utility agents. + """ + + @solver + def monitor() -> Solver: + async def solve(state: TaskState, generate: Generate) -> TaskState: + model = get_model(MODEL) + calls = [ + ("x-extract-1", "You are the x-extract checker."), + ("x-extract-2", "You are the x-extract checker."), + ("y-extract-1", "You are the y-extract checker."), + ("y-ground-1", "You are the y-ground checker."), + ("y-ground-2", "You are the y-ground checker."), + ] + for call_name, prompt in calls: + await model.generate( + [ + ChatMessageSystem(content=prompt), + ChatMessageUser(content=f"{call_name}: {state.input_text}"), + ] + ) + return state + + return solve + + model = get_model(MODEL) + task = Task(name="workflow_generate", dataset=DATASET, solver=monitor()) + return "workflow_generate", task, model + + +def scenario_bridge_utility() -> tuple[str, Task, Any]: + """Foreign-prompt helper call within a tool-calling loop → utility. + + A bridge-style agent emits a tool-calling main loop plus one extraction + call with a different system prompt and no tool calls. The extraction + call is wrapped as a utility span; main-loop calls stay visible. + """ + main_prompt = "You are the main assistant." + + @agent + def bridge() -> Agent: + async def execute(state: AgentState) -> AgentState: + model = get_model( + MODEL, + custom_outputs=[ + ModelOutput.for_tool_call(MODEL, "search", {"query": "1 + 1"}), + ModelOutput.from_content(MODEL, "two"), + ModelOutput.from_content(MODEL, "2"), + ], + ) + await model.generate( + [ + ChatMessageSystem(content=main_prompt), + ChatMessageUser(content="What is 1 + 1?"), + ] + ) + await model.generate( + [ + ChatMessageSystem(content="Extract a short title."), + ChatMessageUser(content="Extract."), + ] + ) + state.output = await model.generate( + [ + ChatMessageSystem(content=main_prompt), + ChatMessageUser(content="Answer."), + ] + ) + return state + + return execute + + model = get_model(MODEL) + task = Task(name="bridge_utility", dataset=DATASET, solver=bridge()) + return "bridge_utility", task, model + + def scenario_parallel_heterogeneous() -> tuple[str, Task, Any]: """Parallel agents with different names. @@ -968,6 +1065,50 @@ def validate_deep_utility(timeline: Timeline) -> None: assert_repr_labels(timeline, "main", "dispatcher", "lookup") +def _utility_spans(span: TimelineSpan) -> list[TimelineSpan]: + """Collect all descendant spans classified as utility.""" + results: list[TimelineSpan] = [] + for item in span.content: + if isinstance(item, TimelineSpan): + if item.utility: + results.append(item) + results.extend(_utility_spans(item)) + return results + + +def _direct_model_events(span: TimelineSpan) -> list[Any]: + return [ + item.event + for item in span.content + if not isinstance(item, TimelineSpan) and item.event.event == "model" + ] + + +def validate_workflow_generate(timeline: Timeline) -> None: + root = timeline.root + assert not _utility_spans(root), ( + "Workflow generate() calls must not be classified as utility agents" + ) + model_events = _direct_model_events(root) + assert len(model_events) == 5, ( + f"Expected all 5 workflow calls direct in main, got {len(model_events)}" + ) + + +def validate_bridge_utility(timeline: Timeline) -> None: + root = timeline.root + utility = _utility_spans(root) + assert len(utility) == 1, f"Expected exactly 1 utility span, got {len(utility)}" + wrapped = _direct_model_events(utility[0]) + assert len(wrapped) == 1, "Expected the utility span to wrap a single call" + system = next(m.content for m in wrapped[0].input if m.role == "system") + assert "Extract" in str(system), "Expected the extraction call to be wrapped" + main_events = _direct_model_events(root) + assert len(main_events) == 2, ( + f"Expected 2 main-loop calls direct in main, got {len(main_events)}" + ) + + def validate_parallel_heterogeneous(timeline: Timeline) -> None: root = timeline.root search_spans = find_spans(root, "search") @@ -997,6 +1138,8 @@ def validate_parallel_heterogeneous(timeline: Timeline) -> None: "parallel_with_nesting": validate_parallel_with_nesting, "sequential_and_parallel": validate_sequential_and_parallel, "deep_utility": validate_deep_utility, + "workflow_generate": validate_workflow_generate, + "bridge_utility": validate_bridge_utility, "parallel_heterogeneous": validate_parallel_heterogeneous, } @@ -1017,6 +1160,8 @@ def validate_parallel_heterogeneous(timeline: Timeline) -> None: scenario_parallel_with_nesting, scenario_sequential_and_parallel, scenario_deep_utility, + scenario_workflow_generate, + scenario_bridge_utility, scenario_parallel_heterogeneous, ] diff --git a/tests/timeline/test_timeline.py b/tests/timeline/test_timeline.py index ee528d48b3..bdf71b732a 100644 --- a/tests/timeline/test_timeline.py +++ b/tests/timeline/test_timeline.py @@ -9,6 +9,7 @@ from inspect_ai.event import Timeline, timeline_build from .generate import ( + scenario_bridge_utility, scenario_deep_nesting, scenario_deep_utility, scenario_handoff_and_as_tool, @@ -22,6 +23,8 @@ scenario_sequential_run, scenario_simple_agent, scenario_utility_agent, + scenario_workflow_generate, + validate_bridge_utility, validate_deep_nesting, validate_deep_utility, validate_handoff_and_as_tool, @@ -35,6 +38,7 @@ validate_sequential_run, validate_simple_agent, validate_utility_agent, + validate_workflow_generate, ) @@ -101,5 +105,13 @@ def test_timeline_deep_utility() -> None: _run_and_validate(scenario_deep_utility, validate_deep_utility) +def test_timeline_workflow_generate() -> None: + _run_and_validate(scenario_workflow_generate, validate_workflow_generate) + + +def test_timeline_bridge_utility() -> None: + _run_and_validate(scenario_bridge_utility, validate_bridge_utility) + + def test_timeline_parallel_heterogeneous() -> None: _run_and_validate(scenario_parallel_heterogeneous, validate_parallel_heterogeneous) From b92fe2f2f1507536537fe0e923feb4da20080d99 Mon Sep 17 00:00:00 2001 From: Charles Teague Date: Wed, 15 Jul 2026 14:22:41 -0400 Subject: [PATCH 2/4] timeline: define utility-classification tests as shared JSON fixtures Move the utility-classification test cases into JSON fixtures under tests/timeline/fixtures/events/ ({description, events, expected}) so a single definition drives both timeline implementations: the new test_timeline_fixtures.py validates timeline_build() against them, and ts-mono's fixtureTimeline.test.ts discovers the same directory to validate buildTimeline(). The scenario-based duplicates (workflow_generate, bridge_utility) are removed from generate.py/test_timeline.py, and two additional cases cover single-turn sub-agent classification with and without a parent tool-calling loop. Co-Authored-By: Claude Fable 5 --- .../events/utility_bridge_wrapped.json | 195 ++++++++++++++ .../events/utility_subagent_loop.json | 210 +++++++++++++++ .../events/utility_subagent_no_loop.json | 147 +++++++++++ .../events/utility_workflow_generate.json | 243 ++++++++++++++++++ tests/timeline/generate.py | 149 +---------- tests/timeline/test_timeline.py | 12 - tests/timeline/test_timeline_fixtures.py | 155 +++++++++++ 7 files changed, 952 insertions(+), 159 deletions(-) create mode 100644 tests/timeline/fixtures/events/utility_bridge_wrapped.json create mode 100644 tests/timeline/fixtures/events/utility_subagent_loop.json create mode 100644 tests/timeline/fixtures/events/utility_subagent_no_loop.json create mode 100644 tests/timeline/fixtures/events/utility_workflow_generate.json create mode 100644 tests/timeline/test_timeline_fixtures.py diff --git a/tests/timeline/fixtures/events/utility_bridge_wrapped.json b/tests/timeline/fixtures/events/utility_bridge_wrapped.json new file mode 100644 index 0000000000..bd25b9a194 --- /dev/null +++ b/tests/timeline/fixtures/events/utility_bridge_wrapped.json @@ -0,0 +1,195 @@ +{ + "description": "Bridge-style agent: a tool-calling main loop plus one extraction call with a different system prompt and no tool calls. The extraction call is wrapped as a synthetic utility span; main-loop calls stay direct content.", + "events": [ + { + "event": "span_begin", + "id": "solvers", + "name": "solvers", + "type": "solvers", + "parent_id": null, + "timestamp": "2026-01-01T00:00:01Z", + "uuid": "sb-solvers" + }, + { + "event": "span_begin", + "id": "bridge", + "name": "bridge", + "type": "agent", + "parent_id": "solvers", + "timestamp": "2026-01-01T00:00:02Z", + "uuid": "sb-bridge" + }, + { + "event": "model", + "uuid": "m1", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:03Z", + "completed": "2026-01-01T00:00:04Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the main assistant." + }, + { + "role": "user", + "content": "What is 1 + 1?" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "function": "search", + "arguments": { + "query": "1 + 1" + } + } + ] + }, + "stop_reason": "tool_calls" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "tool", + "uuid": "t1", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:05Z", + "completed": "2026-01-01T00:00:06Z", + "id": "call_1", + "function": "search", + "arguments": { + "query": "1 + 1" + }, + "result": "found it", + "events": [] + }, + { + "event": "model", + "uuid": "m2", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:07Z", + "completed": "2026-01-01T00:00:08Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "Extract a short title." + }, + { + "role": "user", + "content": "Extract." + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "two" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "model", + "uuid": "m3", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:09Z", + "completed": "2026-01-01T00:00:10Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the main assistant." + }, + { + "role": "user", + "content": "Answer." + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "2" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "span_end", + "id": "bridge", + "timestamp": "2026-01-01T00:00:11Z", + "uuid": "se-bridge" + }, + { + "event": "span_end", + "id": "solvers", + "timestamp": "2026-01-01T00:00:12Z", + "uuid": "se-solvers" + } + ], + "expected": { + "init": null, + "agent": { + "id": "bridge", + "name": "main", + "event_uuids": [ + "m1", + "t1", + "m3" + ], + "children": [ + { + "id": "utility-m2", + "name": "utility", + "utility": true, + "event_uuids": [ + "m2" + ] + } + ] + }, + "scoring": null + } +} diff --git a/tests/timeline/fixtures/events/utility_subagent_loop.json b/tests/timeline/fixtures/events/utility_subagent_loop.json new file mode 100644 index 0000000000..a0976585dd --- /dev/null +++ b/tests/timeline/fixtures/events/utility_subagent_loop.json @@ -0,0 +1,210 @@ +{ + "description": "Single-turn sub-agent with a different system prompt under a parent that DOES run a tool-calling loop. This is the internal-helper shape the utility heuristic targets, so the sub-agent is classified as utility.", + "events": [ + { + "event": "span_begin", + "id": "solvers", + "name": "solvers", + "type": "solvers", + "parent_id": null, + "timestamp": "2026-01-01T00:00:01Z", + "uuid": "sb-solvers" + }, + { + "event": "span_begin", + "id": "bridge", + "name": "bridge", + "type": "agent", + "parent_id": "solvers", + "timestamp": "2026-01-01T00:00:02Z", + "uuid": "sb-bridge" + }, + { + "event": "model", + "uuid": "m1", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:03Z", + "completed": "2026-01-01T00:00:04Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the main assistant." + }, + { + "role": "user", + "content": "What is 1 + 1?" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "function": "search", + "arguments": { + "query": "1 + 1" + } + } + ] + }, + "stop_reason": "tool_calls" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "tool", + "uuid": "t1", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:05Z", + "completed": "2026-01-01T00:00:06Z", + "id": "call_1", + "function": "search", + "arguments": { + "query": "1 + 1" + }, + "result": "found it", + "events": [] + }, + { + "event": "span_begin", + "id": "helper", + "name": "helper", + "type": "agent", + "parent_id": "bridge", + "timestamp": "2026-01-01T00:00:07Z", + "uuid": "sb-helper" + }, + { + "event": "model", + "uuid": "m2", + "span_id": "helper", + "timestamp": "2026-01-01T00:00:08Z", + "completed": "2026-01-01T00:00:09Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are a helper." + }, + { + "role": "user", + "content": "help" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "ok" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "span_end", + "id": "helper", + "timestamp": "2026-01-01T00:00:10Z", + "uuid": "se-helper" + }, + { + "event": "model", + "uuid": "m3", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:11Z", + "completed": "2026-01-01T00:00:12Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the main assistant." + }, + { + "role": "user", + "content": "Answer." + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "2" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "span_end", + "id": "bridge", + "timestamp": "2026-01-01T00:00:13Z", + "uuid": "se-bridge" + }, + { + "event": "span_end", + "id": "solvers", + "timestamp": "2026-01-01T00:00:14Z", + "uuid": "se-solvers" + } + ], + "expected": { + "init": null, + "agent": { + "id": "bridge", + "name": "main", + "event_uuids": [ + "m1", + "t1", + "m3" + ], + "children": [ + { + "id": "helper", + "name": "helper", + "utility": true, + "event_uuids": [ + "m2" + ] + } + ] + }, + "scoring": null + } +} diff --git a/tests/timeline/fixtures/events/utility_subagent_no_loop.json b/tests/timeline/fixtures/events/utility_subagent_no_loop.json new file mode 100644 index 0000000000..d2f01fe2b7 --- /dev/null +++ b/tests/timeline/fixtures/events/utility_subagent_no_loop.json @@ -0,0 +1,147 @@ +{ + "description": "Single-turn sub-agent with a different system prompt under a parent that has NO tool-calling loop. Without an agentic loop there is no main trajectory, so the sub-agent must not be classified as utility.", + "events": [ + { + "event": "span_begin", + "id": "solvers", + "name": "solvers", + "type": "solvers", + "parent_id": null, + "timestamp": "2026-01-01T00:00:01Z", + "uuid": "sb-solvers" + }, + { + "event": "span_begin", + "id": "monitor", + "name": "monitor", + "type": "agent", + "parent_id": "solvers", + "timestamp": "2026-01-01T00:00:02Z", + "uuid": "sb-monitor" + }, + { + "event": "model", + "uuid": "m1", + "span_id": "monitor", + "timestamp": "2026-01-01T00:00:03Z", + "completed": "2026-01-01T00:00:04Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the main assistant." + }, + { + "role": "user", + "content": "go" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "ok" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "span_begin", + "id": "helper", + "name": "helper", + "type": "agent", + "parent_id": "monitor", + "timestamp": "2026-01-01T00:00:05Z", + "uuid": "sb-helper" + }, + { + "event": "model", + "uuid": "m2", + "span_id": "helper", + "timestamp": "2026-01-01T00:00:06Z", + "completed": "2026-01-01T00:00:07Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are a helper." + }, + { + "role": "user", + "content": "help" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "ok" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "span_end", + "id": "helper", + "timestamp": "2026-01-01T00:00:08Z", + "uuid": "se-helper" + }, + { + "event": "span_end", + "id": "monitor", + "timestamp": "2026-01-01T00:00:09Z", + "uuid": "se-monitor" + }, + { + "event": "span_end", + "id": "solvers", + "timestamp": "2026-01-01T00:00:10Z", + "uuid": "se-solvers" + } + ], + "expected": { + "init": null, + "agent": { + "id": "monitor", + "name": "main", + "event_uuids": [ + "m1" + ], + "children": [ + { + "id": "helper", + "name": "helper", + "utility": false, + "event_uuids": [ + "m2" + ] + } + ] + }, + "scoring": null + } +} diff --git a/tests/timeline/fixtures/events/utility_workflow_generate.json b/tests/timeline/fixtures/events/utility_workflow_generate.json new file mode 100644 index 0000000000..2e972e1d48 --- /dev/null +++ b/tests/timeline/fixtures/events/utility_workflow_generate.json @@ -0,0 +1,243 @@ +{ + "description": "Workflow of generate() calls with mixed system prompts and no tool calls. With no tool-calling loop there is no primary trajectory, so no call may be wrapped as a hidden utility agent \u2014 all five model events stay direct content of the main agent.", + "events": [ + { + "event": "span_begin", + "id": "solvers", + "name": "solvers", + "type": "solvers", + "parent_id": null, + "timestamp": "2026-01-01T00:00:01Z", + "uuid": "sb-solvers" + }, + { + "event": "span_begin", + "id": "solver", + "name": "monitor", + "type": "solver", + "parent_id": "solvers", + "timestamp": "2026-01-01T00:00:02Z", + "uuid": "sb-solver" + }, + { + "event": "model", + "uuid": "m1", + "span_id": "solver", + "timestamp": "2026-01-01T00:00:03Z", + "completed": "2026-01-01T00:00:04Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the x-extract checker." + }, + { + "role": "user", + "content": "x-extract-1" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "ok" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "model", + "uuid": "m2", + "span_id": "solver", + "timestamp": "2026-01-01T00:00:05Z", + "completed": "2026-01-01T00:00:06Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the x-extract checker." + }, + { + "role": "user", + "content": "x-extract-2" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "ok" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "model", + "uuid": "m3", + "span_id": "solver", + "timestamp": "2026-01-01T00:00:07Z", + "completed": "2026-01-01T00:00:08Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the y-extract checker." + }, + { + "role": "user", + "content": "y-extract-1" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "ok" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "model", + "uuid": "m4", + "span_id": "solver", + "timestamp": "2026-01-01T00:00:09Z", + "completed": "2026-01-01T00:00:10Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the y-ground checker." + }, + { + "role": "user", + "content": "y-ground-1" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "ok" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "model", + "uuid": "m5", + "span_id": "solver", + "timestamp": "2026-01-01T00:00:11Z", + "completed": "2026-01-01T00:00:12Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the y-ground checker." + }, + { + "role": "user", + "content": "y-ground-2" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "ok" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "span_end", + "id": "solver", + "timestamp": "2026-01-01T00:00:13Z", + "uuid": "se-solver" + }, + { + "event": "span_end", + "id": "solvers", + "timestamp": "2026-01-01T00:00:14Z", + "uuid": "se-solvers" + } + ], + "expected": { + "init": null, + "agent": { + "id": "solvers", + "name": "main", + "event_uuids": [ + "sb-solver", + "m1", + "m2", + "m3", + "m4", + "m5", + "se-solver" + ], + "children": [] + }, + "scoring": null + } +} diff --git a/tests/timeline/generate.py b/tests/timeline/generate.py index 2d564a5b83..2648364c06 100644 --- a/tests/timeline/generate.py +++ b/tests/timeline/generate.py @@ -12,8 +12,6 @@ 6. Sequential run – multiple run() calls with inference in between 7. Parallel collect – three parallel sub-agents via collect() 8. Handoff + as_tool – combining handoff() and as_tool() sub-agents - 9. Workflow generate – mixed-prompt generate() calls, no tools → no utility - 10. Bridge utility – foreign-prompt call within a tool loop → utility """ from __future__ import annotations @@ -32,21 +30,9 @@ from inspect_ai.event import Timeline, timeline_build from inspect_ai.event._timeline import TimelineSpan from inspect_ai.log import EvalLog -from inspect_ai.model import ( - ChatMessageSystem, - ChatMessageUser, - ModelOutput, - get_model, -) +from inspect_ai.model import ModelOutput, get_model from inspect_ai.scorer import includes -from inspect_ai.solver import ( - Generate, - Solver, - TaskState, - generate, - solver, - system_message, -) +from inspect_ai.solver import generate, system_message from inspect_ai.tool import tool from inspect_ai.util import collect @@ -615,89 +601,6 @@ def scenario_deep_utility() -> tuple[str, Task, Any]: return "deep_utility", task, model -def scenario_workflow_generate() -> tuple[str, Task, Any]: - """Workflow of generate() calls with mixed system prompts and no tools. - - A monitor-style solver makes five auxiliary model calls — two sharing - one system prompt, three using different prompts. With no tool-calling - loop anywhere there is no primary trajectory, so none of the calls may - be hidden as utility agents. - """ - - @solver - def monitor() -> Solver: - async def solve(state: TaskState, generate: Generate) -> TaskState: - model = get_model(MODEL) - calls = [ - ("x-extract-1", "You are the x-extract checker."), - ("x-extract-2", "You are the x-extract checker."), - ("y-extract-1", "You are the y-extract checker."), - ("y-ground-1", "You are the y-ground checker."), - ("y-ground-2", "You are the y-ground checker."), - ] - for call_name, prompt in calls: - await model.generate( - [ - ChatMessageSystem(content=prompt), - ChatMessageUser(content=f"{call_name}: {state.input_text}"), - ] - ) - return state - - return solve - - model = get_model(MODEL) - task = Task(name="workflow_generate", dataset=DATASET, solver=monitor()) - return "workflow_generate", task, model - - -def scenario_bridge_utility() -> tuple[str, Task, Any]: - """Foreign-prompt helper call within a tool-calling loop → utility. - - A bridge-style agent emits a tool-calling main loop plus one extraction - call with a different system prompt and no tool calls. The extraction - call is wrapped as a utility span; main-loop calls stay visible. - """ - main_prompt = "You are the main assistant." - - @agent - def bridge() -> Agent: - async def execute(state: AgentState) -> AgentState: - model = get_model( - MODEL, - custom_outputs=[ - ModelOutput.for_tool_call(MODEL, "search", {"query": "1 + 1"}), - ModelOutput.from_content(MODEL, "two"), - ModelOutput.from_content(MODEL, "2"), - ], - ) - await model.generate( - [ - ChatMessageSystem(content=main_prompt), - ChatMessageUser(content="What is 1 + 1?"), - ] - ) - await model.generate( - [ - ChatMessageSystem(content="Extract a short title."), - ChatMessageUser(content="Extract."), - ] - ) - state.output = await model.generate( - [ - ChatMessageSystem(content=main_prompt), - ChatMessageUser(content="Answer."), - ] - ) - return state - - return execute - - model = get_model(MODEL) - task = Task(name="bridge_utility", dataset=DATASET, solver=bridge()) - return "bridge_utility", task, model - - def scenario_parallel_heterogeneous() -> tuple[str, Task, Any]: """Parallel agents with different names. @@ -1065,50 +968,6 @@ def validate_deep_utility(timeline: Timeline) -> None: assert_repr_labels(timeline, "main", "dispatcher", "lookup") -def _utility_spans(span: TimelineSpan) -> list[TimelineSpan]: - """Collect all descendant spans classified as utility.""" - results: list[TimelineSpan] = [] - for item in span.content: - if isinstance(item, TimelineSpan): - if item.utility: - results.append(item) - results.extend(_utility_spans(item)) - return results - - -def _direct_model_events(span: TimelineSpan) -> list[Any]: - return [ - item.event - for item in span.content - if not isinstance(item, TimelineSpan) and item.event.event == "model" - ] - - -def validate_workflow_generate(timeline: Timeline) -> None: - root = timeline.root - assert not _utility_spans(root), ( - "Workflow generate() calls must not be classified as utility agents" - ) - model_events = _direct_model_events(root) - assert len(model_events) == 5, ( - f"Expected all 5 workflow calls direct in main, got {len(model_events)}" - ) - - -def validate_bridge_utility(timeline: Timeline) -> None: - root = timeline.root - utility = _utility_spans(root) - assert len(utility) == 1, f"Expected exactly 1 utility span, got {len(utility)}" - wrapped = _direct_model_events(utility[0]) - assert len(wrapped) == 1, "Expected the utility span to wrap a single call" - system = next(m.content for m in wrapped[0].input if m.role == "system") - assert "Extract" in str(system), "Expected the extraction call to be wrapped" - main_events = _direct_model_events(root) - assert len(main_events) == 2, ( - f"Expected 2 main-loop calls direct in main, got {len(main_events)}" - ) - - def validate_parallel_heterogeneous(timeline: Timeline) -> None: root = timeline.root search_spans = find_spans(root, "search") @@ -1138,8 +997,6 @@ def validate_parallel_heterogeneous(timeline: Timeline) -> None: "parallel_with_nesting": validate_parallel_with_nesting, "sequential_and_parallel": validate_sequential_and_parallel, "deep_utility": validate_deep_utility, - "workflow_generate": validate_workflow_generate, - "bridge_utility": validate_bridge_utility, "parallel_heterogeneous": validate_parallel_heterogeneous, } @@ -1160,8 +1017,6 @@ def validate_parallel_heterogeneous(timeline: Timeline) -> None: scenario_parallel_with_nesting, scenario_sequential_and_parallel, scenario_deep_utility, - scenario_workflow_generate, - scenario_bridge_utility, scenario_parallel_heterogeneous, ] diff --git a/tests/timeline/test_timeline.py b/tests/timeline/test_timeline.py index bdf71b732a..ee528d48b3 100644 --- a/tests/timeline/test_timeline.py +++ b/tests/timeline/test_timeline.py @@ -9,7 +9,6 @@ from inspect_ai.event import Timeline, timeline_build from .generate import ( - scenario_bridge_utility, scenario_deep_nesting, scenario_deep_utility, scenario_handoff_and_as_tool, @@ -23,8 +22,6 @@ scenario_sequential_run, scenario_simple_agent, scenario_utility_agent, - scenario_workflow_generate, - validate_bridge_utility, validate_deep_nesting, validate_deep_utility, validate_handoff_and_as_tool, @@ -38,7 +35,6 @@ validate_sequential_run, validate_simple_agent, validate_utility_agent, - validate_workflow_generate, ) @@ -105,13 +101,5 @@ def test_timeline_deep_utility() -> None: _run_and_validate(scenario_deep_utility, validate_deep_utility) -def test_timeline_workflow_generate() -> None: - _run_and_validate(scenario_workflow_generate, validate_workflow_generate) - - -def test_timeline_bridge_utility() -> None: - _run_and_validate(scenario_bridge_utility, validate_bridge_utility) - - def test_timeline_parallel_heterogeneous() -> None: _run_and_validate(scenario_parallel_heterogeneous, validate_parallel_heterogeneous) diff --git a/tests/timeline/test_timeline_fixtures.py b/tests/timeline/test_timeline_fixtures.py new file mode 100644 index 0000000000..025fa4dd21 --- /dev/null +++ b/tests/timeline/test_timeline_fixtures.py @@ -0,0 +1,155 @@ +"""Fixture-driven tests for timeline_build(). + +Each JSON file in fixtures/events/ defines a flat event stream plus the +expected timeline structure. The same fixtures drive the TypeScript +implementation (ts-mono's fixtureTimeline.test.ts), keeping the Python and +TypeScript timeline builders consistent without duplicating test logic. + +Fixture format: {"description": ..., "events": [...], "expected": {...}} +with expected sections "init", "agent", and "scoring" (see the assertion +helpers below for the supported expectation fields). +""" + +import json +from pathlib import Path +from typing import Any + +import pytest +from pydantic import TypeAdapter + +from inspect_ai.event import Event, Timeline, timeline_build +from inspect_ai.event._timeline import TimelineEvent, TimelineSpan + +FIXTURES_DIR = Path(__file__).parent / "fixtures" / "events" + +_events_adapter = TypeAdapter(list[Event]) + + +def _fixture_names() -> list[str]: + return sorted(f.stem for f in FIXTURES_DIR.glob("*.json")) + + +def _load_fixture(name: str) -> dict[str, Any]: + with open(FIXTURES_DIR / f"{name}.json") as f: + data: dict[str, Any] = json.load(f) + return data + + +# ============================================================================= +# Assertion helpers (mirror ts-mono fixtureTimeline.test.ts) +# ============================================================================= + + +def _direct_event_uuids(span: TimelineSpan) -> list[str]: + return [ + item.event.uuid + for item in span.content + if isinstance(item, TimelineEvent) and item.event.uuid is not None + ] + + +def _child_spans(span: TimelineSpan) -> list[TimelineSpan]: + return [item for item in span.content if isinstance(item, TimelineSpan)] + + +def _all_event_uuids(span: TimelineSpan) -> list[str]: + uuids: list[str] = [] + for item in span.content: + if isinstance(item, TimelineEvent): + if item.event.uuid is not None: + uuids.append(item.event.uuid) + else: + uuids.extend(_all_event_uuids(item)) + return uuids + + +def _assert_span_matches(actual: TimelineSpan | None, expected: dict[str, Any]) -> None: + assert actual is not None + assert actual.id == expected["id"] + assert actual.name == expected["name"] + + if "event_uuids" in expected: + assert _direct_event_uuids(actual) == expected["event_uuids"] + + if "nested_uuids" in expected: + assert _all_event_uuids(actual) == expected["nested_uuids"] + + if "total_tokens" in expected: + assert actual.total_tokens() == expected["total_tokens"] + + if "utility" in expected: + assert actual.utility == expected["utility"] + + if "agent_result" in expected: + assert actual.agent_result == expected["agent_result"] + + if "children" in expected: + children = _child_spans(actual) + assert len(children) == len(expected["children"]), ( + f"Expected {len(expected['children'])} child spans of " + f"'{actual.name}', got {[c.name for c in children]}" + ) + for child, expected_child in zip(children, expected["children"]): + _assert_span_matches(child, expected_child) + + +def _assert_section_matches( + root: TimelineSpan, span_type: str, expected: dict[str, Any] | None +) -> None: + sections = [ + item + for item in root.content + if isinstance(item, TimelineSpan) and item.span_type == span_type + ] + if expected is None: + assert len(sections) == 0, f"Expected no '{span_type}' section" + return + assert len(sections) == 1 + if "event_uuids" in expected: + assert _direct_event_uuids(sections[0]) == expected["event_uuids"] + + +def _assert_timeline_matches(timeline: Timeline, expected: dict[str, Any]) -> None: + root = timeline.root + + _assert_section_matches(root, "init", expected["init"]) + _assert_section_matches(root, "scorers", expected["scoring"]) + + expected_agent = expected["agent"] + if expected_agent is None: + return + + assert root.id == expected_agent["id"] + assert root.name == expected_agent["name"] + + if "event_uuids" in expected_agent: + assert _direct_event_uuids(root) == expected_agent["event_uuids"] + + if "utility" in expected_agent: + assert root.utility == expected_agent["utility"] + + if "children" in expected_agent: + children = [ + item + for item in _child_spans(root) + if item.span_type not in ("init", "scorers") + ] + assert len(children) == len(expected_agent["children"]), ( + f"Expected {len(expected_agent['children'])} child spans of root, " + f"got {[c.name for c in children]}" + ) + for child, expected_child in zip(children, expected_agent["children"]): + _assert_span_matches(child, expected_child) + + +# ============================================================================= +# Tests +# ============================================================================= + + +@pytest.mark.parametrize("fixture_name", _fixture_names()) +def test_timeline_fixture(fixture_name: str) -> None: + fixture = _load_fixture(fixture_name) + events = _events_adapter.validate_python(fixture["events"]) + timeline = timeline_build(events) + _assert_timeline_matches(timeline, fixture["expected"]) From 1f0e064926219d4f5f1c1233a204060619eb6ca3 Mon Sep 17 00:00:00 2001 From: Charles Teague Date: Wed, 15 Jul 2026 15:33:54 -0400 Subject: [PATCH 3/4] timeline: fix utility-wrapper recursion; run viewer tests in CI Two fixes from code review: - _wrap_utility_events recursed into the synthetic wrapper spans it had just created; a wrapper's warmup call was wrapped again inside it, forever (RecursionError on any transcript with a cache-priming call). Recursion now visits only the span's original children. Also skips system-prompt extraction when there is no primary prompt to compare against. Adds a utility_warmup fixture covering warmup wrapping (which is prompt-independent and applies without an agentic loop). - log_viewer.yml gains a viewer-tests job running the inspect-components vitest suite: the fixture-driven timeline tests consume this repo's fixtures and are skipped in ts-mono's standalone CI, so without this job they ran nowhere. Co-Authored-By: Claude Fable 5 --- .github/workflows/log_viewer.yml | 30 ++++++ src/inspect_ai/event/_timeline.py | 40 ++++--- .../fixtures/events/utility_warmup.json | 102 ++++++++++++++++++ 3 files changed, 151 insertions(+), 21 deletions(-) create mode 100644 tests/timeline/fixtures/events/utility_warmup.json diff --git a/.github/workflows/log_viewer.yml b/.github/workflows/log_viewer.yml index 93f58e4ebc..b85690b3ea 100644 --- a/.github/workflows/log_viewer.yml +++ b/.github/workflows/log_viewer.yml @@ -121,6 +121,36 @@ jobs: exit 1 fi + # Runs the viewer test suite in the inspect_ai embedding. This includes the + # fixture-driven timeline tests (packages/inspect-components), which consume + # tests/timeline/fixtures/events/ from this repo and are skipped in ts-mono's + # own standalone CI — without this job they run nowhere. + viewer-tests: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + with: + submodules: true + + - uses: pnpm/action-setup@v6 + with: + package_json_file: src/inspect_ai/_view/ts-mono/package.json + + - name: Set up Node.js + uses: actions/setup-node@v6 + with: + node-version: "22.x" + cache: pnpm + cache-dependency-path: src/inspect_ai/_view/ts-mono/pnpm-lock.yaml + + - name: Install dependencies + working-directory: src/inspect_ai/_view/ts-mono + run: pnpm install --frozen-lockfile + + - name: Run inspect-components tests + working-directory: src/inspect_ai/_view/ts-mono + run: pnpm --filter @tsmono/inspect-components test + dist-validation: runs-on: ubuntu-latest steps: diff --git a/src/inspect_ai/event/_timeline.py b/src/inspect_ai/event/_timeline.py index 5cac7d4970..b0c314857e 100644 --- a/src/inspect_ai/event/_timeline.py +++ b/src/inspect_ai/event/_timeline.py @@ -1179,6 +1179,7 @@ def _wrap_utility_events(agent: TimelineSpan) -> None: break # --- Scan and wrap utility candidates --- + original_spans = [item for item in agent.content if isinstance(item, TimelineSpan)] new_content: list[TimelineEvent | TimelineSpan] = [] for item in agent.content: if isinstance(item, TimelineEvent) and isinstance(item.event, ModelEvent): @@ -1194,31 +1195,28 @@ def _wrap_utility_events(agent: TimelineSpan) -> None: new_content.append(wrapper) continue - evt_prompt = _get_system_prompt_for_event(item.event) - if ( - primary_prompt is not None - and evt_prompt is not None - and evt_prompt != primary_prompt - and not _has_tool_calls(item.event) - ): - # Wrap in a synthetic utility span - wrapper = TimelineSpan( - id=f"utility-{item.event.uuid or id(item)}", - name="utility", - span_type="agent", - content=[item], - ) - wrapper.utility = True - new_content.append(wrapper) - continue + if primary_prompt is not None and not _has_tool_calls(item.event): + evt_prompt = _get_system_prompt_for_event(item.event) + if evt_prompt is not None and evt_prompt != primary_prompt: + # Wrap in a synthetic utility span + wrapper = TimelineSpan( + id=f"utility-{item.event.uuid or id(item)}", + name="utility", + span_type="agent", + content=[item], + ) + wrapper.utility = True + new_content.append(wrapper) + continue new_content.append(item) agent.content = new_content - # --- Recurse into child spans and branches --- - for item in agent.content: - if isinstance(item, TimelineSpan): - _wrap_utility_events(item) + # Recurse into the span's original children only — not the synthetic + # wrappers created above: a wrapper holds a single already-processed + # event, and re-entering it would wrap a warmup call again, forever. + for item in original_spans: + _wrap_utility_events(item) for branch in agent.branches: for item in branch.content: if isinstance(item, TimelineSpan): diff --git a/tests/timeline/fixtures/events/utility_warmup.json b/tests/timeline/fixtures/events/utility_warmup.json new file mode 100644 index 0000000000..dcdc240025 --- /dev/null +++ b/tests/timeline/fixtures/events/utility_warmup.json @@ -0,0 +1,102 @@ +{ + "description": "Cache-priming warmup call (max_tokens=1, single-word user message) alongside a normal call, with no tool-calling loop. Warmup detection is prompt-independent, so the warmup is wrapped as a utility span even without a primary trajectory — and wrapping must not recurse into its own synthetic wrapper.", + "events": [ + { + "event": "span_begin", + "id": "solvers", + "name": "solvers", + "type": "solvers", + "parent_id": null, + "timestamp": "2026-01-01T00:00:01Z", + "uuid": "sb-solvers" + }, + { + "event": "span_begin", + "id": "monitor", + "name": "monitor", + "type": "agent", + "parent_id": "solvers", + "timestamp": "2026-01-01T00:00:02Z", + "uuid": "sb-monitor" + }, + { + "event": "model", + "uuid": "w1", + "span_id": "monitor", + "timestamp": "2026-01-01T00:00:03Z", + "completed": "2026-01-01T00:00:04Z", + "model": "mockllm/model", + "input": [ + { "role": "system", "content": "You are the main assistant." }, + { "role": "user", "content": "warmup" } + ], + "tools": [], + "tool_choice": "none", + "config": { "max_tokens": 1 }, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { "role": "assistant", "content": "w" }, + "stop_reason": "max_tokens" + } + ], + "usage": { "input_tokens": 5, "output_tokens": 1, "total_tokens": 6 } + } + }, + { + "event": "model", + "uuid": "m1", + "span_id": "monitor", + "timestamp": "2026-01-01T00:00:05Z", + "completed": "2026-01-01T00:00:06Z", + "model": "mockllm/model", + "input": [ + { "role": "system", "content": "You are the main assistant." }, + { "role": "user", "content": "What is 1 + 1?" } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { "role": "assistant", "content": "2" }, + "stop_reason": "stop" + } + ], + "usage": { "input_tokens": 5, "output_tokens": 1, "total_tokens": 6 } + } + }, + { + "event": "span_end", + "id": "monitor", + "timestamp": "2026-01-01T00:00:07Z", + "uuid": "se-monitor" + }, + { + "event": "span_end", + "id": "solvers", + "timestamp": "2026-01-01T00:00:08Z", + "uuid": "se-solvers" + } + ], + "expected": { + "init": null, + "agent": { + "id": "monitor", + "name": "main", + "event_uuids": ["m1"], + "children": [ + { + "id": "utility-w1", + "name": "utility", + "utility": true, + "event_uuids": ["w1"] + } + ] + }, + "scoring": null + } +} From eacfbe68ff81c348725d86b87d32ae24de40b7c7 Mon Sep 17 00:00:00 2001 From: Charles Teague Date: Wed, 15 Jul 2026 15:52:49 -0400 Subject: [PATCH 4/4] timeline: align utility classification edge cases with the TS port Four parity fixes between timeline_build and ts-mono's buildTimeline, each now covered by a shared JSON fixture (or unit test where the fixture format can't express the case): - _get_system_prompt (classification) now delegates to _get_system_prompt_for_event, so billing-header normalization applies to span-level prompt comparison (previously raw prompts were compared, classifying Claude Code sub-agents differently than the viewer). - Prompt extraction returns None instead of "" when a prompt is empty after normalization (e.g. only the billing header line), keeping empty prompts out of comparisons and prompt inheritance. - _wrap_utility_events processes a branch's own direct events (branches are trajectories like any span), matching the TS port. - Wrapper spans for uuid-less (legacy) events get unique deterministic position-derived ids (utility--) instead of id(item). Fixtures: utility_billing_header, utility_header_only_prompt, utility_branch_warmup (the fixture harness gains branch assertions). Co-Authored-By: Claude Fable 5 --- src/inspect_ai/event/_timeline.py | 69 +++-- .../events/utility_billing_header.json | 210 ++++++++++++++ .../events/utility_branch_warmup.json | 267 ++++++++++++++++++ .../events/utility_header_only_prompt.json | 210 ++++++++++++++ tests/timeline/test_timeline.py | 56 ++++ tests/timeline/test_timeline_fixtures.py | 26 ++ 6 files changed, 802 insertions(+), 36 deletions(-) create mode 100644 tests/timeline/fixtures/events/utility_billing_header.json create mode 100644 tests/timeline/fixtures/events/utility_branch_warmup.json create mode 100644 tests/timeline/fixtures/events/utility_header_only_prompt.json diff --git a/src/inspect_ai/event/_timeline.py b/src/inspect_ai/event/_timeline.py index b0c314857e..20d6990b55 100644 --- a/src/inspect_ai/event/_timeline.py +++ b/src/inspect_ai/event/_timeline.py @@ -1128,15 +1128,21 @@ def _get_system_prompt_for_event(event: ModelEvent) -> str | None: event: The ModelEvent to inspect. Returns: - The normalized system prompt text, or None if no system message found. + The normalized system prompt text, or None if no system message is + found or the prompt is empty after normalization (e.g. it consisted + only of the billing header). Returning None rather than "" keeps + empty prompts out of primary-trajectory comparisons and prompt + inheritance. """ for msg in event.input: if isinstance(msg, ChatMessageSystem): if isinstance(msg.content, str): - return _normalize_system_prompt(msg.content) - parts = [c.text for c in msg.content if hasattr(c, "text")] - raw = "\n".join(parts) if parts else None - return _normalize_system_prompt(raw) if raw else None + raw = msg.content + else: + parts = [c.text for c in msg.content if hasattr(c, "text")] + raw = "\n".join(parts) + normalized = _normalize_system_prompt(raw) + return normalized if normalized else None return None @@ -1178,35 +1184,32 @@ def _wrap_utility_events(agent: TimelineSpan) -> None: primary_prompt = _get_system_prompt_for_event(item.event) break + def utility_wrapper(item: TimelineEvent, index: int) -> TimelineSpan: + # Fall back to a position-derived id for uuid-less (legacy) events so + # ids stay unique and deterministic across rebuilds. + wrapper = TimelineSpan( + id=f"utility-{item.event.uuid or f'{agent.id}-{index}'}", + name="utility", + span_type="agent", + content=[item], + ) + wrapper.utility = True + return wrapper + # --- Scan and wrap utility candidates --- original_spans = [item for item in agent.content if isinstance(item, TimelineSpan)] new_content: list[TimelineEvent | TimelineSpan] = [] - for item in agent.content: + for index, item in enumerate(agent.content): if isinstance(item, TimelineEvent) and isinstance(item.event, ModelEvent): # Warmup/cache-priming call (max_tokens=1) if _is_warmup_call(item.event): - wrapper = TimelineSpan( - id=f"utility-{item.event.uuid or id(item)}", - name="utility", - span_type="agent", - content=[item], - ) - wrapper.utility = True - new_content.append(wrapper) + new_content.append(utility_wrapper(item, index)) continue if primary_prompt is not None and not _has_tool_calls(item.event): evt_prompt = _get_system_prompt_for_event(item.event) if evt_prompt is not None and evt_prompt != primary_prompt: - # Wrap in a synthetic utility span - wrapper = TimelineSpan( - id=f"utility-{item.event.uuid or id(item)}", - name="utility", - span_type="agent", - content=[item], - ) - wrapper.utility = True - new_content.append(wrapper) + new_content.append(utility_wrapper(item, index)) continue new_content.append(item) @@ -1215,12 +1218,12 @@ def _wrap_utility_events(agent: TimelineSpan) -> None: # Recurse into the span's original children only — not the synthetic # wrappers created above: a wrapper holds a single already-processed # event, and re-entering it would wrap a warmup call again, forever. + # Branches are trajectories like any span: process their own direct + # events too (mirrors the TS port), which also covers their children. for item in original_spans: _wrap_utility_events(item) for branch in agent.branches: - for item in branch.content: - if isinstance(item, TimelineSpan): - _wrap_utility_events(item) + _wrap_utility_events(branch) def _is_warmup_call(event: ModelEvent) -> bool: @@ -1242,24 +1245,18 @@ def _is_warmup_call(event: ModelEvent) -> bool: def _get_system_prompt(agent: TimelineSpan) -> str | None: - """Extract system prompt from the first ModelEvent in agent's direct content. + """Extract the system prompt from the first ModelEvent in agent's direct content. Args: agent: The span node to extract the system prompt from. Returns: - The system prompt text, or None if no system message found. + The normalized system prompt text (see ``_get_system_prompt_for_event``), + or None if the span has no ModelEvent or it carries no system message. """ for item in agent.content: if isinstance(item, TimelineEvent) and isinstance(item.event, ModelEvent): - for msg in item.event.input: - if isinstance(msg, ChatMessageSystem): - if isinstance(msg.content, str): - return msg.content - # Content is list of Content objects - parts = [c.text for c in msg.content if hasattr(c, "text")] - return "\n".join(parts) if parts else None - return None # ModelEvent found but no system message + return _get_system_prompt_for_event(item.event) return None # No ModelEvent found diff --git a/tests/timeline/fixtures/events/utility_billing_header.json b/tests/timeline/fixtures/events/utility_billing_header.json new file mode 100644 index 0000000000..9970ea2feb --- /dev/null +++ b/tests/timeline/fixtures/events/utility_billing_header.json @@ -0,0 +1,210 @@ +{ + "description": "Parent and single-turn child whose system prompts are identical except for the varying x-anthropic-billing-header first line (Claude Code-style logs). Both implementations must normalize the header away before comparing, so the child is NOT classified as utility.", + "events": [ + { + "event": "span_begin", + "uuid": "sb-solvers", + "id": "solvers", + "name": "solvers", + "type": "solvers", + "parent_id": null, + "timestamp": "2026-01-01T00:00:01Z" + }, + { + "event": "span_begin", + "uuid": "sb-bridge", + "id": "bridge", + "name": "bridge", + "type": "agent", + "parent_id": "solvers", + "timestamp": "2026-01-01T00:00:02Z" + }, + { + "event": "model", + "uuid": "m1", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:03Z", + "completed": "2026-01-01T00:00:03Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "x-anthropic-billing-header: cch=aaa111\nYou are the main assistant." + }, + { + "role": "user", + "content": "go" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "function": "search", + "arguments": { + "query": "1 + 1" + } + } + ] + }, + "stop_reason": "tool_calls" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "tool", + "uuid": "t1", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:04Z", + "completed": "2026-01-01T00:00:04Z", + "id": "call_1", + "function": "search", + "arguments": { + "query": "1 + 1" + }, + "result": "found it", + "events": [] + }, + { + "event": "span_begin", + "uuid": "sb-helper", + "id": "helper", + "name": "helper", + "type": "agent", + "parent_id": "bridge", + "timestamp": "2026-01-01T00:00:05Z" + }, + { + "event": "model", + "uuid": "m2", + "span_id": "helper", + "timestamp": "2026-01-01T00:00:06Z", + "completed": "2026-01-01T00:00:06Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "x-anthropic-billing-header: cch=bbb222\nYou are the main assistant." + }, + { + "role": "user", + "content": "go" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "ok" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "span_end", + "uuid": "se-helper", + "id": "helper", + "timestamp": "2026-01-01T00:00:07Z" + }, + { + "event": "model", + "uuid": "m3", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:08Z", + "completed": "2026-01-01T00:00:08Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "x-anthropic-billing-header: cch=ccc333\nYou are the main assistant." + }, + { + "role": "user", + "content": "go" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "ok" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "span_end", + "uuid": "se-bridge", + "id": "bridge", + "timestamp": "2026-01-01T00:00:09Z" + }, + { + "event": "span_end", + "uuid": "se-solvers", + "id": "solvers", + "timestamp": "2026-01-01T00:00:10Z" + } + ], + "expected": { + "init": null, + "agent": { + "id": "bridge", + "name": "main", + "event_uuids": [ + "m1", + "t1", + "m3" + ], + "children": [ + { + "id": "helper", + "name": "helper", + "utility": false, + "event_uuids": [ + "m2" + ] + } + ] + }, + "scoring": null + } +} diff --git a/tests/timeline/fixtures/events/utility_branch_warmup.json b/tests/timeline/fixtures/events/utility_branch_warmup.json new file mode 100644 index 0000000000..f716cfbc3a --- /dev/null +++ b/tests/timeline/fixtures/events/utility_branch_warmup.json @@ -0,0 +1,267 @@ +{ + "description": "A branch (reroll trajectory) whose direct content contains a warmup call. Branches are trajectories like any span: both implementations must process a branch's own events, wrapping the warmup as a utility span inside the branch.", + "events": [ + { + "event": "span_begin", + "uuid": "sb-solvers", + "id": "solvers", + "name": "solvers", + "type": "solvers", + "parent_id": null, + "timestamp": "2026-01-01T00:00:01Z" + }, + { + "event": "span_begin", + "uuid": "sb-bridge", + "id": "bridge", + "name": "bridge", + "type": "agent", + "parent_id": "solvers", + "timestamp": "2026-01-01T00:00:02Z" + }, + { + "event": "model", + "uuid": "m1", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:03Z", + "completed": "2026-01-01T00:00:03Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the main assistant." + }, + { + "role": "user", + "content": "go" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "function": "search", + "arguments": { + "query": "1 + 1" + } + } + ] + }, + "stop_reason": "tool_calls" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "tool", + "uuid": "t1", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:04Z", + "completed": "2026-01-01T00:00:04Z", + "id": "call_1", + "function": "search", + "arguments": { + "query": "1 + 1" + }, + "result": "found it", + "events": [] + }, + { + "event": "span_begin", + "uuid": "sb-br1", + "id": "br1", + "name": "branch 1", + "type": "branch", + "parent_id": "bridge", + "timestamp": "2026-01-01T00:00:05Z" + }, + { + "event": "branch", + "uuid": "b1", + "span_id": "br1", + "timestamp": "2026-01-01T00:00:06Z", + "from_anchor": "a1" + }, + { + "event": "model", + "uuid": "w1", + "span_id": "br1", + "timestamp": "2026-01-01T00:00:07Z", + "completed": "2026-01-01T00:00:07Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the main assistant." + }, + { + "role": "user", + "content": "warmup" + } + ], + "tools": [], + "tool_choice": "none", + "config": { + "max_tokens": 1 + }, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "w" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "model", + "uuid": "m2", + "span_id": "br1", + "timestamp": "2026-01-01T00:00:08Z", + "completed": "2026-01-01T00:00:08Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the main assistant." + }, + { + "role": "user", + "content": "go" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "ok" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "span_end", + "uuid": "se-br1", + "id": "br1", + "timestamp": "2026-01-01T00:00:09Z" + }, + { + "event": "model", + "uuid": "m3", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:10Z", + "completed": "2026-01-01T00:00:10Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the main assistant." + }, + { + "role": "user", + "content": "go" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "ok" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "span_end", + "uuid": "se-bridge", + "id": "bridge", + "timestamp": "2026-01-01T00:00:11Z" + }, + { + "event": "span_end", + "uuid": "se-solvers", + "id": "solvers", + "timestamp": "2026-01-01T00:00:12Z" + } + ], + "expected": { + "init": null, + "agent": { + "id": "bridge", + "name": "main", + "event_uuids": [ + "m1", + "t1", + "m3" + ], + "children": [], + "branches": [ + { + "branched_from": "a1", + "event_uuids": [ + "b1", + "m2" + ], + "children": [ + { + "id": "utility-w1", + "name": "utility", + "utility": true, + "event_uuids": [ + "w1" + ] + } + ] + } + ] + }, + "scoring": null + } +} diff --git a/tests/timeline/fixtures/events/utility_header_only_prompt.json b/tests/timeline/fixtures/events/utility_header_only_prompt.json new file mode 100644 index 0000000000..e316369b30 --- /dev/null +++ b/tests/timeline/fixtures/events/utility_header_only_prompt.json @@ -0,0 +1,210 @@ +{ + "description": "Single-turn child whose system prompt is ONLY the billing header line, which normalizes to empty. An empty prompt is treated as no prompt at all (None), so the child is NOT classified as utility \u2014 there is no evidence its prompt differs from the parent's.", + "events": [ + { + "event": "span_begin", + "uuid": "sb-solvers", + "id": "solvers", + "name": "solvers", + "type": "solvers", + "parent_id": null, + "timestamp": "2026-01-01T00:00:01Z" + }, + { + "event": "span_begin", + "uuid": "sb-bridge", + "id": "bridge", + "name": "bridge", + "type": "agent", + "parent_id": "solvers", + "timestamp": "2026-01-01T00:00:02Z" + }, + { + "event": "model", + "uuid": "m1", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:03Z", + "completed": "2026-01-01T00:00:03Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the main assistant." + }, + { + "role": "user", + "content": "go" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "function": "search", + "arguments": { + "query": "1 + 1" + } + } + ] + }, + "stop_reason": "tool_calls" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "tool", + "uuid": "t1", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:04Z", + "completed": "2026-01-01T00:00:04Z", + "id": "call_1", + "function": "search", + "arguments": { + "query": "1 + 1" + }, + "result": "found it", + "events": [] + }, + { + "event": "span_begin", + "uuid": "sb-helper", + "id": "helper", + "name": "helper", + "type": "agent", + "parent_id": "bridge", + "timestamp": "2026-01-01T00:00:05Z" + }, + { + "event": "model", + "uuid": "m2", + "span_id": "helper", + "timestamp": "2026-01-01T00:00:06Z", + "completed": "2026-01-01T00:00:06Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "x-anthropic-billing-header: cch=zzz999" + }, + { + "role": "user", + "content": "go" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "ok" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "span_end", + "uuid": "se-helper", + "id": "helper", + "timestamp": "2026-01-01T00:00:07Z" + }, + { + "event": "model", + "uuid": "m3", + "span_id": "bridge", + "timestamp": "2026-01-01T00:00:08Z", + "completed": "2026-01-01T00:00:08Z", + "model": "mockllm/model", + "input": [ + { + "role": "system", + "content": "You are the main assistant." + }, + { + "role": "user", + "content": "go" + } + ], + "tools": [], + "tool_choice": "none", + "config": {}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": { + "role": "assistant", + "content": "ok" + }, + "stop_reason": "stop" + } + ], + "usage": { + "input_tokens": 5, + "output_tokens": 1, + "total_tokens": 6 + } + } + }, + { + "event": "span_end", + "uuid": "se-bridge", + "id": "bridge", + "timestamp": "2026-01-01T00:00:09Z" + }, + { + "event": "span_end", + "uuid": "se-solvers", + "id": "solvers", + "timestamp": "2026-01-01T00:00:10Z" + } + ], + "expected": { + "init": null, + "agent": { + "id": "bridge", + "name": "main", + "event_uuids": [ + "m1", + "t1", + "m3" + ], + "children": [ + { + "id": "helper", + "name": "helper", + "utility": false, + "event_uuids": [ + "m2" + ] + } + ] + }, + "scoring": null + } +} diff --git a/tests/timeline/test_timeline.py b/tests/timeline/test_timeline.py index ee528d48b3..35926dbb2b 100644 --- a/tests/timeline/test_timeline.py +++ b/tests/timeline/test_timeline.py @@ -103,3 +103,59 @@ def test_timeline_deep_utility() -> None: def test_timeline_parallel_heterogeneous() -> None: _run_and_validate(scenario_parallel_heterogeneous, validate_parallel_heterogeneous) + + +def test_wrap_utility_uuidless_ids_unique() -> None: + """Wrapper spans for uuid-less (legacy) events get unique deterministic ids. + + The JSON fixtures can't express this case — event parsing auto-assigns + uuids — so construct the events directly and clear them. + """ + from pydantic import TypeAdapter + + from inspect_ai.event import Event + from inspect_ai.event._timeline import ( + TimelineEvent, + TimelineSpan, + _wrap_utility_events, + ) + + def warmup_event(ts: str) -> Event: + event: Event = TypeAdapter(Event).validate_python( + { + "event": "model", + "timestamp": ts, + "completed": ts, + "model": "mockllm/model", + "input": [{"role": "user", "content": "warmup"}], + "tools": [], + "tool_choice": "none", + "config": {"max_tokens": 1}, + "output": { + "model": "mockllm/model", + "choices": [ + { + "message": {"role": "assistant", "content": "w"}, + "stop_reason": "max_tokens", + } + ], + }, + } + ) + event.uuid = None + return event + + span = TimelineSpan( + id="agent", + name="agent", + span_type="agent", + content=[ + TimelineEvent(event=warmup_event("2026-01-01T00:00:01Z")), + TimelineEvent(event=warmup_event("2026-01-01T00:00:02Z")), + ], + ) + _wrap_utility_events(span) + + wrappers = [item for item in span.content if isinstance(item, TimelineSpan)] + assert [w.id for w in wrappers] == ["utility-agent-0", "utility-agent-1"] + assert all(w.utility for w in wrappers) diff --git a/tests/timeline/test_timeline_fixtures.py b/tests/timeline/test_timeline_fixtures.py index 025fa4dd21..68a9ba0629 100644 --- a/tests/timeline/test_timeline_fixtures.py +++ b/tests/timeline/test_timeline_fixtures.py @@ -92,6 +92,29 @@ def _assert_span_matches(actual: TimelineSpan | None, expected: dict[str, Any]) for child, expected_child in zip(children, expected["children"]): _assert_span_matches(child, expected_child) + if "branches" in expected: + _assert_branches_match(actual, expected["branches"]) + + +def _assert_branches_match( + actual: TimelineSpan, expected_branches: list[dict[str, Any]] +) -> None: + assert len(actual.branches) == len(expected_branches), ( + f"Expected {len(expected_branches)} branches of '{actual.name}', " + f"got {len(actual.branches)}" + ) + for branch, expected in zip(actual.branches, expected_branches): + assert branch.branched_from == expected["branched_from"] + if "event_uuids" in expected: + assert _direct_event_uuids(branch) == expected["event_uuids"] + if "children" in expected: + children = _child_spans(branch) + assert len(children) == len(expected["children"]) + for child, expected_child in zip(children, expected["children"]): + _assert_span_matches(child, expected_child) + if "branches" in expected: + _assert_branches_match(branch, expected["branches"]) + def _assert_section_matches( root: TimelineSpan, span_type: str, expected: dict[str, Any] | None @@ -141,6 +164,9 @@ def _assert_timeline_matches(timeline: Timeline, expected: dict[str, Any]) -> No for child, expected_child in zip(children, expected_agent["children"]): _assert_span_matches(child, expected_child) + if "branches" in expected_agent: + _assert_branches_match(root, expected_agent["branches"]) + # ============================================================================= # Tests