data-designer-engine 0.4.0rc1__tar.gz → 0.4.0rc3__tar.gz

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-designer-engine
-Version: 0.4.0rc1
+Version: 0.4.0rc3
 Summary: Generation engine for DataDesigner synthetic data generation
 License-Expression: Apache-2.0
 Classifier: Development Status :: 4 - Beta

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc3}/src/data_designer/engine/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.4.0rc1'
-__version_tuple__ = version_tuple = (0, 4, 0, 'rc1')
+__version__ = version = '0.4.0rc3'
+__version_tuple__ = version_tuple = (0, 4, 0, 'rc3')
 
 __commit_id__ = commit_id = None

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc3}/src/data_designer/engine/column_generators/generators/llm_completion.py

@@ -12,7 +12,7 @@ from data_designer.config.column_configs import (
     LLMStructuredColumnConfig,
     LLMTextColumnConfig,
 )
-from data_designer.config.utils.constants import REASONING_TRACE_COLUMN_POSTFIX
+from data_designer.config.utils.constants import TRACE_COLUMN_POSTFIX
 from data_designer.engine.column_generators.generators.base import ColumnGeneratorWithModel, GenerationStrategy
 from data_designer.engine.column_generators.utils.prompt_renderer import (
     PromptType,
@@ -62,11 +62,11 @@ class ColumnGeneratorWithModelChatCompletion(ColumnGeneratorWithModel[TaskConfig
 
         multi_modal_context = None
         if self.config.multi_modal_context is not None and len(self.config.multi_modal_context) > 0:
-            multi_modal_context = [
-                context.get_context(deserialized_record) for context in self.config.multi_modal_context
-            ]
+            multi_modal_context = []
+            for context in self.config.multi_modal_context:
+                multi_modal_context.extend(context.get_contexts(deserialized_record))
 
-        response, reasoning_trace = self.model.generate(
+        response, trace = self.model.generate(
             prompt=self.prompt_renderer.render(
                 record=deserialized_record,
                 prompt_template=self.config.prompt,
@@ -87,8 +87,11 @@ class ColumnGeneratorWithModelChatCompletion(ColumnGeneratorWithModel[TaskConfig
         serialized_output = self.response_recipe.serialize_output(response)
         data[self.config.name] = self._process_serialized_output(serialized_output)
 
-        if reasoning_trace:
-            data[self.config.name + REASONING_TRACE_COLUMN_POSTFIX] = reasoning_trace
+        should_save_trace = (
+            self.config.with_trace or self.resource_provider.run_config.debug_override_save_all_column_traces
+        )
+        if should_save_trace:
+            data[self.config.name + TRACE_COLUMN_POSTFIX] = [message.to_dict() for message in trace]
 
         return data
 

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc3}/src/data_designer/engine/dataset_builders/column_wise_builder.py

@@ -34,6 +34,7 @@ from data_designer.engine.dataset_builders.multi_column_configs import MultiColu
 from data_designer.engine.dataset_builders.utils.concurrency import ConcurrentThreadExecutor
 from data_designer.engine.dataset_builders.utils.config_compiler import compile_dataset_builder_column_configs
 from data_designer.engine.dataset_builders.utils.dataset_batch_manager import DatasetBatchManager
+from data_designer.engine.dataset_builders.utils.progress_tracker import ProgressTracker
 from data_designer.engine.models.telemetry import InferenceEvent, NemoSourceEnum, TaskStatusEnum, TelemetryHandler
 from data_designer.engine.processing.processors.base import Processor
 from data_designer.engine.processing.processors.drop_columns import DropColumnsProcessor
@@ -221,16 +222,18 @@ class ColumnWiseDatasetBuilder:
                 "generator so concurrency through threads is not supported."
             )
 
-        logger.info(
-            f"🐙 Processing {generator.config.column_type} column '{generator.config.name}' "
-            f"with {max_workers} concurrent workers"
+        progress_tracker = ProgressTracker(
+            total_records=self.batch_manager.num_records_batch,
+            label=f"{generator.config.column_type} column '{generator.config.name}'",
         )
+        progress_tracker.log_start(max_workers)
+
         settings = self._resource_provider.run_config
         with ConcurrentThreadExecutor(
             max_workers=max_workers,
             column_name=generator.config.name,
-            result_callback=self._worker_result_callback,
-            error_callback=self._worker_error_callback,
+            result_callback=self._make_result_callback(progress_tracker),
+            error_callback=self._make_error_callback(progress_tracker),
             shutdown_error_rate=settings.shutdown_error_rate,
             shutdown_error_window=settings.shutdown_error_window,
             disable_early_shutdown=settings.disable_early_shutdown,
@@ -238,10 +241,26 @@ class ColumnWiseDatasetBuilder:
             for i, record in self.batch_manager.iter_current_batch():
                 executor.submit(lambda record: generator.generate(record), record, context={"index": i})
 
+        progress_tracker.log_final()
+
         if len(self._records_to_drop) > 0:
             self.batch_manager.drop_records(self._records_to_drop)
             self._records_to_drop.clear()
 
+    def _make_result_callback(self, progress_tracker: ProgressTracker) -> Callable[[dict], None]:
+        def callback(result: dict, *, context: dict | None = None) -> None:
+            self._worker_result_callback(result, context=context)
+            progress_tracker.record_success()
+
+        return callback
+
+    def _make_error_callback(self, progress_tracker: ProgressTracker) -> Callable[[Exception], None]:
+        def callback(exc: Exception, *, context: dict | None = None) -> None:
+            self._worker_error_callback(exc, context=context)
+            progress_tracker.record_failure()
+
+        return callback
+
     def _write_processed_batch(self, dataframe: pd.DataFrame) -> None:
         self.batch_manager.update_records(dataframe.to_dict(orient="records"))
         self.batch_manager.write()

data_designer_engine-0.4.0rc3/src/data_designer/engine/dataset_builders/utils/progress_tracker.py

@@ -0,0 +1,122 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+import time
+from threading import Lock
+
+from data_designer.logging import RandomEmoji
+
+logger = logging.getLogger(__name__)
+
+
+class ProgressTracker:
+    """
+    Thread-safe progress tracker for monitoring concurrent task completion.
+
+    Tracks completed, successful, and failed task counts and logs progress
+    at configurable intervals. Designed for use with ConcurrentThreadExecutor
+    to provide visibility into long-running batch operations.
+
+    Example usage:
+        tracker = ProgressTracker(total_records=100, label="LLM_TEXT column 'response'")
+        tracker.log_start(max_workers=8)
+
+        # In callbacks from ConcurrentThreadExecutor:
+        tracker.record_success()  # or tracker.record_failure()
+
+        # After executor completes:
+        tracker.log_final()
+    """
+
+    def __init__(self, total_records: int, label: str, log_interval_percent: int = 10):
+        """
+        Initialize the progress tracker.
+
+        Args:
+            total_records: Total number of records to process.
+            label: Human-readable label for log messages (e.g., "LLM_TEXT column 'response'").
+            log_interval_percent: How often to log progress as a percentage (default 10%).
+        """
+        self.total_records = total_records
+        self.label = label
+
+        self.completed = 0
+        self.success = 0
+        self.failed = 0
+
+        interval_fraction = max(1, log_interval_percent) / 100.0
+        self.log_interval = max(1, int(total_records * interval_fraction)) if total_records > 0 else 1
+        self.next_log_at = self.log_interval
+
+        self.start_time = time.perf_counter()
+        self.lock = Lock()
+        self._random_emoji = RandomEmoji()
+
+    def log_start(self, max_workers: int) -> None:
+        """Log the start of processing with worker count and interval information."""
+        logger.info(
+            "🐙 Processing %s with %d concurrent workers",
+            self.label,
+            max_workers,
+        )
+        logger.info(
+            "🧭 %s will report progress every %d record(s).",
+            self.label,
+            self.log_interval,
+        )
+
+    def record_success(self) -> None:
+        """Record a successful task completion and log progress if at interval."""
+        self._record_completion(success=True)
+
+    def record_failure(self) -> None:
+        """Record a failed task completion and log progress if at interval."""
+        self._record_completion(success=False)
+
+    def log_final(self) -> None:
+        """Log final progress summary."""
+        with self.lock:
+            if self.completed > 0:
+                self._log_progress_unlocked()
+
+    def _record_completion(self, *, success: bool) -> None:
+        should_log = False
+        with self.lock:
+            self.completed += 1
+            if success:
+                self.success += 1
+            else:
+                self.failed += 1
+
+            if self.completed >= self.next_log_at and self.completed < self.total_records:
+                should_log = True
+                while self.next_log_at <= self.completed:
+                    self.next_log_at += self.log_interval
+
+        if should_log:
+            with self.lock:
+                self._log_progress_unlocked()
+
+    def _log_progress_unlocked(self) -> None:
+        """Log current progress. Must be called while holding the lock."""
+        elapsed = time.perf_counter() - self.start_time
+        rate = self.completed / elapsed if elapsed > 0 else 0.0
+        remaining = max(0, self.total_records - self.completed)
+        eta = f"{(remaining / rate):.1f}s" if rate > 0 else "unknown"
+        percent = (self.completed / self.total_records) * 100 if self.total_records else 100.0
+
+        logger.info(
+            "  |-- %s %s progress: %d/%d (%.0f%%) complete, %d ok, %d failed, %.2f rec/s, eta %s",
+            self._random_emoji.progress(percent),
+            self.label,
+            self.completed,
+            self.total_records,
+            percent,
+            self.success,
+            self.failed,
+            rate,
+            eta,
+        )

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc3}/src/data_designer/engine/models/facade.py

@@ -18,7 +18,7 @@ from data_designer.engine.models.errors import (
 from data_designer.engine.models.litellm_overrides import CustomRouter, LiteLLMRouterDefaultKwargs
 from data_designer.engine.models.parsers.errors import ParserException
 from data_designer.engine.models.usage import ModelUsageStats, RequestUsageStats, TokenUsageStats
-from data_designer.engine.models.utils import prompt_to_messages, str_to_message
+from data_designer.engine.models.utils import ChatMessage, prompt_to_messages
 from data_designer.engine.secret_resolver import SecretResolver
 from data_designer.lazy_heavy_imports import litellm
 
@@ -67,16 +67,17 @@ class ModelFacade:
         return self._usage_stats
 
     def completion(
-        self, messages: list[dict[str, str]], skip_usage_tracking: bool = False, **kwargs
+        self, messages: list[ChatMessage], skip_usage_tracking: bool = False, **kwargs
     ) -> litellm.ModelResponse:
+        message_payloads = [message.to_dict() for message in messages]
         logger.debug(
             f"Prompting model {self.model_name!r}...",
-            extra={"model": self.model_name, "messages": messages},
+            extra={"model": self.model_name, "messages": message_payloads},
         )
         response = None
         kwargs = self.consolidate_kwargs(**kwargs)
         try:
-            response = self._router.completion(model=self.model_name, messages=messages, **kwargs)
+            response = self._router.completion(model=self.model_name, messages=message_payloads, **kwargs)
             logger.debug(
                 f"Received completion from model {self.model_name!r}",
                 extra={
@@ -149,7 +150,7 @@ class ModelFacade:
         skip_usage_tracking: bool = False,
         purpose: str | None = None,
         **kwargs,
-    ) -> tuple[Any, str | None]:
+    ) -> tuple[Any, list[ChatMessage]]:
         """Generate a parsed output with correction steps.
 
         This generation call will attempt to generate an output which is
@@ -182,6 +183,12 @@ class ModelFacade:
                 It is expected to be used by the @catch_llm_exceptions decorator.
             **kwargs: Additional arguments to pass to the model.
 
+        Returns:
+            A tuple containing:
+                - The parsed output object from the parser.
+                - The full trace of ChatMessage entries in the conversation, including any
+                  corrections and reasoning traces. Callers can decide whether to store this.
+
         Raises:
             GenerationValidationFailureError: If the maximum number of retries or
                 correction steps are met and the last response failures on
@@ -190,29 +197,17 @@ class ModelFacade:
         output_obj = None
         curr_num_correction_steps = 0
         curr_num_restarts = 0
-        curr_generation_attempt = 0
-        max_generation_attempts = (max_correction_steps + 1) * (max_conversation_restarts + 1)
 
         starting_messages = prompt_to_messages(
             user_prompt=prompt, system_prompt=system_prompt, multi_modal_context=multi_modal_context
         )
-        messages = deepcopy(starting_messages)
+        messages: list[ChatMessage] = deepcopy(starting_messages)
 
         while True:
-            curr_generation_attempt += 1
-            logger.debug(
-                f"Starting generation attempt {curr_generation_attempt} of {max_generation_attempts} attempts."
-            )
-
             completion_response = self.completion(messages, skip_usage_tracking=skip_usage_tracking, **kwargs)
             response = completion_response.choices[0].message.content or ""
             reasoning_trace = getattr(completion_response.choices[0].message, "reasoning_content", None)
-
-            if reasoning_trace:
-                ## There are generally some extra newlines with how these get parsed.
-                response = response.strip()
-                reasoning_trace = reasoning_trace.strip()
-
+            messages.append(ChatMessage.as_assistant(content=response, reasoning_content=reasoning_trace or None))
             curr_num_correction_steps += 1
 
             try:
@@ -223,21 +218,23 @@ class ModelFacade:
                     raise GenerationValidationFailureError(
                         "Unsuccessful generation attempt. No retries were attempted."
                     ) from exc
+
                 if curr_num_correction_steps <= max_correction_steps:
-                    ## Add turns to loop-back errors for correction
-                    messages += [
-                        str_to_message(content=response, role="assistant"),
-                        str_to_message(content=str(get_exception_primary_cause(exc)), role="user"),
-                    ]
+                    # Add user message with error for correction
+                    messages.append(ChatMessage.as_user(content=str(get_exception_primary_cause(exc))))
+
                 elif curr_num_restarts < max_conversation_restarts:
                     curr_num_correction_steps = 0
                     curr_num_restarts += 1
                     messages = deepcopy(starting_messages)
+
                 else:
                     raise GenerationValidationFailureError(
-                        f"Unsuccessful generation attempt despite {max_generation_attempts} attempts."
+                        f"Unsuccessful generation despite {max_correction_steps} correction steps "
+                        f"and {max_conversation_restarts} conversation restarts."
                     ) from exc
-        return output_obj, reasoning_trace
+
+        return output_obj, messages
 
     def _get_litellm_deployment(self, model_config: ModelConfig) -> litellm.DeploymentTypedDict:
         provider = self._model_provider_registry.get_provider(model_config.provider)

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc3}/src/data_designer/engine/models/registry.py

@@ -107,6 +107,11 @@ class ModelRegistry:
     def run_health_check(self, model_aliases: list[str]) -> None:
         logger.info("🩺 Running health checks for models...")
         for model_alias in model_aliases:
+            model_config = self.get_model_config(model_alias=model_alias)
+            if model_config.skip_health_check:
+                logger.info(f"  |-- ⏭️  Skipping health check for model alias {model_alias!r} (skip_health_check=True)")
+                continue
+
             model = self.get_model(model_alias=model_alias)
             logger.info(
                 f"  |-- 👀 Checking {model.model_name!r} in provider named {model.model_provider_name!r} for model alias {model.model_alias!r}..."

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc3}/src/data_designer/engine/models/telemetry.py

@@ -8,6 +8,7 @@ Environment variables:
 - NEMO_TELEMETRY_ENABLED: Whether telemetry is enabled.
 - NEMO_DEPLOYMENT_TYPE: The deployment type the event came from.
 - NEMO_TELEMETRY_ENDPOINT: The endpoint to send the telemetry events to.
+- NEMO_SESSION_PREFIX: Optional prefix to add to session IDs.
 """
 
 from __future__ import annotations
@@ -18,15 +19,12 @@ import platform
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from enum import Enum
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import Any, ClassVar
 
 from pydantic import BaseModel, Field
 
 from data_designer.lazy_heavy_imports import httpx
 
-if TYPE_CHECKING:
-    import httpx
-
 TELEMETRY_ENABLED = os.getenv("NEMO_TELEMETRY_ENABLED", "true").lower() in ("1", "true", "yes")
 CLIENT_ID = "184482118588404"
 NEMO_TELEMETRY_VERSION = "nemo-telemetry/1.0"
@@ -35,6 +33,7 @@ NEMO_TELEMETRY_ENDPOINT = os.getenv(
     "NEMO_TELEMETRY_ENDPOINT", "https://events.telemetry.data.nvidia.com/v1.1/events/json"
 ).lower()
 CPU_ARCHITECTURE = platform.uname().machine
+SESSION_PREFIX = os.getenv("NEMO_SESSION_PREFIX")
 
 
 class NemoSourceEnum(str, Enum):
@@ -231,7 +230,11 @@ class TelemetryHandler:
         self._timer_task: asyncio.Task | None = None
         self._running = False
         self._source_client_version = source_client_version
-        self._session_id = session_id
+        # Apply session prefix if environment variable is set
+        if SESSION_PREFIX:
+            self._session_id = f"{SESSION_PREFIX}{session_id}"
+        else:
+            self._session_id = session_id
 
     async def astart(self) -> None:
         if self._running:

data_designer_engine-0.4.0rc3/src/data_designer/engine/models/utils.py

@@ -0,0 +1,101 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+
+@dataclass
+class ChatMessage:
+    """A chat message in an LLM conversation.
+
+    This dataclass represents messages exchanged in a conversation with an LLM,
+    supporting various message types including user prompts, assistant responses,
+    system instructions, and tool interactions.
+
+    Attributes:
+        role: The role of the message sender. One of 'user', 'assistant', 'system', or 'tool'.
+        content: The message content. Can be a string or a list of content blocks
+            for multimodal messages (e.g., text + images).
+        reasoning_content: Optional reasoning/thinking content from the assistant,
+            typically from extended thinking or chain-of-thought models.
+        tool_calls: Optional list of tool calls requested by the assistant.
+            Each tool call contains 'id', 'type', and 'function' keys.
+        tool_call_id: Optional ID linking a tool response to its corresponding
+            tool call. Required for messages with role='tool'.
+    """
+
+    role: Literal["user", "assistant", "system", "tool"]
+    content: str | list[dict[str, Any]] = ""
+    reasoning_content: str | None = None
+    tool_calls: list[dict[str, Any]] = field(default_factory=list)
+    tool_call_id: str | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert the message to a dictionary format for API calls.
+
+        Returns:
+            A dictionary containing the message fields. Only includes non-empty
+            optional fields to keep the output clean.
+        """
+        result: dict[str, Any] = {"role": self.role, "content": self.content}
+        if self.reasoning_content:
+            result["reasoning_content"] = self.reasoning_content
+        if self.tool_calls:
+            result["tool_calls"] = self.tool_calls
+        if self.tool_call_id:
+            result["tool_call_id"] = self.tool_call_id
+        return result
+
+    @classmethod
+    def as_user(cls, content: str | list[dict[str, Any]]) -> ChatMessage:
+        """Create a user message."""
+        return cls(role="user", content=content)
+
+    @classmethod
+    def as_assistant(
+        cls,
+        content: str = "",
+        reasoning_content: str | None = None,
+        tool_calls: list[dict[str, Any]] | None = None,
+    ) -> ChatMessage:
+        """Create an assistant message."""
+        return cls(
+            role="assistant",
+            content=content,
+            reasoning_content=reasoning_content,
+            tool_calls=tool_calls or [],
+        )
+
+    @classmethod
+    def as_system(cls, content: str) -> ChatMessage:
+        """Create a system message."""
+        return cls(role="system", content=content)
+
+    @classmethod
+    def as_tool(cls, content: str, tool_call_id: str) -> ChatMessage:
+        """Create a tool response message."""
+        return cls(role="tool", content=content, tool_call_id=tool_call_id)
+
+
+def prompt_to_messages(
+    *,
+    user_prompt: str,
+    system_prompt: str | None = None,
+    multi_modal_context: list[dict[str, Any]] | None = None,
+) -> list[ChatMessage]:
+    """Convert a user and system prompt into ChatMessage list.
+
+    Args:
+        user_prompt (str): A user prompt.
+        system_prompt (str, optional): An optional system prompt.
+    """
+    user_content: str | list[dict[str, Any]] = user_prompt
+    if multi_modal_context:
+        user_content = [*multi_modal_context, {"type": "text", "text": user_prompt}]
+
+    if system_prompt:
+        return [ChatMessage.as_system(system_prompt), ChatMessage.as_user(user_content)]
+    return [ChatMessage.as_user(user_content)]

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc3}/src/data_designer/engine/processing/processors/schema_transform.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 
 import json
 import logging
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from data_designer.config.processors import SchemaTransformProcessorConfig
 from data_designer.engine.dataset_builders.artifact_storage import BatchStage
@@ -20,6 +20,26 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def _json_escape_record(record: dict[str, Any]) -> dict[str, Any]:
+    """Escape record values for safe insertion into a JSON template."""
+
+    def escape_for_json_string(s: str) -> str:
+        """Use json.dumps to escape, then strip the surrounding quotes."""
+        return json.dumps(s)[1:-1]
+
+    escaped = {}
+    for key, value in record.items():
+        if isinstance(value, str):
+            escaped[key] = escape_for_json_string(value)
+        elif isinstance(value, (dict, list)):
+            escaped[key] = escape_for_json_string(json.dumps(value))
+        elif value is None:
+            escaped[key] = "null"
+        else:
+            escaped[key] = str(value)
+    return escaped
+
+
 class SchemaTransformProcessor(WithJinja2UserTemplateRendering, Processor[SchemaTransformProcessorConfig]):
     @property
     def template_as_str(self) -> str:
@@ -27,10 +47,12 @@ class SchemaTransformProcessor(WithJinja2UserTemplateRendering, Processor[Schema
 
     def process(self, data: pd.DataFrame, *, current_batch_number: int | None = None) -> pd.DataFrame:
         self.prepare_jinja2_template_renderer(self.template_as_str, data.columns.to_list())
-        formatted_records = [
-            json.loads(self.render_template(deserialize_json_values(record)).replace("\n", "\\n"))
-            for record in data.to_dict(orient="records")
-        ]
+        formatted_records = []
+        for record in data.to_dict(orient="records"):
+            deserialized = deserialize_json_values(record)
+            escaped = _json_escape_record(deserialized)
+            rendered = self.render_template(escaped)
+            formatted_records.append(json.loads(rendered))
         formatted_data = pd.DataFrame(formatted_records)
         if current_batch_number is not None:
             self.artifact_storage.write_batch_to_parquet_file(

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc3}/tests/engine/column_generators/generators/test_llm_completion_generators.py

@@ -12,14 +12,15 @@ from data_designer.config.column_configs import (
     LLMTextColumnConfig,
 )
 from data_designer.config.run_config import RunConfig
+from data_designer.config.utils.constants import TRACE_COLUMN_POSTFIX
 from data_designer.engine.column_generators.generators.base import GenerationStrategy
 from data_designer.engine.column_generators.generators.llm_completion import (
-    REASONING_TRACE_COLUMN_POSTFIX,
     LLMCodeCellGenerator,
     LLMJudgeCellGenerator,
     LLMStructuredCellGenerator,
     LLMTextCellGenerator,
 )
+from data_designer.engine.models.utils import ChatMessage
 
 
 def _create_generator_with_mocks(config_class=LLMTextColumnConfig, **config_kwargs):
@@ -67,14 +68,14 @@ def _create_generator_with_mocks(config_class=LLMTextColumnConfig, **config_kwar
     )
 
 
-def _setup_generate_mocks(mock_prompt_renderer, mock_response_recipe, mock_model, output="test_output", reasoning=None):
+def _setup_generate_mocks(mock_prompt_renderer, mock_response_recipe, mock_model, output="test_output"):
     """Helper function to setup common generate method mocks."""
     mock_prompt_renderer.render.side_effect = ["rendered_user_prompt", "rendered_system_prompt"]
     mock_response_recipe.serialize_output.return_value = {"result": output}
-    mock_model.generate.return_value = ({"result": output}, reasoning)
+    mock_model.generate.return_value = ({"result": output}, [])
 
 
-def test_generate_method():
+def test_generate_method() -> None:
     generator, _, mock_model, _, _, mock_prompt_renderer, mock_response_recipe = _create_generator_with_mocks()
 
     # Test basic generation
@@ -87,16 +88,19 @@ def test_generate_method():
     assert mock_model.generate.call_args[1]["max_correction_steps"] == 2
     assert mock_model.generate.call_args[1]["max_conversation_restarts"] == 7
     assert result["test_column"] == {"result": "test_output"}
-    assert "test_column" + REASONING_TRACE_COLUMN_POSTFIX not in result
+    assert "test_column" + TRACE_COLUMN_POSTFIX not in result
 
-    # Test with reasoning trace
+    # Test with full trace enabled
     mock_model.reset_mock()
     mock_prompt_renderer.reset_mock()
-    _setup_generate_mocks(mock_prompt_renderer, mock_response_recipe, mock_model, reasoning="reasoning_trace")
+    generator.resource_provider.run_config.debug_override_save_all_column_traces = True
+    mock_prompt_renderer.render.side_effect = ["rendered_user_prompt", "rendered_system_prompt"]
+    mock_response_recipe.serialize_output.return_value = {"result": "test_output"}
+    mock_model.generate.return_value = ({"result": "test_output"}, [ChatMessage.as_user("x")])
     result = generator.generate(data)
 
     assert result["test_column"] == {"result": "test_output"}
-    assert result["test_column" + REASONING_TRACE_COLUMN_POSTFIX] == "reasoning_trace"
+    assert result["test_column" + TRACE_COLUMN_POSTFIX] == [{"role": "user", "content": "x"}]
 
     # Test multi-modal context is None
     call_args = mock_model.generate.call_args
@@ -235,7 +239,7 @@ def test_generate_with_errors(error_type, error_message):
 
     if error_type == "serialization":
         mock_response_recipe.serialize_output.side_effect = Exception(error_message)
-        mock_model.generate.return_value = ({"result": "test_output"}, None)
+        mock_model.generate.return_value = ({"result": "test_output"}, [])
     elif error_type == "model":
         mock_model.generate.side_effect = Exception(error_message)
     elif error_type == "prompt_render":
@@ -249,13 +253,12 @@ def test_generate_with_errors(error_type, error_message):
 
 def test_generate_with_complex_data():
     generator, _, mock_model, _, _, mock_prompt_renderer, mock_response_recipe = _create_generator_with_mocks()
-    _setup_generate_mocks(mock_prompt_renderer, mock_response_recipe, mock_model, "complex_output", "complex_reasoning")
+    _setup_generate_mocks(mock_prompt_renderer, mock_response_recipe, mock_model, "complex_output")
 
     data = {"input": "test_input", "nested": {"key": "value"}, "list": [1, 2, 3], "json_string": '{"key": "value"}'}
     result = generator.generate(data)
 
     assert result["test_column"] == {"result": "complex_output"}
-    assert result["test_column" + REASONING_TRACE_COLUMN_POSTFIX] == "complex_reasoning"
     assert result["input"] == "test_input"
     assert result["nested"] == {"key": "value"}
     assert result["list"] == [1, 2, 3]
@@ -341,7 +344,7 @@ def test_generator_output_type_handling(
     mock_response_recipe.serialize_output.return_value = serialized_output
     stub_resource_provider.model_registry.get_model.return_value.generate.return_value = (
         {"result": "raw_output"},
-        None,
+        [],
     )
 
     data = {"input": "test_input"}