feat: add trace visualization to display_sample_record (#396)

[nabinchha]

📋 Summary

Adds LLM conversation trace visualization to display_sample_record, allowing users to inspect the full chain-of-thought, tool calls, and tool results for
LLM-generated columns. Supports both terminal (Rich) and Jupyter notebook (HTML) rendering.

🔄 Changes

✨ Added

• New TraceRenderer class in [trace_renderer.py](https://github.com/NVIDIA-NeMo/DataDesigner/blob/nmulepati/feat/396-trace-visualization/packages/data-designer
  -config/src/data_designer/config/utils/trace_renderer.py) with Rich terminal and Jupyter HTML rendering
• is_notebook_environment() helper in misc.py for centralized notebook detection
• include_traces parameter on display_sample_record and WithRecordSamplerMixin (default True)
• Comprehensive test suite in [test_trace_renderer.py](https://github.com/NVIDIA-NeMo/DataDesigner/blob/nmulepati/feat/396-trace-visualization/packages/data-desi
  gner-config/tests/config/utils/test_trace_renderer.py) (338 lines, covering Rich rendering, HTML rendering, edge cases, and integration)

🔧 Changed

• visualization.py: consolidated notebook detection to use shared is_notebook_environment(), added trace rendering for all LLM column types, moved record index
  display to top, normalized padding defaults

🔍 Attention Areas

⚠️ Reviewers: Please pay special attention to the following:

• [trace_renderer.py](https://github.com/NVIDIA-NeMo/DataDesigner/blob/nmulepati/feat/396-trace-visualization/packages/data-designer-config/src/data_designer/con
  fig/utils/trace_renderer.py) — New module: core rendering logic for both Rich and HTML output, typed trace message structures
• [visualization.py](https://github.com/NVIDIA-NeMo/DataDesigner/blob/nmulepati/feat/396-trace-visualization/packages/data-designer-config/src/data_designer/conf
  ig/utils/visualization.py) — Integration point: trace column discovery, rendering order, and padding changes

In ipynb notebook

In console

---

🤖 Generated with AI

[greptile-apps]

Greptile Summary

This PR adds LLM conversation trace visualization to display_sample_record, introducing a new TraceRenderer class that renders chain-of-thought, tool calls, and tool results as Rich panels (terminal) or IPython HTML blocks (Jupyter). It also centralises notebook detection into a shared is_notebook_environment() helper and wires include_traces through WithRecordSamplerMixin.

Key changes:

• trace_renderer.py — new TraceRenderer class with render_rich and render_notebook_html methods, typed TraceMessage TypedDicts, and private HTML-building helpers.
• visualization.py — detects trace side-effect columns for all LLM column types, defers notebook HTML rendering until after the Rich console output, and moves the record-index label to the top.
• misc.py — adds is_notebook_environment() checking ZMQInteractiveShell to avoid false positives in plain IPython terminals.
• test_trace_renderer.py — 338-line suite covering Rich/HTML rendering, edge cases, and integration scenarios.

Issues found:

• Turn-count logic bug (trace_renderer.py:157): turn_ids accumulates one entry per tool-call ID, so a single assistant message with N parallel tool calls inflates turn_count to N instead of 1, producing a misleading summary (e.g. "2 tool calls in 2 turns" when only 1 LLM inference turn occurred).
• Trace content omitted from saved HTML in notebook mode (visualization.py:361): When in_notebook=True, trace panels are excluded from render_list and only sent to IPython.display. If save_path is also provided, the saved file will silently contain no trace data. The integration test does not cover this combination because it runs outside a notebook.

Confidence Score: 3/5

• Safe to merge with minor fixes; no critical runtime failures, but two display-correctness issues should be addressed before broader rollout.
• The new TraceRenderer is well-structured, properly escapes HTML, and has good test coverage for the happy paths. The two identified issues are display-accuracy bugs (wrong turn count for parallel tool calls) and a silent data-omission (traces not in saved HTML when called from a notebook with save_path), neither of which causes a crash. However, both would produce confusing or incomplete output for users, so they merit fixes before shipping.
• Pay close attention to trace_renderer.py (turn-count logic) and visualization.py (save_path + notebook interaction).

Important Files Changed

Filename	Overview
packages/data-designer-config/src/data_designer/config/utils/trace_renderer.py	New module providing Rich terminal and Jupyter HTML trace rendering; contains a logic error in the tool-call turn-count calculation that inflates "turns" for parallel tool calls.
packages/data-designer-config/src/data_designer/config/utils/visualization.py	Integrates trace rendering into display_sample_record; trace panels are excluded from render_list in notebook mode, causing them to be silently omitted from any saved HTML file when save_path is provided in a notebook.
packages/data-designer-config/src/data_designer/config/utils/misc.py	Adds is_notebook_environment() helper that correctly gates on ZMQInteractiveShell to avoid false positives; clean centralisation of a previously inline check.
packages/data-designer-config/tests/config/utils/test_trace_renderer.py	Comprehensive 338-line test suite covering Rich and HTML rendering paths; the saved-HTML integration test only runs in non-notebook mode so it does not cover the notebook+save_path omission case.

Sequence Diagram

sequenceDiagram
    participant Caller
    participant display_sample_record
    participant is_notebook_environment
    participant render_list
    participant TraceRenderer
    participant IPython

    Caller->>display_sample_record: record, config_builder, include_traces, save_path
    display_sample_record->>is_notebook_environment: check
    is_notebook_environment-->>display_sample_record: in_notebook (bool)

    alt in_notebook == True
        display_sample_record->>IPython: display HTML index label
    else
        display_sample_record->>render_list: append Text index label
    end

    display_sample_record->>render_list: append seed/generated/code/validation tables

    alt include_traces == True
        display_sample_record->>TraceRenderer: instantiate
        loop each LLM column with trace side-effect
            alt in_notebook == False
                TraceRenderer->>render_list: append Rich Panel (render_rich)
            end
            display_sample_record->>display_sample_record: traces_to_display_later.append(...)
        end
    end

    alt save_path set
        display_sample_record->>render_list: Console.print → save HTML/SVG
        Note over display_sample_record: traces NOT in saved file when in_notebook
    else
        display_sample_record->>render_list: Console.print to terminal
    end

    alt in_notebook and include_traces and traces_to_display_later
        TraceRenderer->>IPython: render_notebook_html → display(HTML(...))
    end

Loading

Prompt To Fix All With AI

This is a comment left during a code review.
Path: packages/data-designer-config/src/data_designer/config/utils/trace_renderer.py
Line: 157-160

Comment:
**Turn count is inflated for parallel tool calls**

`turn_ids` is populated with one entry per individual tool-call ID (line 124), so when a single assistant message issues *N* parallel tool calls (all IDs are distinct), `len(turn_ids)` equals N rather than 1. The summary will therefore read "2 tool calls in 2 turns" for the `multi_turn_tool_trace` fixture, which is a single LLM inference turn that returned two calls at once.

The "turns" semantic should count the number of assistant messages that contained tool calls — i.e. the number of separate LLM inference requests — not the number of unique call IDs. A simple fix:

```python
turn_count = sum(1 for msg in traces if msg.get("tool_calls")) if tool_call_count > 0 else 0
call_word = "call" if tool_call_count == 1 else "calls"
turn_word = "turn" if turn_count == 1 else "turns"
summary = f"{tool_call_count} tool {call_word} in {turn_count} {turn_word}" if tool_call_count > 0 else ""
```

This also makes the `turn_ids` set — which is currently only used to derive `turn_count` — unnecessary, and it can be removed.

How can I resolve this? If you propose a fix, please make it concise.

---

This is a comment left during a code review.
Path: packages/data-designer-config/src/data_designer/config/utils/visualization.py
Line: 361-363

Comment:
**Trace content is silently dropped from saved files in notebook mode**

In a Jupyter notebook (`in_notebook=True`) trace panels are intentionally excluded from `render_list` and deferred to `traces_to_display_later` for IPython HTML rendering. The problem is that the `save_path` code path runs *before* that deferred display, and it only captures what is in `render_list`:

```python
# visualization.py – save path captures render_list only
if save_path is not None:
    recording_console.print(Group(*render_list), markup=False)   # ← no trace panels here
    _save_console_output(recording_console, save_path, theme=theme)

# Traces are displayed afterwards, only to the live notebook cell
if in_notebook and include_traces and len(traces_to_display_later) > 0:
    trace_renderer.render_notebook_html(...)
```

A user who calls `display_sample_record(record, builder, save_path="report.html")` from a notebook cell will see traces in the cell output but their saved file will contain no trace data at all. `test_display_sample_record_with_trace_in_saved_html` does not catch this because it runs in a non-notebook environment where traces *do* enter `render_list`.

Consider also adding the Rich trace panel to `render_list` when `save_path` is explicitly requested (regardless of `in_notebook`), so saved files are always complete:

```python
if not in_notebook or save_path is not None:
    render_list.append(pad_console_element(trace_renderer.render_rich(traces, side_col)))
```

Alternatively, document the limitation clearly in the `include_traces` docstring so callers are not surprised.

How can I resolve this? If you propose a fix, please make it concise.

_{Last reviewed commit: "Merge branch 'main' ..."}