data-designer-config 0.4.0rc2__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

data_designer/config/column_configs.py

@@ -3,66 +3,28 @@
 
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
-from typing import Annotated, Literal
+from enum import Enum
+from typing import Annotated, Any, Literal
 
-from pydantic import BaseModel, Discriminator, Field, model_validator
+from pydantic import BaseModel, Discriminator, Field, field_serializer, field_validator, model_validator
 from typing_extensions import Self
 
-from data_designer.config.base import ConfigBase
+from data_designer.config.base import ConfigBase, SingleColumnConfig
 from data_designer.config.errors import InvalidConfigError
 from data_designer.config.models import ImageContext
 from data_designer.config.sampler_params import SamplerParamsT, SamplerType
 from data_designer.config.utils.code_lang import CodeLang
-from data_designer.config.utils.constants import REASONING_TRACE_COLUMN_POSTFIX
+from data_designer.config.utils.constants import REASONING_CONTENT_COLUMN_POSTFIX, 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
 
 
-class SingleColumnConfig(ConfigBase, ABC):
-    """Abstract base class for all single-column configuration types.
+class GenerationStrategy(str, Enum):
+    """Strategy for custom column generation."""
 
-    This class serves as the foundation for all column configurations in DataDesigner,
-    defining shared fields and properties across all column types.
-
-    Attributes:
-        name: Unique name of the column to be generated.
-        drop: If True, the column will be generated but removed from the final dataset.
-            Useful for intermediate columns that are dependencies for other columns.
-        column_type: Discriminator field that identifies the specific column type.
-            Subclasses must override this field to specify the column type with a `Literal` value.
-    """
-
-    name: str
-    drop: bool = False
-    column_type: str
-
-    @staticmethod
-    def get_column_emoji() -> str:
-        return "🎨"
-
-    @property
-    @abstractmethod
-    def required_columns(self) -> list[str]:
-        """Returns a list of column names that must exist before this column can be generated.
-
-        Returns:
-            List of column names that this column depends on. Empty list indicates
-            no dependencies. Override in subclasses to specify dependencies.
-        """
-
-    @property
-    @abstractmethod
-    def side_effect_columns(self) -> list[str]:
-        """Returns a list of additional columns that this column will create as a side effect.
-
-        Some column types generate additional metadata or auxiliary columns alongside
-        the primary column (e.g., reasoning traces for LLM columns).
-
-        Returns:
-            List of column names that this column will create as a side effect. Empty list
-            indicates no side effect columns. Override in subclasses to specify side effects.
-        """
+    CELL_BY_CELL = "cell_by_cell"
+    FULL_COLUMN = "full_column"
 
 
 class SamplerColumnConfig(SingleColumnConfig):
@@ -143,8 +105,8 @@ class LLMTextColumnConfig(SingleColumnConfig):
 
     LLM text columns generate free-form text content using language models via LiteLLM.
     Prompts support Jinja2 templating to reference values from other columns, enabling
-    context-aware generation. The generated text can optionally include reasoning traces
-    when models support extended thinking.
+    context-aware generation. The generated text can optionally include message traces
+    capturing the full conversation history.
 
     Attributes:
         prompt: Prompt template for text generation. Supports Jinja2 syntax to
@@ -159,6 +121,18 @@ class LLMTextColumnConfig(SingleColumnConfig):
             `LLMStructuredColumnConfig` for structured output, `LLMCodeColumnConfig` for code.
         multi_modal_context: Optional list of image contexts for multi-modal generation.
             Enables vision-capable models to generate text based on image inputs.
+        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: 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).
+        extract_reasoning_content: If True, creates a `{column_name}__reasoning_content` column
+            containing only the reasoning_content from the final assistant response. This is
+            useful for models that expose chain-of-thought reasoning separately from the main
+            response. Defaults to False.
         column_type: Discriminator field, always "llm-text" for this configuration type.
     """
 
@@ -166,6 +140,9 @@ class LLMTextColumnConfig(SingleColumnConfig):
     model_alias: str
     system_prompt: str | None = None
     multi_modal_context: list[ImageContext] | None = None
+    tool_alias: str | None = None
+    with_trace: TraceType = TraceType.NONE
+    extract_reasoning_content: bool = False
     column_type: Literal["llm-text"] = "llm-text"
 
     @staticmethod
@@ -186,14 +163,20 @@ class LLMTextColumnConfig(SingleColumnConfig):
 
     @property
     def side_effect_columns(self) -> list[str]:
-        """Returns the reasoning trace column, which may be generated alongside the main column.
+        """Returns side-effect columns that may be generated alongside the main column.
 
-        Reasoning traces are only returned if the served model parses and returns reasoning content.
+        Side-effect columns include:
+        - `{name}__trace`: Generated when `with_trace` is not `TraceType.NONE` on the column
+          config.
+        - `{name}__reasoning_content`: Generated when `extract_reasoning_content=True`.
 
         Returns:
-            List containing the reasoning trace column name.
+            List of side-effect column names.
         """
-        return [f"{self.name}{REASONING_TRACE_COLUMN_POSTFIX}"]
+        return [
+            *([f"{self.name}{TRACE_COLUMN_POSTFIX}"] if self.with_trace != TraceType.NONE else []),
+            *([f"{self.name}{REASONING_CONTENT_COLUMN_POSTFIX}"] if self.extract_reasoning_content else []),
+        ]
 
     @model_validator(mode="after")
     def assert_prompt_valid_jinja(self) -> Self:
@@ -216,7 +199,7 @@ class LLMCodeColumnConfig(LLMTextColumnConfig):
 
     Extends LLMTextColumnConfig to generate code snippets in specific programming languages
     or SQL dialects. The generated code is automatically extracted from markdown code blocks
-    for the specified language. Inherits all prompt templating capabilities.
+    for the specified language. Inherits all prompt templating capabilities from LLMTextColumnConfig.
 
     Attributes:
         code_lang: Programming language or SQL dialect for code generation. Supported
@@ -224,6 +207,16 @@ class LLMCodeColumnConfig(LLMTextColumnConfig):
             "rust", "ruby", "scala", "swift", "sql:sqlite", "sql:postgres", "sql:mysql",
             "sql:tsql", "sql:bigquery", "sql:ansi". See CodeLang enum for complete list.
         column_type: Discriminator field, always "llm-code" for this configuration type.
+
+    Inherited Attributes:
+        prompt: Prompt template for code generation (supports Jinja2).
+        model_alias: Alias of the model configuration to use.
+        system_prompt: Optional system prompt (supports Jinja2).
+        multi_modal_context: Optional image contexts for multi-modal generation.
+        tool_alias: Optional tool configuration alias for MCP tool calls.
+        with_trace: If True, creates a `{column_name}__trace` column with message history.
+        extract_reasoning_content: If True, creates a `{column_name}__reasoning_content`
+            column containing the reasoning content from the final assistant response.
     """
 
     code_lang: CodeLang
@@ -239,13 +232,24 @@ class LLMStructuredColumnConfig(LLMTextColumnConfig):
 
     Extends LLMTextColumnConfig to generate structured data conforming to a specified schema.
     Uses JSON schema or Pydantic models to define the expected output structure, enabling
-    type-safe and validated structured output generation. Inherits prompt templating capabilities.
+    type-safe and validated structured output generation. Inherits prompt templating capabilities
+    from LLMTextColumnConfig.
 
     Attributes:
         output_format: The schema defining the expected output structure. Can be either:
             - A Pydantic BaseModel class (recommended)
             - A JSON schema dictionary
         column_type: Discriminator field, always "llm-structured" for this configuration type.
+
+    Inherited Attributes:
+        prompt: Prompt template for structured generation (supports Jinja2).
+        model_alias: Alias of the model configuration to use.
+        system_prompt: Optional system prompt (supports Jinja2).
+        multi_modal_context: Optional image contexts for multi-modal generation.
+        tool_alias: Optional tool configuration alias for MCP tool calls.
+        with_trace: If True, creates a `{column_name}__trace` column with message history.
+        extract_reasoning_content: If True, creates a `{column_name}__reasoning_content`
+            column containing the reasoning content from the final assistant response.
     """
 
     output_format: dict | type[BaseModel]
@@ -293,13 +297,24 @@ class LLMJudgeColumnConfig(LLMTextColumnConfig):
 
     Extends LLMTextColumnConfig to create judge columns that evaluate and score other
     generated content based on the defined criteria. Useful for quality assessment, preference
-    ranking, and multi-dimensional evaluation of generated data.
+    ranking, and multi-dimensional evaluation of generated data. Inherits prompt templating
+    capabilities from LLMTextColumnConfig.
 
     Attributes:
         scores: List of Score objects defining the evaluation dimensions. Each score
             represents a different aspect to evaluate (e.g., accuracy, relevance, fluency).
             Must contain at least one score.
         column_type: Discriminator field, always "llm-judge" for this configuration type.
+
+    Inherited Attributes:
+        prompt: Prompt template for the judge evaluation (supports Jinja2).
+        model_alias: Alias of the model configuration to use.
+        system_prompt: Optional system prompt (supports Jinja2).
+        multi_modal_context: Optional image contexts for multi-modal generation.
+        tool_alias: Optional tool configuration alias for MCP tool calls.
+        with_trace: If True, creates a `{column_name}__trace` column with message history.
+        extract_reasoning_content: If True, creates a `{column_name}__reasoning_content`
+            column containing the reasoning content from the final assistant response.
     """
 
     scores: list[Score] = Field(..., min_length=1)
@@ -468,3 +483,80 @@ class EmbeddingColumnConfig(SingleColumnConfig):
     @property
     def side_effect_columns(self) -> list[str]:
         return []
+
+
+class CustomColumnConfig(SingleColumnConfig):
+    """Configuration for custom user-defined column generators.
+
+    Custom columns allow users to provide their own generation logic via a callable function
+    decorated with `@custom_column_generator`. Two strategies are supported: cell_by_cell
+    (default, row-based) and full_column (batch-based with DataFrame access).
+
+    Attributes:
+        generator_function: A callable decorated with @custom_column_generator.
+        generation_strategy: "cell_by_cell" (row-based) or "full_column" (batch-based).
+        generator_params: Optional typed configuration object (Pydantic BaseModel) passed
+            as the second argument to the generator function.
+        column_type: Discriminator field, always "custom" for this configuration type.
+    """
+
+    generator_function: Any = Field(description="Function decorated with @custom_column_generator")
+    generation_strategy: GenerationStrategy = Field(
+        default=GenerationStrategy.CELL_BY_CELL,
+        description="Generation strategy: 'cell_by_cell' for row-based or 'full_column' for batch-based",
+    )
+    generator_params: BaseModel | None = Field(
+        default=None,
+        description="Optional typed configuration object passed as second argument to generator function",
+    )
+    column_type: Literal["custom"] = "custom"
+
+    @field_validator("generator_function")
+    @classmethod
+    def _validate_generator_function(cls, v: Any) -> Any:
+        if not callable(v):
+            raise ValueError("generator_function must be callable")
+        if not hasattr(v, "custom_column_metadata"):
+            raise ValueError("generator_function must be decorated with @custom_column_generator")
+        return v
+
+    @staticmethod
+    def get_column_emoji() -> str:
+        return "🔧"
+
+    @property
+    def required_columns(self) -> list[str]:
+        """Returns the columns required for custom generation (from decorator metadata)."""
+        metadata = getattr(self.generator_function, "custom_column_metadata", {})
+        return metadata.get("required_columns", [])
+
+    @property
+    def side_effect_columns(self) -> list[str]:
+        """Returns additional columns created by this generator (from decorator metadata)."""
+        metadata = getattr(self.generator_function, "custom_column_metadata", {})
+        return metadata.get("side_effect_columns", [])
+
+    @property
+    def model_aliases(self) -> list[str]:
+        """Returns model aliases for LLM access and health checks (from decorator metadata)."""
+        metadata = getattr(self.generator_function, "custom_column_metadata", {})
+        return metadata.get("model_aliases", [])
+
+    @field_serializer("generator_function")
+    def serialize_generator_function(self, v: Any) -> str:
+        return getattr(v, "__name__", repr(v))
+
+    @field_serializer("generator_params")
+    def serialize_generator_params(self, v: BaseModel | None) -> dict[str, Any] | None:
+        if v is None:
+            return None
+        return v.model_dump()
+
+    @model_validator(mode="after")
+    def validate_generator_function(self) -> Self:
+        if not callable(self.generator_function):
+            raise InvalidConfigError(
+                f"🛑 `generator_function` must be a callable for custom column '{self.name}'. "
+                f"Expected a function decorated with @custom_column_generator."
+            )
+        return self

data_designer/config/column_types.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 from typing_extensions import TypeAlias
 
 from data_designer.config.column_configs import (
+    CustomColumnConfig,
     EmbeddingColumnConfig,
     ExpressionColumnConfig,
     LLMCodeColumnConfig,
@@ -28,7 +29,8 @@ from data_designer.plugin_manager import PluginManager
 plugin_manager = PluginManager()
 
 ColumnConfigT: TypeAlias = (
-    ExpressionColumnConfig
+    CustomColumnConfig
+    | ExpressionColumnConfig
     | LLMCodeColumnConfig
     | LLMJudgeColumnConfig
     | LLMStructuredColumnConfig
@@ -87,6 +89,7 @@ def get_column_display_order() -> list[DataDesignerColumnType]:
         DataDesignerColumnType.EMBEDDING,
         DataDesignerColumnType.VALIDATION,
         DataDesignerColumnType.EXPRESSION,
+        DataDesignerColumnType.CUSTOM,
     ]
     display_order.extend(plugin_manager.get_plugin_column_types(DataDesignerColumnType))
     return display_order
@@ -129,6 +132,7 @@ def _resolve_sampler_kwargs(name: str, kwargs: dict) -> dict:
 
 
 _COLUMN_TYPE_CONFIG_CLS_MAP = {
+    DataDesignerColumnType.CUSTOM: CustomColumnConfig,
     DataDesignerColumnType.LLM_TEXT: LLMTextColumnConfig,
     DataDesignerColumnType.LLM_CODE: LLMCodeColumnConfig,
     DataDesignerColumnType.LLM_STRUCTURED: LLMStructuredColumnConfig,

data_designer/config/config_builder.py

@@ -13,7 +13,6 @@ from pygments.lexers import PythonLexer
 from typing_extensions import Self
 
 from data_designer.config.analysis.column_profilers import ColumnProfilerConfigT
-from data_designer.config.base import ExportableConfigBase
 from data_designer.config.column_configs import SeedDatasetColumnConfig
 from data_designer.config.column_types import (
     ColumnConfigT,
@@ -25,6 +24,8 @@ from data_designer.config.data_designer_config import DataDesignerConfig
 from data_designer.config.dataset_builders import BuildStage
 from data_designer.config.default_model_settings import get_default_model_configs
 from data_designer.config.errors import BuilderConfigurationError, BuilderSerializationError, InvalidColumnTypeError
+from data_designer.config.exportable_config import ExportableConfigBase
+from data_designer.config.mcp import ToolConfig
 from data_designer.config.models import ModelConfig, load_model_configs
 from data_designer.config.processors import ProcessorConfigT, ProcessorType, get_processor_config_from_kwargs
 from data_designer.config.sampler_constraints import (
@@ -93,7 +94,10 @@ class DataDesignerConfigBuilder:
             json_config = json.loads(serialize_data(smart_load_yaml(config)))
             builder_config = BuilderConfig.model_validate(json_config)
 
-        builder = cls(model_configs=builder_config.data_designer.model_configs)
+        builder = cls(
+            model_configs=builder_config.data_designer.model_configs,
+            tool_configs=builder_config.data_designer.tool_configs,
+        )
         data_designer_config = builder_config.data_designer
 
         for col in data_designer_config.columns:
@@ -111,7 +115,11 @@ class DataDesignerConfigBuilder:
 
         return builder
 
-    def __init__(self, model_configs: list[ModelConfig] | str | Path | None = None):
+    def __init__(
+        self,
+        model_configs: list[ModelConfig] | str | Path | None = None,
+        tool_configs: list[ToolConfig] | None = None,
+    ):
         """Initialize a new DataDesignerConfigBuilder instance.
 
         Args:
@@ -119,9 +127,13 @@ class DataDesignerConfigBuilder:
                 - None to use default model configurations in local mode
                 - A list of ModelConfig objects
                 - A string or Path to a model configuration file
+            tool_configs: Tool configurations for MCP tool calling. Can be:
+                - None if no tool configs are needed
+                - A list of ToolConfig objects
         """
         self._column_configs = {}
         self._model_configs = _load_model_configs(model_configs)
+        self._tool_configs: list[ToolConfig] = tool_configs or []
         self._processor_configs: list[ProcessorConfigT] = []
         self._seed_config: SeedConfig | None = None
         self._constraints: list[ColumnConstraintT] = []
@@ -136,6 +148,15 @@ class DataDesignerConfigBuilder:
         """
         return self._model_configs
 
+    @property
+    def tool_configs(self) -> list[ToolConfig]:
+        """Get the tool configurations for this builder.
+
+        Returns:
+            A list of ToolConfig objects used for MCP tool calling.
+        """
+        return self._tool_configs
+
     @property
     def allowed_references(self) -> list[str]:
         """Get all referenceable variables allowed in prompt templates and expressions.
@@ -184,6 +205,38 @@ class DataDesignerConfigBuilder:
             )
         return self
 
+    def add_tool_config(self, tool_config: ToolConfig) -> Self:
+        """Add a tool configuration to the current Data Designer configuration.
+
+        Args:
+            tool_config: The tool configuration to add.
+
+        Returns:
+            The current Data Designer config builder instance.
+
+        Raises:
+            BuilderConfigurationError: If a tool configuration with the same alias already exists.
+        """
+        if tool_config.tool_alias in {tc.tool_alias for tc in self._tool_configs}:
+            raise BuilderConfigurationError(
+                f"Tool configuration with alias {tool_config.tool_alias} already exists. "
+                "Please delete the existing tool configuration or choose a different alias."
+            )
+        self._tool_configs.append(tool_config)
+        return self
+
+    def delete_tool_config(self, alias: str) -> Self:
+        """Delete a tool configuration from the current Data Designer configuration by alias.
+
+        Args:
+            alias: The alias of the tool configuration to delete.
+
+        Returns:
+            The current Data Designer config builder instance.
+        """
+        self._tool_configs = [tc for tc in self._tool_configs if tc.tool_alias != alias]
+        return self
+
     def add_column(
         self,
         column_config: ColumnConfigT | None = None,
@@ -350,9 +403,14 @@ class DataDesignerConfigBuilder:
 
         Returns:
             The current Data Designer config object.
+
+        Raises:
+            BuilderConfigurationError: If any ToolConfig has duplicate tool names in its allow_tools list.
         """
+        self._validate_tool_configs_no_duplicates()
         return DataDesignerConfig(
             model_configs=self._model_configs,
+            tool_configs=self._tool_configs,
             seed_config=self._seed_config,
             columns=list(self._column_configs.values()),
             constraints=self._constraints or None,
@@ -360,6 +418,31 @@ class DataDesignerConfigBuilder:
             processors=self._processor_configs or None,
         )
 
+    def _validate_tool_configs_no_duplicates(self) -> None:
+        """Validate that no ToolConfig has duplicate tool names in its allow_tools list.
+
+        This is a static validation that catches obvious duplicates at config build time,
+        before providers are queried. Full validation (including duplicates across providers)
+        happens at resource provider creation time.
+
+        Raises:
+            BuilderConfigurationError: If any ToolConfig has duplicate tool names in allow_tools.
+        """
+        for tool_config in self._tool_configs:
+            if tool_config.allow_tools is None:
+                continue
+            seen: set[str] = set()
+            duplicates: list[str] = []
+            for tool_name in tool_config.allow_tools:
+                if tool_name in seen:
+                    duplicates.append(tool_name)
+                seen.add(tool_name)
+            if duplicates:
+                raise BuilderConfigurationError(
+                    f"🛑 ToolConfig '{tool_config.tool_alias}' has duplicate tool names in allow_tools: "
+                    f"{sorted(set(duplicates))!r}. Each tool name must be unique within a ToolConfig."
+                )
+
     def delete_constraints(self, target_column: str) -> Self:
         """Delete all constraints for the given target column.
 
@@ -411,6 +494,23 @@ class DataDesignerConfigBuilder:
         """
         return list(self._column_configs.values())
 
+    def get_tool_config(self, alias: str) -> ToolConfig:
+        """Get a tool configuration by alias.
+
+        Args:
+            alias: The alias of the tool configuration to retrieve.
+
+        Returns:
+            The tool configuration object.
+
+        Raises:
+            KeyError: If no tool configuration with the given alias exists.
+        """
+        for tc in self._tool_configs:
+            if tc.tool_alias == alias:
+                return tc
+        raise KeyError(f"No tool configuration with alias {alias!r} found")
+
     def get_constraints(self, target_column: str) -> list[ColumnConstraintT]:
         """Get all constraints for the given target column.
 

data_designer/config/custom_column.py

@@ -0,0 +1,64 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""User-facing utilities for custom column generation."""
+
+from __future__ import annotations
+
+import functools
+import inspect
+from typing import Any, Callable, TypeVar
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+# Expected parameter names by position (first param validated as row/df at runtime based on strategy)
+EXPECTED_PARAMS = ({"row", "df"}, {"generator_params"}, {"models"})
+
+
+def validate_generator_signature(fn: Callable[..., Any]) -> list[inspect.Parameter]:
+    """Validate generator function signature. Returns positional params if valid."""
+    params = [
+        p
+        for p in inspect.signature(fn).parameters.values()
+        if p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+    ]
+    n = len(params)
+    if n == 0 or n > 3:
+        raise TypeError(f"Generator '{fn.__name__}' must have 1-3 parameters, got {n}.")
+    for i, param in enumerate(params):
+        if param.name not in EXPECTED_PARAMS[i]:
+            expected = " or ".join(f"'{p}'" for p in sorted(EXPECTED_PARAMS[i]))
+            raise TypeError(f"Generator '{fn.__name__}' param {i + 1} must be {expected}, got '{param.name}'.")
+    return params
+
+
+def custom_column_generator(
+    required_columns: list[str] | None = None,
+    side_effect_columns: list[str] | None = None,
+    model_aliases: list[str] | None = None,
+) -> Callable[[F], F]:
+    """Decorator to define metadata for a custom column generator function.
+
+    Args:
+        required_columns: Columns that must exist before this column runs (DAG ordering).
+        side_effect_columns: Additional columns the function will create.
+        model_aliases: Model aliases to include in the `models` dict (required for LLM access).
+    """
+
+    def decorator(fn: F) -> F:
+        validate_generator_signature(fn)
+        fn.custom_column_metadata = {  # type: ignore[attr-defined]
+            "required_columns": required_columns or [],
+            "side_effect_columns": side_effect_columns or [],
+            "model_aliases": model_aliases or [],
+        }
+
+        @functools.wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            return fn(*args, **kwargs)
+
+        # Copy metadata to wrapper
+        wrapper.custom_column_metadata = fn.custom_column_metadata  # type: ignore[attr-defined]
+        return wrapper  # type: ignore[return-value]
+
+    return decorator

data_designer/config/data_designer_config.py

@@ -8,8 +8,9 @@ from typing import Annotated
 from pydantic import Field
 
 from data_designer.config.analysis.column_profilers import ColumnProfilerConfigT
-from data_designer.config.base import ExportableConfigBase
 from data_designer.config.column_types import ColumnConfigT
+from data_designer.config.exportable_config import ExportableConfigBase
+from data_designer.config.mcp import ToolConfig
 from data_designer.config.models import ModelConfig
 from data_designer.config.processors import ProcessorConfigT
 from data_designer.config.sampler_constraints import ColumnConstraintT
@@ -27,6 +28,8 @@ class DataDesignerConfig(ExportableConfigBase):
             should be generated. Must contain at least one column.
         model_configs: Optional list of model configurations for LLM-based generation.
             Each model config defines the model, provider, and inference parameters.
+        tool_configs: Optional list of tool configurations for MCP tool calling.
+            Each tool config defines the provider, allowed tools, and execution limits.
         seed_config: Optional seed dataset settings to use for generation.
         constraints: Optional list of column constraints.
         profilers: Optional list of column profilers for analyzing generated data characteristics.
@@ -34,6 +37,7 @@ class DataDesignerConfig(ExportableConfigBase):
 
     columns: list[Annotated[ColumnConfigT, Field(discriminator="column_type")]] = Field(min_length=1)
     model_configs: list[ModelConfig] | None = None
+    tool_configs: list[ToolConfig] | None = None
     seed_config: SeedConfig | None = None
     constraints: list[ColumnConstraintT] | None = None
     profilers: list[ColumnProfilerConfigT] | None = None

data_designer/config/exportable_config.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from data_designer.config.base import ConfigBase
+from data_designer.config.utils.io_helpers import serialize_data
+
+
+class ExportableConfigBase(ConfigBase):
+    def to_dict(self) -> dict[str, Any]:
+        """Convert the configuration to a dictionary.
+
+        Returns:
+            A dictionary representation of the configuration using JSON-compatible
+            serialization.
+        """
+        return self.model_dump(mode="json")
+
+    def to_yaml(self, path: str | Path | None = None, *, indent: int | None = 2, **kwargs) -> str | None:
+        """Convert the configuration to a YAML string or file.
+
+        Args:
+            path: Optional file path to write the YAML to. If None, returns the
+                YAML string instead of writing to file.
+            indent: Number of spaces for YAML indentation. Defaults to 2.
+            **kwargs: Additional keyword arguments passed to yaml.dump().
+
+        Returns:
+            The YAML string if path is None, otherwise None (file is written).
+        """
+        yaml_str = yaml.dump(self.to_dict(), indent=indent, **kwargs)
+        if path is None:
+            return yaml_str
+        with open(path, "w") as f:
+            f.write(yaml_str)
+
+    def to_json(self, path: str | Path | None = None, *, indent: int | None = 2, **kwargs) -> str | None:
+        """Convert the configuration to a JSON string or file.
+
+        Args:
+            path: Optional file path to write the JSON to. If None, returns the
+                JSON string instead of writing to file.
+            indent: Number of spaces for JSON indentation. Defaults to 2.
+            **kwargs: Additional keyword arguments passed to json.dumps().
+
+        Returns:
+            The JSON string if path is None, otherwise None (file is written).
+        """
+        json_str = serialize_data(self.to_dict(), indent=indent, **kwargs)
+        if path is None:
+            return json_str
+        with open(path, "w") as f:
+            f.write(json_str)