data-designer-engine 0.4.0rc3__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

data_designer/engine/analysis/column_profilers/base.py

@@ -10,8 +10,7 @@ from typing import TYPE_CHECKING
 from pydantic import BaseModel, model_validator
 from typing_extensions import Self
 
-from data_designer.config.base import ConfigBase
-from data_designer.config.column_configs import SingleColumnConfig
+from data_designer.config.base import ConfigBase, SingleColumnConfig
 from data_designer.config.column_types import DataDesignerColumnType
 from data_designer.engine.configurable_task import ConfigurableTask, TaskConfigT
 from data_designer.lazy_heavy_imports import pd

data_designer/engine/analysis/dataset_profiler.py

@@ -12,8 +12,7 @@ from pydantic import Field, field_validator
 
 from data_designer.config.analysis.column_profilers import ColumnProfilerConfigT
 from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
-from data_designer.config.base import ConfigBase
-from data_designer.config.column_configs import SingleColumnConfig
+from data_designer.config.base import ConfigBase, SingleColumnConfig
 from data_designer.config.column_types import ColumnConfigT
 from data_designer.engine.analysis.column_profilers.base import ColumnConfigWithDataFrame, ColumnProfiler
 from data_designer.engine.analysis.column_statistics import get_column_statistics_calculator

data_designer/engine/column_generators/generators/base.py

@@ -6,9 +6,9 @@ from __future__ import annotations
 import functools
 import logging
 from abc import ABC, abstractmethod
-from enum import Enum
 from typing import TYPE_CHECKING, overload
 
+from data_designer.config.column_configs import GenerationStrategy
 from data_designer.engine.configurable_task import ConfigurableTask, DataT, TaskConfigT
 from data_designer.lazy_heavy_imports import pd
 
@@ -22,11 +22,6 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-class GenerationStrategy(str, Enum):
-    CELL_BY_CELL = "cell_by_cell"
-    FULL_COLUMN = "full_column"
-
-
 class ColumnGenerator(ConfigurableTask[TaskConfigT], ABC):
     @property
     def can_generate_from_scratch(self) -> bool:

data_designer/engine/column_generators/generators/custom.py

@@ -0,0 +1,195 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Custom column generator using user-provided callable functions."""
+
+from __future__ import annotations
+
+import inspect
+import logging
+from typing import TYPE_CHECKING, Any
+
+from data_designer.config.column_configs import CustomColumnConfig, GenerationStrategy
+from data_designer.engine.column_generators.generators.base import ColumnGenerator
+from data_designer.engine.column_generators.utils.errors import CustomColumnGenerationError
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+logger = logging.getLogger(__name__)
+
+
+class CustomColumnGenerator(ColumnGenerator[CustomColumnConfig]):
+    """Column generator that uses a user-provided callable function.
+
+    Supports two strategies based on config.strategy:
+        - cell_by_cell: Processes rows one at a time (dict -> dict), parallelized by framework.
+        - full_column: Processes entire batch (DataFrame -> DataFrame) for vectorized ops.
+
+    Supported function signatures (validated by parameter name):
+        - fn(row) -> dict                              # cell_by_cell, simple transform
+        - fn(row, generator_params) -> dict            # cell_by_cell, with typed params
+        - fn(row, generator_params, models) -> dict    # cell_by_cell, with LLM access
+        - fn(df) -> DataFrame                          # full_column, simple transform
+        - fn(df, generator_params) -> DataFrame        # full_column, with typed params
+        - fn(df, generator_params, models) -> DataFrame  # full_column, with LLM access
+
+    The models dict provides direct access to ModelFacade instances keyed by alias.
+    """
+
+    def get_generation_strategy(self) -> GenerationStrategy:
+        """Return strategy based on config."""
+        return self.config.generation_strategy
+
+    def generate(self, data: dict | pd.DataFrame) -> dict | pd.DataFrame:
+        """Generate column value(s) for a row (dict) or batch (DataFrame)."""
+        is_full_column = self.config.generation_strategy == GenerationStrategy.FULL_COLUMN
+        is_dataframe = not isinstance(data, dict)
+
+        # Validate data type matches strategy
+        if is_full_column and not is_dataframe:
+            raise CustomColumnGenerationError(
+                f"Custom generator {self.config.name!r} is configured for 'full_column' strategy "
+                "but received a dict. Expected a DataFrame."
+            )
+        if not is_full_column and is_dataframe:
+            raise CustomColumnGenerationError(
+                f"Custom generator {self.config.name!r} is configured for 'cell_by_cell' strategy "
+                "but received a DataFrame. Expected a dict."
+            )
+
+        return self._generate(data, is_dataframe)
+
+    def _generate(self, data: dict | pd.DataFrame, is_dataframe: bool) -> dict | pd.DataFrame:
+        """Unified generation logic for both strategies."""
+        # Get columns/keys using unified accessor
+        get_keys = (lambda d: set(d.columns)) if is_dataframe else (lambda d: set(d.keys()))
+        expected_type = pd.DataFrame if is_dataframe else dict
+        type_name = "DataFrame" if is_dataframe else "dict"
+
+        # Check required columns
+        missing = set(self.config.required_columns) - get_keys(data)
+        if missing:
+            raise CustomColumnGenerationError(
+                f"Missing required columns for custom generator '{self.config.name}': {sorted(missing)}"
+            )
+
+        keys_before = get_keys(data)
+
+        # Invoke generator
+        try:
+            result = self._invoke_generator_function(data)
+        except CustomColumnGenerationError:
+            raise
+        except Exception as e:
+            raise CustomColumnGenerationError(
+                f"Custom generator function failed for column '{self.config.name}': {e}"
+            ) from e
+
+        # Validate return type
+        if not isinstance(result, expected_type):
+            raise CustomColumnGenerationError(
+                f"Custom generator for column '{self.config.name}' must return a {type_name}, "
+                f"got {type(result).__name__}"
+            )
+
+        return self._validate_output(result, keys_before, is_dataframe)
+
+    def _validate_output(
+        self, result: dict | pd.DataFrame, keys_before: set[str], is_dataframe: bool
+    ) -> dict | pd.DataFrame:
+        """Validate output columns and remove undeclared ones."""
+        # Unified accessors
+        get_keys = (lambda d: set(d.columns)) if is_dataframe else (lambda d: set(d.keys()))
+        container_name = "DataFrame" if is_dataframe else "row"
+
+        expected_new = {self.config.name} | set(self.config.side_effect_columns)
+        result_keys = get_keys(result)
+
+        # Check primary column exists
+        if self.config.name not in result_keys:
+            raise CustomColumnGenerationError(
+                f"Custom generator for column '{self.config.name}' did not create the expected column. "
+                f"The generator_function must add a key named '{self.config.name}' to the {container_name}."
+            )
+
+        # Check side effect columns exist
+        missing = set(self.config.side_effect_columns) - result_keys
+        if missing:
+            raise CustomColumnGenerationError(
+                f"Custom generator for column '{self.config.name}' did not create declared side_effect_columns: "
+                f"{sorted(missing)}. Declared side_effect_columns must be added to the {container_name}."
+            )
+
+        # Check no pre-existing columns removed
+        removed = keys_before - result_keys
+        if removed:
+            raise CustomColumnGenerationError(
+                f"Custom generator for column '{self.config.name}' removed pre-existing columns: "
+                f"{sorted(removed)}. The generator_function must not remove any existing columns."
+            )
+
+        # Remove undeclared columns with warning
+        undeclared = (result_keys - keys_before) - expected_new
+        if undeclared:
+            logger.warning(
+                f"⚠️ Custom generator for column '{self.config.name}' created undeclared columns: "
+                f"{sorted(undeclared)}. These columns will be removed. "
+                f"To keep additional columns, declare them in @custom_column_generator(side_effect_columns=[...])."
+            )
+            if is_dataframe:
+                result = result.drop(columns=list(undeclared))
+            else:
+                for key in undeclared:
+                    del result[key]
+
+        return result
+
+    def _invoke_generator_function(self, data: dict | pd.DataFrame) -> dict | pd.DataFrame:
+        """Invoke the user's generate function with appropriate arguments based on signature."""
+        params = self._get_validated_params()
+
+        if len(params) == 1:
+            return self.config.generator_function(data)
+        elif len(params) == 2:
+            return self.config.generator_function(data, self.config.generator_params)
+        else:
+            models = self._build_models_dict()
+            return self.config.generator_function(data, self.config.generator_params, models)
+
+    def _build_models_dict(self) -> dict[str, Any]:
+        """Build a dict of ModelFacade instances from model_aliases."""
+        return {
+            alias: self.resource_provider.model_registry.get_model(model_alias=alias)
+            for alias in self.config.model_aliases
+        }
+
+    def _get_validated_params(self) -> list[inspect.Parameter]:
+        """Get positional params and validate first param matches generation strategy."""
+        params = [
+            p
+            for p in inspect.signature(self.config.generator_function).parameters.values()
+            if p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+        ]
+        # Decorator validated param names; here we only check strategy match
+        is_full = self.config.generation_strategy == GenerationStrategy.FULL_COLUMN
+        expected = "df" if is_full else "row"
+        if params[0].name != expected:
+            raise CustomColumnGenerationError(
+                f"Generator '{self.config.name}': strategy is {'full_column' if is_full else 'cell_by_cell'}, "
+                f"first parameter must be '{expected}', got '{params[0].name}'."
+            )
+        return params
+
+    def log_pre_generation(self) -> None:
+        logger.info(f"{self.config.get_column_emoji()} Custom column config for column '{self.config.name}'")
+        logger.info(f"  |-- generator_function: {self.config.generator_function.__name__!r}")
+        logger.info(f"  |-- generation_strategy: {self.config.generation_strategy!r}")
+        logger.info(f"  |-- required_columns: {self.config.required_columns}")
+        if self.config.side_effect_columns:
+            logger.info(f"  |-- side_effect_columns: {self.config.side_effect_columns}")
+        if self.config.model_aliases:
+            logger.info(f"  |-- model_aliases: {self.config.model_aliases}")
+        if self.config.generator_params:
+            logger.info(f"  |-- generator_params: {self.config.generator_params}")

data_designer/engine/column_generators/generators/llm_completion.py

@@ -12,7 +12,8 @@ from data_designer.config.column_configs import (
     LLMStructuredColumnConfig,
     LLMTextColumnConfig,
 )
-from data_designer.config.utils.constants import TRACE_COLUMN_POSTFIX
+from data_designer.config.utils.constants import REASONING_CONTENT_COLUMN_POSTFIX, TRACE_COLUMN_POSTFIX
+from data_designer.config.utils.trace_type import TraceType
 from data_designer.engine.column_generators.generators.base import ColumnGeneratorWithModel, GenerationStrategy
 from data_designer.engine.column_generators.utils.prompt_renderer import (
     PromptType,
@@ -79,6 +80,7 @@ class ColumnGeneratorWithModelChatCompletion(ColumnGeneratorWithModel[TaskConfig
             ),
             parser=self.response_recipe.parse,
             multi_modal_context=multi_modal_context,
+            tool_alias=self.config.tool_alias,
             max_correction_steps=self.max_conversation_correction_steps,
             max_conversation_restarts=self.max_conversation_restarts,
             purpose=f"running generation for column '{self.config.name}'",
@@ -87,14 +89,39 @@ class ColumnGeneratorWithModelChatCompletion(ColumnGeneratorWithModel[TaskConfig
         serialized_output = self.response_recipe.serialize_output(response)
         data[self.config.name] = self._process_serialized_output(serialized_output)
 
-        should_save_trace = (
-            self.config.with_trace or self.resource_provider.run_config.debug_override_save_all_column_traces
-        )
-        if should_save_trace:
+        effective_trace_type = self.config.with_trace
+
+        if effective_trace_type == TraceType.ALL_MESSAGES:
             data[self.config.name + TRACE_COLUMN_POSTFIX] = [message.to_dict() for message in trace]
+        elif effective_trace_type == TraceType.LAST_MESSAGE:
+            last_assistant = next((m for m in reversed(trace) if m.role == "assistant"), None)
+            data[self.config.name + TRACE_COLUMN_POSTFIX] = [last_assistant.to_dict()] if last_assistant else []
+
+        if self.config.extract_reasoning_content:
+            data[self.config.name + REASONING_CONTENT_COLUMN_POSTFIX] = self._extract_reasoning_content(trace)
 
         return data
 
+    def _extract_reasoning_content(self, trace: list) -> str | None:
+        """Extract reasoning_content from the final assistant message in the trace.
+
+        Args:
+            trace: List of ChatMessage objects from the generation.
+
+        Returns:
+            The stripped reasoning_content from the final assistant message, or None if not present.
+        """
+        reasoning_value: str | None = None
+        for message in reversed(trace):
+            if message.role == "assistant":
+                reasoning_value = message.reasoning_content
+                break
+
+        if reasoning_value is not None:
+            reasoning_value = reasoning_value.strip() or None
+
+        return reasoning_value
+
     def _process_serialized_output(self, serialized_output: str) -> str | dict | list:
         """Process the serialized output from the model. Subclasses can override to customize deserialization."""
         return serialized_output

data_designer/engine/column_generators/registry.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.column_configs import (
+    CustomColumnConfig,
     EmbeddingColumnConfig,
     ExpressionColumnConfig,
     LLMCodeColumnConfig,
@@ -15,6 +16,7 @@ from data_designer.config.column_configs import (
 )
 from data_designer.config.column_types import DataDesignerColumnType
 from data_designer.engine.column_generators.generators.base import ColumnGenerator
+from data_designer.engine.column_generators.generators.custom import CustomColumnGenerator
 from data_designer.engine.column_generators.generators.embedding import EmbeddingCellGenerator
 from data_designer.engine.column_generators.generators.expression import ExpressionColumnGenerator
 from data_designer.engine.column_generators.generators.llm_completion import (
@@ -40,6 +42,7 @@ class ColumnGeneratorRegistry(TaskRegistry[DataDesignerColumnType, ColumnGenerat
 
 def create_default_column_generator_registry(with_plugins: bool = True) -> ColumnGeneratorRegistry:
     registry = ColumnGeneratorRegistry()
+    registry.register(DataDesignerColumnType.CUSTOM, CustomColumnGenerator, CustomColumnConfig)
     registry.register(DataDesignerColumnType.LLM_TEXT, LLMTextCellGenerator, LLMTextColumnConfig)
     registry.register(DataDesignerColumnType.LLM_CODE, LLMCodeCellGenerator, LLMCodeColumnConfig)
     registry.register(DataDesignerColumnType.LLM_JUDGE, LLMJudgeCellGenerator, LLMJudgeColumnConfig)

data_designer/engine/column_generators/utils/errors.py

@@ -13,3 +13,6 @@ class ExpressionTemplateRenderError(DataDesignerError): ...
 
 
 class SeedDatasetError(DataDesignerError): ...
+
+
+class CustomColumnGenerationError(DataDesignerError): ...

data_designer/engine/column_generators/utils/prompt_renderer.py

@@ -6,7 +6,7 @@ from __future__ import annotations
 import json
 import logging
 
-from data_designer.config.column_configs import SingleColumnConfig
+from data_designer.config.base import SingleColumnConfig
 from data_designer.config.column_types import DataDesignerColumnType
 from data_designer.config.models import ModelConfig
 from data_designer.config.utils.code_lang import CodeLang

data_designer/engine/dataset_builders/column_wise_builder.py

@@ -12,6 +12,7 @@ import uuid
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable
 
+from data_designer.config.column_configs import CustomColumnConfig
 from data_designer.config.column_types import ColumnConfigT
 from data_designer.config.config_builder import BuilderConfig
 from data_designer.config.data_designer_config import DataDesignerConfig
@@ -97,6 +98,7 @@ class ColumnWiseDatasetBuilder:
         on_batch_complete: Callable[[Path], None] | None = None,
     ) -> Path:
         self._run_model_health_check_if_needed()
+        self._run_mcp_tool_check_if_needed()
         self._write_builder_config()
         generators = self._initialize_generators()
         start_time = time.perf_counter()
@@ -125,6 +127,7 @@ class ColumnWiseDatasetBuilder:
 
     def build_preview(self, *, num_records: int) -> pd.DataFrame:
         self._run_model_health_check_if_needed()
+        self._run_mcp_tool_check_if_needed()
 
         generators = self._initialize_generators()
         group_id = uuid.uuid4().hex
@@ -209,11 +212,26 @@ class ColumnWiseDatasetBuilder:
         df = generator.generate(self.batch_manager.get_current_batch(as_dataframe=True))
         self.batch_manager.update_records(df.to_dict(orient="records"))
 
-    def _run_model_health_check_if_needed(self) -> bool:
-        if any(column_type_is_model_generated(config.column_type) for config in self.single_column_configs):
-            self._resource_provider.model_registry.run_health_check(
-                list(set(config.model_alias for config in self.llm_generated_column_configs))
-            )
+    def _run_model_health_check_if_needed(self) -> None:
+        model_aliases: set[str] = set()
+        for config in self.single_column_configs:
+            if column_type_is_model_generated(config.column_type):
+                model_aliases.add(config.model_alias)
+            if isinstance(config, CustomColumnConfig) and config.model_aliases:
+                model_aliases.update(config.model_aliases)
+
+        if model_aliases:
+            self._resource_provider.model_registry.run_health_check(list(model_aliases))
+
+    def _run_mcp_tool_check_if_needed(self) -> None:
+        tool_aliases = sorted(
+            {config.tool_alias for config in self.llm_generated_column_configs if getattr(config, "tool_alias", None)}
+        )
+        if not tool_aliases:
+            return
+        if self._resource_provider.mcp_registry is None:
+            raise DatasetGenerationError(f"Tool alias(es) {tool_aliases!r} specified but no MCPRegistry configured.")
+        self._resource_provider.mcp_registry.run_health_check(tool_aliases)
 
     def _fan_out_with_threads(self, generator: ColumnGeneratorWithModelRegistry, max_workers: int) -> None:
         if generator.get_generation_strategy() != GenerationStrategy.CELL_BY_CELL:

data_designer/engine/dataset_builders/multi_column_configs.py

@@ -8,8 +8,8 @@ from typing import TypeAlias
 
 from pydantic import Field, field_validator
 
-from data_designer.config.base import ConfigBase
-from data_designer.config.column_configs import SamplerColumnConfig, SeedDatasetColumnConfig, SingleColumnConfig
+from data_designer.config.base import ConfigBase, SingleColumnConfig
+from data_designer.config.column_configs import SamplerColumnConfig, SeedDatasetColumnConfig
 from data_designer.config.column_types import ColumnConfigT, DataDesignerColumnType
 from data_designer.config.sampler_constraints import ColumnConstraintT
 from data_designer.config.seed import SeedConfig

data_designer/engine/mcp/__init__.py

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from data_designer.engine.mcp import io
+from data_designer.engine.mcp.errors import (
+    DuplicateToolNameError,
+    MCPClientUnavailableError,
+    MCPConfigurationError,
+    MCPError,
+    MCPToolError,
+)
+from data_designer.engine.mcp.facade import MCPFacade
+from data_designer.engine.mcp.factory import create_mcp_registry
+from data_designer.engine.mcp.registry import MCPRegistry, MCPToolDefinition, MCPToolResult
+
+__all__ = [
+    "DuplicateToolNameError",
+    "MCPClientUnavailableError",
+    "MCPConfigurationError",
+    "MCPError",
+    "MCPFacade",
+    "MCPRegistry",
+    "MCPToolDefinition",
+    "MCPToolError",
+    "MCPToolResult",
+    "create_mcp_registry",
+    "io",
+]

data_designer/engine/mcp/errors.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from data_designer.errors import DataDesignerError
+
+
+class MCPError(DataDesignerError): ...
+
+
+class MCPConfigurationError(MCPError): ...
+
+
+class MCPClientUnavailableError(MCPError): ...
+
+
+class MCPToolError(MCPError): ...
+
+
+class DuplicateToolNameError(MCPConfigurationError):
+    """Raised when the same tool name exists in multiple MCP providers or tool configs."""