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/CHANGELOG.md b/CHANGELOG.md index 7d402286b2..bdb35329a6 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. - Bugfix: Nested `score_reducer` and `task_source` values in serialized task args now restore as instances rather than their registered factories. (#4374) - 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..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 @@ -1157,84 +1163,67 @@ 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 + 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 - evt_prompt = _get_system_prompt_for_event(item.event) - if ( - 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: + new_content.append(utility_wrapper(item, index)) + 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. + # 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: @@ -1256,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 @@ -1316,17 +1299,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 +1336,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/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_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_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/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_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 + } +} 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/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 new file mode 100644 index 0000000000..68a9ba0629 --- /dev/null +++ b/tests/timeline/test_timeline_fixtures.py @@ -0,0 +1,181 @@ +"""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) + + 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 +) -> 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) + + if "branches" in expected_agent: + _assert_branches_match(root, expected_agent["branches"]) + + +# ============================================================================= +# 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"])