feat: Add TraceType enum for granular trace control (#284)

Author: eric-tramel

docs/assets/recipes/mcp_and_tooluse/basic_mcp.py

@@ -132,7 +132,7 @@ def build_config(model_alias: str, provider_name: str) -> dd.DataDesignerConfigB
             ),
             system_prompt="You must call the get_fact tool before answering. Only use information from tool results.",
             tool_alias="basic-tools",
-            with_trace=True,
+            with_trace=dd.TraceType.ALL_MESSAGES,
         )
     )
 
@@ -163,7 +163,7 @@ def build_config(model_alias: str, provider_name: str) -> dd.DataDesignerConfigB
             ),
             system_prompt="You must call the add_numbers tool to perform the calculation. Report the exact result.",
             tool_alias="basic-tools",
-            with_trace=True,
+            with_trace=dd.TraceType.ALL_MESSAGES,
         )
     )
 

docs/assets/recipes/mcp_and_tooluse/pdf_qa.py

@@ -312,7 +312,7 @@ def build_config(model_alias: str, provider_name: str) -> dd.DataDesignerConfigB
             ),
             output_format=TopicList,
             tool_alias="doc-search",
-            with_trace=True,  # Enable trace to capture tool call history
+            with_trace=dd.TraceType.ALL_MESSAGES,  # Enable trace to capture tool call history
         )
     )
 
@@ -341,7 +341,7 @@ def build_config(model_alias: str, provider_name: str) -> dd.DataDesignerConfigB
             ),
             output_format=QAPair,
             tool_alias="doc-search",
-            with_trace=True,  # Enable trace to capture tool call history
+            with_trace=dd.TraceType.ALL_MESSAGES,  # Enable trace to capture tool call history
         )
     )
 

docs/concepts/columns.md

@@ -39,7 +39,7 @@ LLM-Text columns generate natural language text: product descriptions, customer
 Use **Jinja2 templating** in prompts to reference other columns. Data Designer automatically manages dependencies and injects the referenced column values into the prompt.
 
 !!! note "Generation Traces"
-    LLM columns can optionally capture a full message trace in a separate `{column_name}__trace` column. Enable traces per-column via `with_trace=True` on the column config, or globally for all columns via `RunConfig(debug_override_save_all_column_traces=True)`. The trace includes the ordered message history for the final generation attempt (system/user/assistant/tool calls/tool results), and may include model reasoning fields when the provider exposes them.
+    LLM columns can optionally capture message traces in a separate `{column_name}__trace` column. Set `with_trace` on the column config to control what's captured: `TraceType.NONE` (default, no trace), `TraceType.LAST_MESSAGE` (final assistant message only), or `TraceType.ALL_MESSAGES` (full conversation history). Override globally via `RunConfig(debug_trace_override=TraceType.ALL_MESSAGES)`. The trace includes the ordered message history for the final generation attempt (system/user/assistant/tool calls/tool results), and may include model reasoning fields when the provider exposes them.
 
 !!! tip "Tool Use in LLM Columns"
     LLM columns can invoke external tools during generation via MCP (Model Context Protocol). Enable tools by setting `tool_alias` to reference a configured `ToolConfig`:
@@ -50,7 +50,7 @@ Use **Jinja2 templating** in prompts to reference other columns. Data Designer a
         model_alias="nvidia-text",
         prompt="Search for information and answer: {{ question }}",
         tool_alias="search-tools",  # References a ToolConfig
-        with_trace=True,  # Capture tool call history
+        with_trace=dd.TraceType.ALL_MESSAGES,  # Capture tool call history
     )
     ```
 
@@ -162,6 +162,6 @@ You read this property for introspection but never set it—always computed from
 
 ### `side_effect_columns`
 
-Computed property listing columns created implicitly alongside the primary column. Currently, only LLM columns produce side effects (trace columns like `{name}__trace` when `with_trace=True` is set on the column or `debug_override_save_all_column_traces` is enabled globally).
+Computed property listing columns created implicitly alongside the primary column. Currently, only LLM columns produce side effects (trace columns like `{name}__trace` when `with_trace` is not `TraceType.NONE` on the column or `debug_trace_override` is set globally).
 
 For detailed information on each column type, refer to the [column configuration code reference](../code_reference/column_configs.md).

docs/concepts/mcp/enabling-tools.md

@@ -90,7 +90,7 @@ builder.add_column(
         prompt="Use the available tools to research and answer: {{ question }}",
         model_alias="nvidia-text",
         tool_alias="my-tools",  # Enable tools
-        with_trace=True,        # Capture tool call history
+        with_trace=dd.TraceType.ALL_MESSAGES,  # Capture tool call history
     )
 )
 

docs/concepts/traces.md

@@ -1,6 +1,6 @@
 # Message Traces
 
-Traces capture the full conversation history during LLM generation, including system prompts, user prompts, model reasoning, tool calls, tool results, and the final response. This visibility is essential for understanding model behavior, debugging generation issues, and iterating on prompts.
+Traces capture the conversation history during LLM generation, including system prompts, user prompts, model reasoning, tool calls, tool results, and the final response. This visibility is essential for understanding model behavior, debugging generation issues, and iterating on prompts.
 
 Traces are also useful in certain scenarios as the target output of the workflow, e.g. producing an SFT dataset for fine-tuning tool-use capability, for instance.
 
@@ -19,39 +19,74 @@ When generating content with LLM columns, you often need to understand what happ
 
 Traces provide this visibility by capturing the ordered message history for each generation, including any multi-turn conversations that occur during tool use or retry scenarios.
 
+## Trace Types
+
+Data Designer supports three trace modes via the `TraceType` enum:
+
+| 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) |
+
 ## Enabling Traces
 
 ### Per-Column (Recommended)
 
-Enable `with_trace=True` on specific LLM columns:
+Set `with_trace` on specific LLM columns:
 
 ```python
 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=True,  # Enable trace for this column
+        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
 
-Enable traces for ALL LLM columns (useful during development):
+Override trace settings for ALL LLM columns (useful during development):
 
 ```python
 import data_designer.config as dd
 from data_designer.interface import DataDesigner
 
 data_designer = DataDesigner()
+
+# Enable full traces for all columns
+data_designer.set_run_config(
+    dd.RunConfig(debug_trace_override=dd.TraceType.ALL_MESSAGES)
+)
+
+# Or capture only last messages for all columns
 data_designer.set_run_config(
-    dd.RunConfig(debug_override_save_all_column_traces=True)
+    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.
+
 ## Trace Column Naming
 
 When enabled, LLM columns produce an additional side-effect column:
@@ -161,4 +196,4 @@ When an assistant message includes tool calls:
 ## See Also
 
 - **[Safety and Limits](mcp/safety-and-limits.md)**: Understand turn limits and timeout behavior
-- **[Run Config](../code_reference/run_config.md)**: Runtime options including `debug_override_save_all_column_traces`
+- **[Run Config](../code_reference/run_config.md)**: Runtime options including `debug_trace_override`

packages/data-designer-config/src/data_designer/config/__init__.py

@@ -74,6 +74,7 @@
 )
 from data_designer.config.utils.code_lang import CodeLang
 from data_designer.config.utils.info import InfoType
+from data_designer.config.utils.trace_type import TraceType
 from data_designer.config.validator_params import (
     CodeValidatorParams,
     LocalCallableValidatorParams,
@@ -144,6 +145,7 @@ def get_config_exports() -> list[str]:
         SeedDatasetColumnConfig.__name__,
         SubcategorySamplerParams.__name__,
         TimeDeltaSamplerParams.__name__,
+        TraceType.__name__,
         UniformDistribution.__name__,
         UniformDistributionParams.__name__,
         UniformSamplerParams.__name__,

packages/data-designer-config/src/data_designer/config/column_configs.py

@@ -16,6 +16,7 @@
 from data_designer.config.utils.code_lang import CodeLang
 from data_designer.config.utils.constants import TRACE_COLUMN_POSTFIX
 from data_designer.config.utils.misc import assert_valid_jinja2_template, extract_keywords_from_jinja2_template
+from data_designer.config.utils.trace_type import TraceType
 from data_designer.config.validator_params import ValidatorParamsT, ValidatorType
 
 
@@ -162,10 +163,12 @@ class LLMTextColumnConfig(SingleColumnConfig):
         tool_alias: Optional alias of the tool configuration to use for MCP tool calls.
             Must match a tool alias defined when initializing the DataDesignerConfigBuilder.
             When provided, the model may call permitted tools during generation.
-        with_trace: If True, creates a `{column_name}__trace` column containing the full
-            ordered message history (system/user/assistant/tool) for the generation.
-            Can be overridden globally via `RunConfig.debug_override_save_all_column_traces`.
-            Defaults to False.
+        with_trace: Specifies what trace information to capture in a `{column_name}__trace`
+            column. Options are:
+            - `TraceType.NONE` (default): No trace is captured.
+            - `TraceType.LAST_MESSAGE`: Only the final assistant message is captured.
+            - `TraceType.ALL_MESSAGES`: Full conversation history (system/user/assistant/tool).
+            Can be overridden globally via `RunConfig.debug_trace_override`.
         column_type: Discriminator field, always "llm-text" for this configuration type.
     """
 
@@ -174,7 +177,7 @@ class LLMTextColumnConfig(SingleColumnConfig):
     system_prompt: str | None = None
     multi_modal_context: list[ImageContext] | None = None
     tool_alias: str | None = None
-    with_trace: bool = False
+    with_trace: TraceType = TraceType.NONE
     column_type: Literal["llm-text"] = "llm-text"
 
     @staticmethod
@@ -197,8 +200,8 @@ def required_columns(self) -> list[str]:
     def side_effect_columns(self) -> list[str]:
         """Returns the trace column, which may be generated alongside the main column.
 
-        Traces are generated when `with_trace=True` on the column config or
-        when `RunConfig.debug_override_save_all_column_traces=True` globally.
+        Traces are generated when `with_trace` is not `TraceType.NONE` on the column config
+        or when `RunConfig.debug_trace_override` is set globally.
 
         Returns:
             List containing the trace column name.

packages/data-designer-config/src/data_designer/config/run_config.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations
@@ -7,6 +7,7 @@
 from typing_extensions import Self
 
 from data_designer.config.base import ConfigBase
+from data_designer.config.utils.trace_type import TraceType
 
 
 class RunConfig(ConfigBase):
@@ -33,10 +34,13 @@ class RunConfig(ConfigBase):
         max_conversation_correction_steps: Maximum number of correction rounds permitted within a
             single conversation when generation tasks call `ModelFacade.generate(...)`. Must be >= 0.
             Default is 0.
-        debug_override_save_all_column_traces: If True, overrides per-column `with_trace` settings
-            and includes `__trace` columns for ALL LLM generations, containing the full ordered
-            message history (system/user/assistant/tool) for the final generation attempt.
-            Useful for debugging. Default is False.
+        debug_trace_override: If set, overrides per-column `with_trace` settings for ALL LLM
+            generations. Options are:
+            - `None` (default): Use per-column `with_trace` settings.
+            - `TraceType.NONE`: Disable all traces, ignoring per-column settings.
+            - `TraceType.LAST_MESSAGE`: Capture only the final assistant message for all columns.
+            - `TraceType.ALL_MESSAGES`: Capture full conversation history for all columns.
+            Useful for debugging or bulk trace collection.
     """
 
     disable_early_shutdown: bool = False
@@ -46,7 +50,7 @@ class RunConfig(ConfigBase):
     non_inference_max_parallel_workers: int = Field(default=4, ge=1)
     max_conversation_restarts: int = Field(default=5, ge=0)
     max_conversation_correction_steps: int = Field(default=0, ge=0)
-    debug_override_save_all_column_traces: bool = False
+    debug_trace_override: TraceType | None = None
 
     @model_validator(mode="after")
     def normalize_shutdown_settings(self) -> Self:

packages/data-designer-config/src/data_designer/config/utils/trace_type.py

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from data_designer.config.utils.type_helpers import StrEnum
+
+
+class TraceType(StrEnum):
+    """Specifies the type of reasoning trace to capture for LLM columns.
+
+    Traces capture the conversation history during LLM generation, which is
+    useful for debugging, analysis, and understanding model behavior.
+
+    Attributes:
+        NONE: No trace is captured. This is the default.
+        LAST_MESSAGE: Only the final assistant message is captured.
+        ALL_MESSAGES: The full conversation history (system/user/assistant/tool)
+            is captured.
+    """
+
+    NONE = "none"
+    LAST_MESSAGE = "last_message"
+    ALL_MESSAGES = "all_messages"

packages/data-designer-config/tests/config/test_columns.py

@@ -36,6 +36,7 @@
 )
 from data_designer.config.utils.code_lang import CodeLang
 from data_designer.config.utils.errors import UserJinjaTemplateSyntaxError
+from data_designer.config.utils.trace_type import TraceType
 from data_designer.config.validator_params import CodeValidatorParams
 
 stub_prompt = "test_prompt {{some_column}}"
@@ -86,6 +87,7 @@ def test_llm_text_column_config():
     assert llm_text_column_config.column_type == DataDesignerColumnType.LLM_TEXT
     assert set(llm_text_column_config.required_columns) == {"some_column", "some_other_column"}
     assert llm_text_column_config.side_effect_columns == ["test_llm_text__trace"]
+    assert llm_text_column_config.with_trace == TraceType.NONE
 
     # invalid prompt
     with pytest.raises(
@@ -110,6 +112,35 @@ def test_llm_text_column_config():
         )
 
 
+def test_llm_text_column_config_with_trace_serialization() -> None:
+    """Test that with_trace field serializes and deserializes correctly."""
+    config = LLMTextColumnConfig(
+        name="test_llm_text",
+        prompt=stub_prompt,
+        model_alias=stub_model_alias,
+        with_trace=TraceType.ALL_MESSAGES,
+    )
+    assert config.with_trace == TraceType.ALL_MESSAGES
+
+    # Serialize
+    serialized = config.model_dump()
+    assert serialized["with_trace"] == "all_messages"
+
+    # Deserialize
+    deserialized = LLMTextColumnConfig(**serialized)
+    assert deserialized.with_trace == TraceType.ALL_MESSAGES
+
+    # Test with LAST_MESSAGE
+    config_last = LLMTextColumnConfig(
+        name="test_llm_text",
+        prompt=stub_prompt,
+        model_alias=stub_model_alias,
+        with_trace=TraceType.LAST_MESSAGE,
+    )
+    assert config_last.with_trace == TraceType.LAST_MESSAGE
+    assert config_last.model_dump()["with_trace"] == "last_message"
+
+
 def test_llm_code_column_config():
     llm_code_column_config = LLMCodeColumnConfig(
         name="test_llm_code",