feat: Add TraceType enum for granular trace control

[eric-tramel]

Summary

This PR converts the boolean with_trace field in LLM column configs and debug_override_save_all_column_traces in RunConfig to use a new TraceType enum, providing more granular control over what trace information is captured.

New TraceType Enum

The TraceType enum provides three modes for trace capture:

TraceType	Description
TraceType.NONE	No trace captured (default)
TraceType.LAST_MESSAGE	Only the final assistant message is captured
TraceType.ALL_MESSAGES	Full conversation history (system/user/assistant/tool)

Breaking Changes

• with_trace field type changed: The with_trace field on LLM column configs (LLMTextColumnConfig, LLMCodeColumnConfig, LLMStructuredColumnConfig, LLMJudgeColumnConfig) now accepts TraceType instead of bool.

  • with_trace=True → with_trace=TraceType.ALL_MESSAGES
  • with_trace=False → with_trace=TraceType.NONE (default)

• debug_override_save_all_column_traces renamed: The RunConfig field is now debug_trace_override and accepts TraceType | None.

  • debug_override_save_all_column_traces=True → debug_trace_override=TraceType.ALL_MESSAGES
  • debug_override_save_all_column_traces=False → debug_trace_override=None (default, uses per-column settings)

Usage

Per-Column Trace Configuration

import data_designer.config as dd

# Capture full conversation history
builder.add_column(
    dd.LLMTextColumnConfig(
        name="answer",
        prompt="Answer: {{ question }}",
        model_alias="nvidia-text",
        with_trace=dd.TraceType.ALL_MESSAGES,  # Full trace
    )
)

# Capture only the final assistant response
builder.add_column(
    dd.LLMTextColumnConfig(
        name="summary",
        prompt="Summarize: {{ text }}",
        model_alias="nvidia-text",
        with_trace=dd.TraceType.LAST_MESSAGE,  # Just the final response
    )
)

Global Debug Override

import data_designer.config as dd
from data_designer.interface import DataDesigner

data_designer = DataDesigner()

# Enable full traces for all columns (overrides per-column settings)
data_designer.set_run_config(
    dd.RunConfig(debug_trace_override=dd.TraceType.ALL_MESSAGES)
)

# Capture only last messages for all columns
data_designer.set_run_config(
    dd.RunConfig(debug_trace_override=dd.TraceType.LAST_MESSAGE)
)

# Disable all traces (overrides per-column settings)
data_designer.set_run_config(
    dd.RunConfig(debug_trace_override=dd.TraceType.NONE)
)

When debug_trace_override is set (not None), it takes precedence over per-column with_trace settings.

Changes

New Files

• packages/data-designer-config/src/data_designer/config/utils/trace_type.py - New TraceType enum

Modified Files

• packages/data-designer-config/src/data_designer/config/column_configs.py - Updated with_trace field type
• packages/data-designer-config/src/data_designer/config/run_config.py - Renamed and retyped debug override field
• packages/data-designer-config/src/data_designer/config/__init__.py - Export TraceType
• packages/data-designer-engine/src/data_designer/engine/column_generators/generators/llm_completion.py - Updated trace logic for three modes
• Documentation files updated with new usage patterns
• Test files updated with comprehensive coverage for all trace modes

Test plan

• All existing tests pass (1797 tests)
• New tests added for TraceType.LAST_MESSAGE behavior
• New tests added for override precedence logic
• New tests added for with_trace field serialization
• Linting passes
• License headers updated

🤖 Generated with Claude Code

[greptile-apps]

Greptile Overview

Greptile Summary

This PR introduces a TraceType enum to replace the boolean with_trace field, enabling more granular control over trace capture for LLM column generations.

Key Changes:

• New TraceType enum with three modes: NONE (default), LAST_MESSAGE, and ALL_MESSAGES
• Breaking change: with_trace field type changed from bool to TraceType on all LLM column configs
• Breaking change: debug_override_save_all_column_traces renamed to debug_trace_override with type TraceType | None
• Implementation correctly handles override precedence (debug_trace_override takes precedence over per-column settings)
• LAST_MESSAGE mode extracts only the final assistant message from the trace using reverse iteration
• Comprehensive test coverage for all three modes, override precedence, and edge cases
• All documentation, examples, and E2E tests updated consistently

Implementation Quality:

• Clean enum implementation using StrEnum for JSON serialization
• Backward compatible migration path is clear in PR description
• All 1797 existing tests pass
• New tests cover serialization, all three trace modes, override behavior, and edge case (no assistant message)

Confidence Score: 5/5

• This PR is safe to merge with high confidence
• Well-implemented breaking change with comprehensive test coverage (1797 existing tests pass, new tests for all trace modes added), clear migration path documented, consistent updates across all files, and follows repository patterns (StrEnum, type annotations, docstrings)
• No files require special attention

Important Files Changed

Filename	Overview
packages/data-designer-config/src/data_designer/config/utils/trace_type.py	Added new TraceType enum with three values: NONE, LAST_MESSAGE, and ALL_MESSAGES for granular trace control
packages/data-designer-config/src/data_designer/config/column_configs.py	Changed with_trace field from bool to TraceType enum with default TraceType.NONE, updated docstrings
packages/data-designer-config/src/data_designer/config/run_config.py	Renamed debug_override_save_all_column_traces to debug_trace_override, changed type from bool to `TraceType
packages/data-designer-engine/src/data_designer/engine/column_generators/generators/llm_completion.py	Implemented three-mode trace logic: saves all messages, only last assistant message, or no trace based on effective TraceType
packages/data-designer-engine/tests/engine/column_generators/generators/test_llm_completion_generators.py	Added extensive tests for all trace modes, override precedence, and edge cases (e.g., no assistant message in trace)

Sequence Diagram

sequenceDiagram
    participant User
    participant ConfigBuilder
    participant LLMColumnConfig
    participant RunConfig
    participant LLMCompletionGenerator
    participant Model
    
    User->>ConfigBuilder: add_column(LLMTextColumnConfig)
    User->>LLMColumnConfig: with_trace=TraceType.LAST_MESSAGE
    LLMColumnConfig-->>ConfigBuilder: Column config created
    
    User->>RunConfig: debug_trace_override=TraceType.ALL_MESSAGES
    RunConfig-->>User: Override configured
    
    User->>LLMCompletionGenerator: generate(data)
    
    LLMCompletionGenerator->>RunConfig: Check debug_trace_override
    RunConfig-->>LLMCompletionGenerator: TraceType.ALL_MESSAGES (override)
    
    Note over LLMCompletionGenerator: Effective trace type = ALL_MESSAGES<br/>(override takes precedence)
    
    LLMCompletionGenerator->>Model: generate(prompt, system_prompt, ...)
    Model-->>LLMCompletionGenerator: (response, trace)
    
    alt effective_trace_type == ALL_MESSAGES
        LLMCompletionGenerator->>LLMCompletionGenerator: Save all messages to trace column
    else effective_trace_type == LAST_MESSAGE
        LLMCompletionGenerator->>LLMCompletionGenerator: Find last assistant message
        LLMCompletionGenerator->>LLMCompletionGenerator: Save only last assistant to trace column
    else effective_trace_type == NONE
        LLMCompletionGenerator->>LLMCompletionGenerator: No trace column created
    end
    
    LLMCompletionGenerator-->>User: Generated data with trace column

Loading

[greptile-apps]

_{5 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[nabinchha]

tiniest of nit, but lgtm! Thanks @eric-tramel!

[greptile-apps]

_{5 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}