feat: add trace visualization to display_sample_record (#396) by nabinchha · Pull Request #438 · NVIDIA-NeMo/DataDesigner

nabinchha · 2026-03-18T20:20:41Z

📋 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

Render LLM conversation traces (produced by `with_trace != TraceType.NONE`) as readable conversation flows in `display_sample_record()`. Two backends: Rich terminal panels (styled by role) and Jupyter HTML block diagrams. - New module `trace_renderer.py` with `TraceMessage` TypedDict and `TraceRenderer` class (`render_rich`, `render_notebook_html`) - `include_traces` parameter on both mixin and standalone function (defaults to True, opt out with `include_traces=False`) - Traces shown after Generated Columns table, before images - Unit tests for various trace shapes and integration tests Made-with: Cursor

Extract is_notebook_environment() utility to replace scattered get_ipython() try/except blocks. Improve Rich trace readability with better colors, separators, text folding, and dedented content. Match HTML trace font to Rich monospace output. Move index label to top of display and reduce inter-table spacing. Made-with: Cursor

greptile-apps · 2026-03-18T20:33:31Z

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 in both Rich terminal and Jupyter HTML formats. The implementation is well-structured with good test coverage, and the two P0/P1 concerns from prior rounds (double-escaping of func_name and the trace_renderer NameError guard) have been correctly resolved.

One remaining issue:

is_notebook_environment() in misc.py compares shell.__class__.__name__ against "google.colab._shell", which is a module path — not a class name. In Google Colab the class name is "Shell", so the Colab branch never matches and is_notebook_environment() silently returns False on Colab, causing users to see Rich terminal output instead of the intended HTML rendering.

Other notable changes:

include_traces=True parameter added to display_sample_record and WithRecordSamplerMixin with clean pass-through wiring
Record index label moved to the top of the output (both terminal and notebook paths)
Padding defaults normalized across table renders ((1, 0, 1, 0) → default)
_display_notebook_html helper extracted to avoid repeating IPython.display import logic

Confidence Score: 4/5

Safe to merge after fixing the __class__.__name__ vs module path mismatch in the Colab detection branch.
The core rendering logic and HTML escaping are correct, all prior review concerns are addressed, and the test suite is thorough. The only remaining concrete bug is the Google Colab environment detection always returning False (module path compared against class name), which is a targeted one-line fix.
packages/data-designer-config/src/data_designer/config/utils/misc.py — the is_notebook_environment() Colab branch

Important Files Changed

Filename	Overview
packages/data-designer-config/src/data_designer/config/utils/misc.py	Adds `is_notebook_environment()` helper; the Google Colab detection branch compares `__class__.__name__` against a module path string ("google.colab._shell") instead of the class name ("Shell"), so Colab is never detected.
packages/data-designer-config/src/data_designer/config/utils/trace_renderer.py	New module implementing Rich and HTML rendering for LLM trace conversations. HTML escaping is handled correctly throughout (double-escape fix from prior review applied), TypedDict schema is well-typed, edge cases like empty traces and malformed JSON are handled.
packages/data-designer-config/src/data_designer/config/utils/visualization.py	Integrates trace rendering into `display_sample_record`; the NameError guard fix from prior review is applied, Rich vs HTML rendering paths are logically correct, and the record-index display is moved to the top as intended.
packages/data-designer-config/tests/config/utils/test_trace_renderer.py	Comprehensive 338-line test suite covering Rich rendering, HTML rendering, edge cases (empty traces, malformed JSON, missing IPython), and integration with `display_sample_record` including `include_traces=False` and `save_path` scenarios.

Sequence Diagram

sequenceDiagram
    participant Caller
    participant display_sample_record
    participant is_notebook_environment
    participant TraceRenderer
    participant IPython

    Caller->>display_sample_record: record, config_builder, include_traces=True
    display_sample_record->>is_notebook_environment: check environment
    is_notebook_environment-->>display_sample_record: in_notebook (bool)

    display_sample_record->>display_sample_record: build render_list (Rich panels)

    alt include_traces=True
        loop each LLM column → side_effect trace columns
            display_sample_record->>display_sample_record: collect traces_to_display_later
            alt not in_notebook OR save_path set
                display_sample_record->>TraceRenderer: render_rich(traces, col_name)
                TraceRenderer-->>display_sample_record: Rich Panel
                display_sample_record->>display_sample_record: append to render_list
            end
        end
    end

    display_sample_record->>display_sample_record: print Rich render_list to console/file

    alt in_notebook AND include_traces AND traces exist
        loop each (col_name, trace_data)
            display_sample_record->>TraceRenderer: render_notebook_html(traces, col_name)
            TraceRenderer->>IPython: display(HTML(container))
        end
    end

Comments Outside Diff (1)

packages/data-designer-config/src/data_designer/config/utils/misc.py, line 13 (link)

Google Colab detection always fails

shell.__class__.__name__ returns the class name (e.g. "Shell"), not the full module path. Comparing it against "google.colab._shell" — which is a module path, not a class name — will never match in a Colab runtime. As a result, is_notebook_environment() always returns False on Colab, and users will silently get Rich terminal output instead of HTML rendering.

The canonical approach is to check the class's module:

Prompt To Fix With AI

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

Comment:
**Google Colab detection always fails**

`shell.__class__.__name__` returns the class **name** (e.g. `"Shell"`), not the full module path. Comparing it against `"google.colab._shell"` — which is a module path, not a class name — will never match in a Colab runtime. As a result, `is_notebook_environment()` always returns `False` on Colab, and users will silently get Rich terminal output instead of HTML rendering.

The canonical approach is to check the class's module:



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

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/misc.py
Line: 13

Comment:
**Google Colab detection always fails**

`shell.__class__.__name__` returns the class **name** (e.g. `"Shell"`), not the full module path. Comparing it against `"google.colab._shell"` — which is a module path, not a class name — will never match in a Colab runtime. As a result, `is_notebook_environment()` always returns `False` on Colab, and users will silently get Rich terminal output instead of HTML rendering.

The canonical approach is to check the class's module:

```suggestion
        return shell is not None and (
            shell.__class__.__name__ == "ZMQInteractiveShell"
            or shell.__class__.__module__ == "google.colab._shell"
        )
```

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

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into nmulepati/feat/..." | Re-trigger Greptile}

- Remove double html.escape() on func_name in render_notebook_html; _build_html_block already escapes the title - Guard notebook trace display with include_traces to prevent potential NameError if trace_renderer is not instantiated - Improve is_notebook_environment() to check shell class name (ZMQInteractiveShell / google.colab._shell) instead of just get_ipython() existence, avoiding false positives in IPython terminals Made-with: Cursor

johnnygreco

Review comments on trace visualization PR.

johnnygreco

A few more review comments.

johnnygreco

One more.

- Fix turn count: count assistant messages with tool calls instead of individual tool-call IDs so parallel calls report 1 turn, not N - Include Rich trace panels in render_list when save_path is set, even in notebook mode, so saved HTML files contain trace data - Move _LLM_COLUMN_TYPES to module level - Revert pad_console_element default padding to (1, 0, 1, 0) - Initialize trace_renderer unconditionally to prevent potential NameError - Consolidate IPython.display imports into _display_notebook_html helper - Restrict bg_color to RoleHtmlColor Literal type - Precise type annotation for sections list in render_rich - Remove unnecessary # noqa: F821 suppression - Move test imports to module level, remove __import__ pattern Made-with: Cursor

nabinchha added 3 commits March 13, 2026 16:44

Merge branch 'main' into nmulepati/feat/396-trace-visualization

1eec0d5

nabinchha requested a review from a team as a code owner March 18, 2026 20:20

move required to import from type_extensions

65e346b

greptile-apps Bot reviewed Mar 18, 2026

View reviewed changes

Comment thread packages/data-designer-config/src/data_designer/config/utils/trace_renderer.py Outdated

Comment thread packages/data-designer-config/src/data_designer/config/utils/visualization.py

Comment thread packages/data-designer-config/src/data_designer/config/utils/misc.py Outdated

nabinchha added 2 commits March 18, 2026 14:40

Merge branch 'main' into nmulepati/feat/396-trace-visualization

30f3913

johnnygreco reviewed Mar 20, 2026

View reviewed changes

Comment thread packages/data-designer-config/src/data_designer/config/utils/visualization.py

Comment thread packages/data-designer-config/src/data_designer/config/utils/visualization.py Outdated

Comment thread packages/data-designer-config/src/data_designer/config/utils/trace_renderer.py Outdated

johnnygreco reviewed Mar 20, 2026

View reviewed changes

Comment thread packages/data-designer-config/src/data_designer/config/utils/trace_renderer.py

nabinchha added 2 commits March 23, 2026 09:49

Merge branch 'main' into nmulepati/feat/396-trace-visualization

843171e

johnnygreco approved these changes Mar 23, 2026

View reviewed changes

Merge branch 'main' into nmulepati/feat/396-trace-visualization

20b3f45

nabinchha merged commit 63c294e into main Mar 23, 2026
47 checks passed

nabinchha deleted the nmulepati/feat/396-trace-visualization branch March 23, 2026 22:57

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: add trace visualization to display_sample_record (#396) #438

feat: add trace visualization to display_sample_record (#396) #438
nabinchha merged 9 commits into
mainfrom
nmulepati/feat/396-trace-visualization

nabinchha commented Mar 18, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented Mar 18, 2026 •

edited

Loading

Confidence Score: 4/5

Sequence Diagram

Comments Outside Diff (1)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

johnnygreco left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

johnnygreco left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

johnnygreco left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

nabinchha commented Mar 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📋 Summary

🔄 Changes

✨ Added

🔧 Changed

🔍 Attention Areas

In ipynb notebook

In console

Uh oh!

greptile-apps Bot commented Mar 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Sequence Diagram

Comments Outside Diff (1)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

johnnygreco left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

johnnygreco left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

johnnygreco left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

nabinchha commented Mar 18, 2026 •

edited

Loading

greptile-apps Bot commented Mar 18, 2026 •

edited

Loading