data-designer 0.1.0__py3-none-any.whl

data_designer/interface/data_designer.py

@@ -0,0 +1,351 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from pathlib import Path
+
+import pandas as pd
+
+from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
+from data_designer.config.config_builder import DataDesignerConfigBuilder
+from data_designer.config.default_model_settings import (
+    get_default_model_configs,
+    get_default_provider_name,
+    get_default_providers,
+    resolve_seed_default_model_settings,
+)
+from data_designer.config.interface import DataDesignerInterface
+from data_designer.config.models import (
+    ModelConfig,
+    ModelProvider,
+)
+from data_designer.config.preview_results import PreviewResults
+from data_designer.config.seed import LocalSeedDatasetReference
+from data_designer.config.utils.constants import (
+    DEFAULT_NUM_RECORDS,
+    MANAGED_ASSETS_PATH,
+    MODEL_CONFIGS_FILE_PATH,
+    MODEL_PROVIDERS_FILE_PATH,
+)
+from data_designer.config.utils.info import InterfaceInfo
+from data_designer.config.utils.io_helpers import write_seed_dataset
+from data_designer.config.utils.misc import can_run_data_designer_locally
+from data_designer.engine.analysis.dataset_profiler import (
+    DataDesignerDatasetProfiler,
+    DatasetProfilerConfig,
+)
+from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
+from data_designer.engine.dataset_builders.column_wise_builder import ColumnWiseDatasetBuilder
+from data_designer.engine.dataset_builders.utils.config_compiler import compile_dataset_builder_column_configs
+from data_designer.engine.model_provider import resolve_model_provider_registry
+from data_designer.engine.models.registry import create_model_registry
+from data_designer.engine.resources.managed_storage import init_managed_blob_storage
+from data_designer.engine.resources.resource_provider import ResourceProvider
+from data_designer.engine.resources.seed_dataset_data_store import (
+    HfHubSeedDatasetDataStore,
+    LocalSeedDatasetDataStore,
+)
+from data_designer.engine.secret_resolver import (
+    CompositeResolver,
+    EnvironmentResolver,
+    PlaintextResolver,
+    SecretResolver,
+)
+from data_designer.interface.errors import (
+    DataDesignerGenerationError,
+    DataDesignerProfilingError,
+    InvalidBufferValueError,
+)
+from data_designer.interface.results import DatasetCreationResults
+from data_designer.logging import RandomEmoji
+
+DEFAULT_BUFFER_SIZE = 1000
+
+logger = logging.getLogger(__name__)
+
+
+# Resolve default model settings on import to ensure they are available when the library is used.
+if can_run_data_designer_locally():
+    resolve_seed_default_model_settings()
+
+
+class DataDesigner(DataDesignerInterface[DatasetCreationResults]):
+    """Main interface for creating datasets with Data Designer.
+
+    This class provides the primary interface for building synthetic datasets using
+    Data Designer configurations. It manages model providers, artifact storage, and
+    orchestrates the dataset creation and profiling processes.
+
+    Args:
+        artifact_path: Path where generated artifacts will be stored.
+        dataset_name: Name for the generated dataset. Defaults to "dataset".
+            This will be used as the dataset folder name in the artifact path.
+        model_providers: Optional list of model providers for LLM generation. If None,
+            uses default providers.
+        secret_resolver: Resolver for handling secrets and credentials. Defaults to
+            EnvironmentResolver which reads secrets from environment variables.
+        managed_assets_path: Path to the managed assets directory. This is used to point
+            to the location of managed datasets and other assets used during dataset generation.
+            If not provided, will check for an environment variable called DATA_DESIGNER_MANAGED_ASSETS_PATH.
+            If the environment variable is not set, will use the default managed assets directory, which
+            is defined in `data_designer.config.utils.constants`.
+    """
+
+    def __init__(
+        self,
+        artifact_path: Path | str | None = None,
+        *,
+        model_providers: list[ModelProvider] | None = None,
+        secret_resolver: SecretResolver | None = None,
+        managed_assets_path: Path | str | None = None,
+    ):
+        self._secret_resolver = secret_resolver or CompositeResolver([EnvironmentResolver(), PlaintextResolver()])
+        self._artifact_path = Path(artifact_path) if artifact_path is not None else Path.cwd() / "artifacts"
+        self._buffer_size = DEFAULT_BUFFER_SIZE
+        self._managed_assets_path = Path(managed_assets_path or MANAGED_ASSETS_PATH)
+        self._model_providers = model_providers or self.get_default_model_providers()
+        self._model_provider_registry = resolve_model_provider_registry(
+            self._model_providers, get_default_provider_name()
+        )
+
+    @staticmethod
+    def make_seed_reference_from_file(file_path: str | Path) -> LocalSeedDatasetReference:
+        """Create a seed dataset reference from an existing file.
+
+        Supported file extensions: .parquet (recommended), .csv, .json, .jsonl
+
+        Args:
+            file_path: Path to an existing dataset file.
+
+        Returns:
+            A LocalSeedDatasetReference pointing to the specified file.
+        """
+        return LocalSeedDatasetReference(dataset=str(file_path))
+
+    @classmethod
+    def make_seed_reference_from_dataframe(
+        cls, dataframe: pd.DataFrame, file_path: str | Path
+    ) -> LocalSeedDatasetReference:
+        """Create a seed dataset reference from a pandas DataFrame.
+
+        This method writes the DataFrame to disk and returns a reference that can
+        be passed to the config builder's `with_seed_dataset` method. If the file
+        already exists, it will be overwritten.
+
+        Supported file extensions: .parquet (recommended), .csv, .json, .jsonl
+
+        Args:
+            dataframe: Pandas DataFrame to use as seed data.
+            file_path: Path where to save dataset.
+
+        Returns:
+            A LocalSeedDatasetReference pointing to the written file.
+        """
+        write_seed_dataset(dataframe, Path(file_path))
+        return cls.make_seed_reference_from_file(file_path)
+
+    @property
+    def info(self) -> InterfaceInfo:
+        """Get information about the Data Designer interface.
+
+        Returns:
+            InterfaceInfo object with information about the Data Designer interface.
+        """
+        return InterfaceInfo(model_providers=self._model_providers)
+
+    def create(
+        self,
+        config_builder: DataDesignerConfigBuilder,
+        *,
+        num_records: int = DEFAULT_NUM_RECORDS,
+        dataset_name: str = "dataset",
+    ) -> DatasetCreationResults:
+        """Create dataset and save results to the local artifact storage.
+
+        This method orchestrates the full dataset creation pipeline including building
+        the dataset according to the configuration, profiling the generated data, and
+        storing artifacts.
+
+        Args:
+            config_builder: The DataDesignerConfigBuilder containing the dataset
+                configuration (columns, constraints, seed data, etc.).
+            num_records: Number of records to generate.
+            dataset_name: Name of the dataset. This name will be used as the dataset
+                folder name in the artifact path directory.
+
+        Returns:
+            DatasetCreationResults object with methods for loading the generated dataset,
+            analysis results, and displaying sample records for inspection.
+
+        Raises:
+            DataDesignerGenerationError: If an error occurs during dataset generation.
+            DataDesignerProfilingError: If an error occurs during dataset profiling.
+        """
+        logger.info("🎨 Creating Data Designer dataset")
+
+        resource_provider = self._create_resource_provider(dataset_name, config_builder)
+
+        builder = self._create_dataset_builder(config_builder, resource_provider)
+
+        try:
+            builder.build(num_records=num_records, buffer_size=self._buffer_size)
+        except Exception as e:
+            raise DataDesignerGenerationError(f"🛑 Error generating dataset: {e}")
+
+        try:
+            profiler = self._create_dataset_profiler(config_builder, resource_provider)
+            analysis = profiler.profile_dataset(
+                num_records,
+                builder.artifact_storage.load_dataset_with_dropped_columns(),
+            )
+        except Exception as e:
+            raise DataDesignerProfilingError(f"🛑 Error profiling dataset: {e}")
+
+        return DatasetCreationResults(
+            artifact_storage=builder.artifact_storage,
+            analysis=analysis,
+            config_builder=config_builder,
+        )
+
+    def preview(
+        self, config_builder: DataDesignerConfigBuilder, *, num_records: int = DEFAULT_NUM_RECORDS
+    ) -> PreviewResults:
+        """Generate preview dataset for fast iteration on your Data Designer configuration.
+
+        All preview results are stored in memory. Once you are satisfied with the preview,
+        use the `create` method to generate data at a larger scale and save results to disk.
+
+        Args:
+            config_builder: The DataDesignerConfigBuilder containing the dataset
+                configuration (columns, constraints, seed data, etc.).
+            num_records: Number of records to generate.
+
+        Returns:
+            PreviewResults object with methods for inspecting the results.
+
+        Raises:
+            DataDesignerGenerationError: If an error occurs during preview dataset generation.
+            DataDesignerProfilingError: If an error occurs during preview dataset profiling.
+        """
+        logger.info(f"{RandomEmoji.previewing()} Preview generation in progress")
+
+        resource_provider = self._create_resource_provider("preview-dataset", config_builder)
+        builder = self._create_dataset_builder(config_builder, resource_provider)
+
+        try:
+            raw_dataset = builder.build_preview(num_records=num_records)
+            processed_dataset = builder.process_preview(raw_dataset)
+        except Exception as e:
+            raise DataDesignerGenerationError(f"🛑 Error generating preview dataset: {e}")
+
+        dropped_columns = raw_dataset.columns.difference(processed_dataset.columns)
+        if len(dropped_columns) > 0:
+            dataset_for_profiler = pd.concat([processed_dataset, raw_dataset[dropped_columns]], axis=1)
+        else:
+            dataset_for_profiler = processed_dataset
+
+        try:
+            profiler = self._create_dataset_profiler(config_builder, resource_provider)
+            analysis = profiler.profile_dataset(num_records, dataset_for_profiler)
+        except Exception as e:
+            raise DataDesignerProfilingError(f"🛑 Error profiling preview dataset: {e}")
+
+        if (
+            len(processed_dataset) > 0
+            and isinstance(analysis, DatasetProfilerResults)
+            and len(analysis.column_statistics) > 0
+        ):
+            logger.info(f"{RandomEmoji.success()} Preview complete!")
+
+        return PreviewResults(
+            dataset=processed_dataset,
+            analysis=analysis,
+            config_builder=config_builder,
+        )
+
+    def get_default_model_configs(self) -> list[ModelConfig]:
+        """Get the default model configurations.
+
+        Returns:
+            List of default model configurations.
+        """
+        logger.info(f"♻️ Using default model configs from {str(MODEL_CONFIGS_FILE_PATH)!r}")
+        return get_default_model_configs()
+
+    def get_default_model_providers(self) -> list[ModelProvider]:
+        """Get the default model providers.
+
+        Returns:
+            List of default model providers.
+        """
+        logger.info(f"♻️ Using default model providers from {str(MODEL_PROVIDERS_FILE_PATH)!r}")
+        return get_default_providers()
+
+    @property
+    def secret_resolver(self) -> SecretResolver:
+        """Get the secret resolver used by this DataDesigner instance.
+
+        Returns:
+            The SecretResolver instance handling credentials and secrets.
+        """
+        return self._secret_resolver
+
+    def set_buffer_size(self, buffer_size: int) -> None:
+        """Set the buffer size for dataset generation.
+
+        The buffer size controls how many records are processed in memory at once
+        during dataset generation using the `create` method. The default value is
+        set to the constant `DEFAULT_BUFFER_SIZE` defined in the data_designer module.
+
+        Args:
+            buffer_size: Number of records to process in each buffer.
+
+        Raises:
+            InvalidBufferValueError: If buffer size is less than or equal to 0.
+        """
+        if buffer_size <= 0:
+            raise InvalidBufferValueError("Buffer size must be greater than 0.")
+        self._buffer_size = buffer_size
+
+    def _create_dataset_builder(
+        self, config_builder: DataDesignerConfigBuilder, resource_provider: ResourceProvider
+    ) -> ColumnWiseDatasetBuilder:
+        return ColumnWiseDatasetBuilder(
+            column_configs=compile_dataset_builder_column_configs(config_builder.build(raise_exceptions=True)),
+            processor_configs=config_builder.get_processor_configs(),
+            resource_provider=resource_provider,
+        )
+
+    def _create_dataset_profiler(
+        self, config_builder: DataDesignerConfigBuilder, resource_provider: ResourceProvider
+    ) -> DataDesignerDatasetProfiler:
+        return DataDesignerDatasetProfiler(
+            config=DatasetProfilerConfig(
+                column_configs=config_builder.get_column_configs(),
+                column_profiler_configs=config_builder.get_profilers(),
+            ),
+            resource_provider=resource_provider,
+        )
+
+    def _create_resource_provider(
+        self, dataset_name: str, config_builder: DataDesignerConfigBuilder
+    ) -> ResourceProvider:
+        model_configs = config_builder.model_configs
+        ArtifactStorage.mkdir_if_needed(self._artifact_path)
+        return ResourceProvider(
+            artifact_storage=ArtifactStorage(artifact_path=self._artifact_path, dataset_name=dataset_name),
+            model_registry=create_model_registry(
+                model_configs=model_configs,
+                model_provider_registry=self._model_provider_registry,
+                secret_resolver=self._secret_resolver,
+            ),
+            blob_storage=init_managed_blob_storage(str(self._managed_assets_path)),
+            datastore=(
+                LocalSeedDatasetDataStore()
+                if (settings := config_builder.get_seed_datastore_settings()) is None
+                else HfHubSeedDatasetDataStore(
+                    endpoint=settings.endpoint,
+                    token=settings.token,
+                )
+            ),
+        )

data_designer/interface/errors.py

@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.errors import DataDesignerError
+
+
+class DataDesignerProfilingError(DataDesignerError):
+    """Raised for errors related to a Data Designer dataset profiling."""
+
+
+class DataDesignerGenerationError(DataDesignerError):
+    """Raised for errors related to a Data Designer dataset generation."""
+
+
+class InvalidBufferValueError(DataDesignerError):
+    """Raised for errors related to an invalid buffer value."""

data_designer/interface/results.py

@@ -0,0 +1,55 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import pandas as pd
+
+from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
+from data_designer.config.config_builder import DataDesignerConfigBuilder
+from data_designer.config.utils.visualization import WithRecordSamplerMixin
+from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
+
+
+class DatasetCreationResults(WithRecordSamplerMixin):
+    """Results container for a Data Designer dataset creation run.
+
+    This class provides access to the generated dataset, profiling analysis, and
+    visualization utilities. It is returned by the DataDesigner.create() method
+    and implements ResultsProtocol of the DataDesigner interface.
+    """
+
+    def __init__(
+        self,
+        *,
+        artifact_storage: ArtifactStorage,
+        analysis: DatasetProfilerResults,
+        config_builder: DataDesignerConfigBuilder,
+    ):
+        """Creates a new instance with results based on a dataset creation run.
+
+        Args:
+            artifact_storage: Storage manager for accessing generated artifacts.
+            analysis: Profiling results for the generated dataset.
+            config_builder: Configuration builder used to create the dataset.
+        """
+        self.artifact_storage = artifact_storage
+        self._analysis = analysis
+        self._config_builder = config_builder
+
+    def load_analysis(self) -> DatasetProfilerResults:
+        """Load the profiling analysis results for the generated dataset.
+
+        Returns:
+            DatasetProfilerResults containing statistical analysis and quality metrics
+                for each column in the generated dataset.
+        """
+        return self._analysis
+
+    def load_dataset(self) -> pd.DataFrame:
+        """Load the generated dataset as a pandas DataFrame.
+
+        Returns:
+            A pandas DataFrame containing the full generated dataset.
+        """
+        return self.artifact_storage.load_dataset()

data_designer/logging.py

@@ -0,0 +1,161 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+import logging
+from pathlib import Path
+import random
+import sys
+from typing import TextIO, Union
+
+from pythonjsonlogger import jsonlogger
+
+
+@dataclass
+class LoggerConfig:
+    name: str
+    level: str
+
+
+@dataclass
+class OutputConfig:
+    destination: Union[TextIO, Path]
+    structured: bool
+
+
+@dataclass
+class LoggingConfig:
+    logger_configs: list[LoggerConfig]
+    output_configs: list[OutputConfig]
+    root_level: str = "INFO"
+    to_silence: list[str] = field(default_factory=lambda: _DEFAULT_NOISY_LOGGERS)
+
+    @classmethod
+    def default(cls):
+        return LoggingConfig(
+            logger_configs=[LoggerConfig(name="data_designer", level="INFO")],
+            output_configs=[OutputConfig(destination=sys.stderr, structured=False)],
+        )
+
+    @classmethod
+    def debug(cls):
+        return LoggingConfig(
+            logger_configs=[LoggerConfig(name="data_designer", level="DEBUG")],
+            output_configs=[OutputConfig(destination=sys.stderr, structured=False)],
+        )
+
+
+class RandomEmoji:
+    """A generator for various themed emoji collections."""
+
+    @staticmethod
+    def cooking() -> str:
+        """Get a random cooking or food preparation emoji."""
+        return random.choice(["👨‍🍳", "👩‍🍳", "🍳", "🥘", "🍲", "🔪", "🥄", "🍴", "⏲️", "🥗"])
+
+    @staticmethod
+    def data() -> str:
+        """Get a random data or analytics emoji."""
+        return random.choice(["📊", "📈", "📉", "💾", "💿", "📀", "🗄️", "📁", "📂", "🗃️"])
+
+    @staticmethod
+    def generating() -> str:
+        """Get a random generating or creating emoji."""
+        return random.choice(["🏭", "⚙️", "🔨", "🛠️", "🏗️", "🎨", "✍️", "📝", "🔧", "⚒️"])
+
+    @staticmethod
+    def loading() -> str:
+        """Get a random loading or waiting emoji."""
+        return random.choice(["⏳", "⌛", "🔄", "♻️", "🔃", "⏰", "⏱️", "⏲️", "📡", "🌀"])
+
+    @staticmethod
+    def magic() -> str:
+        """Get a random magical or special effect emoji."""
+        return random.choice(["✨", "⭐", "🌟", "💫", "🪄", "🔮", "🎩", "🌈", "💎", "🦄"])
+
+    @staticmethod
+    def previewing() -> str:
+        """Get a random previewing or looking ahead emoji."""
+        return random.choice(["👀", "📺", "🔁", "👁️", "🔭", "🕵️", "🧐", "📸", "🎥", "🖼️"])
+
+    @staticmethod
+    def speed() -> str:
+        """Get a random speed or fast emoji."""
+        return random.choice(["⚡", "💨", "🏃", "🏎️", "🚄", "✈️", "💥", "⏩", "🏃‍♂️", "🏃‍♀️"])
+
+    @staticmethod
+    def start() -> str:
+        """Get a random emoji representing starting or launching something."""
+        return random.choice(["🚀", "▶️", "🎬", "🌅", "🏁", "🎯", "🚦", "🔔", "📣", "🎺"])
+
+    @staticmethod
+    def success() -> str:
+        """Get a random success or celebration emoji."""
+        return random.choice(["🎉", "🎊", "👏", "🙌", "🎆", "🍾", "☀️", "🏆", "✅", "🥳"])
+
+    @staticmethod
+    def thinking() -> str:
+        """Get a random thinking or processing emoji."""
+        return random.choice(["🤔", "💭", "🧠", "💡", "🔍", "🔎", "🤨", "🧐", "📝", "🧮"])
+
+    @staticmethod
+    def working() -> str:
+        """Get a random working or in-progress emoji."""
+        return random.choice(["⚙️", "🔧", "🔨", "⚒️", "🛠️", "💼", "👷", "🏗️", "🪛", "👨‍💻"])
+
+
+def configure_logging(config: LoggingConfig) -> None:
+    root_logger = logging.getLogger()
+
+    # Remove all handlers
+    root_logger.handlers.clear()
+
+    # Create and attach handler(s)
+    handlers = [_create_handler(output_config) for output_config in config.output_configs]
+    for handler in handlers:
+        root_logger.addHandler(handler)
+
+    # Set levels
+    root_logger.setLevel(config.root_level)
+    for logger_config in config.logger_configs:
+        logger = logging.getLogger(logger_config.name)
+        logger.setLevel(logger_config.level)
+
+    # Adjust noisy loggers
+    for name in config.to_silence:
+        quiet_noisy_logger(name)
+
+
+def quiet_noisy_logger(name: str) -> None:
+    logger = logging.getLogger(name)
+    logger.handlers.clear()
+    logger.setLevel(logging.WARNING)
+
+
+def _create_handler(output_config: OutputConfig) -> logging.Handler:
+    if isinstance(output_config.destination, Path):
+        handler = logging.FileHandler(str(output_config.destination))
+    else:
+        handler = logging.StreamHandler()
+
+    if output_config.structured:
+        formatter = _make_json_formatter()
+    else:
+        formatter = _make_stream_formatter()
+
+    handler.setFormatter(formatter)
+    return handler
+
+
+def _make_json_formatter() -> logging.Formatter:
+    log_format = "%(asctime)s %(levelname)s %(name)s %(message)s"
+    return jsonlogger.JsonFormatter(log_format)
+
+
+def _make_stream_formatter() -> logging.Formatter:
+    log_format = "[%(asctime)s] [%(levelname)s] %(message)s"
+    time_format = "%H:%M:%S"
+    return logging.Formatter(log_format, time_format)
+
+
+_DEFAULT_NOISY_LOGGERS = ["httpx", "matplotlib"]

data_designer/plugin_manager.py

@@ -0,0 +1,83 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import TYPE_CHECKING, Type, TypeAlias
+
+from .config.utils.misc import can_run_data_designer_locally
+
+if TYPE_CHECKING:
+    from data_designer.plugins.plugin import Plugin
+
+
+if can_run_data_designer_locally():
+    from data_designer.plugins.plugin import PluginType
+    from data_designer.plugins.registry import PluginRegistry
+
+
+class PluginManager:
+    def __init__(self):
+        if can_run_data_designer_locally():
+            self._plugins_supported = True
+            self._plugin_registry = PluginRegistry()
+        else:
+            self._plugins_supported = False
+            self._plugin_registry = None
+
+    def get_column_generator_plugins(self) -> list[Plugin]:
+        """Get all column generator plugins.
+
+        Returns:
+            A list of all column generator plugins.
+        """
+        return self._plugin_registry.get_plugins(PluginType.COLUMN_GENERATOR) if self._plugins_supported else []
+
+    def get_column_generator_plugin_if_exists(self, plugin_name: str) -> Plugin | None:
+        """Get a column generator plugin by name if it exists.
+
+        Args:
+            plugin_name: The name of the plugin to retrieve.
+
+        Returns:
+            The plugin if found, otherwise None.
+        """
+        if self._plugins_supported and self._plugin_registry.plugin_exists(plugin_name):
+            return self._plugin_registry.get_plugin(plugin_name)
+        return None
+
+    def get_plugin_column_types(self, enum_type: Type[Enum], required_resources: list[str] | None = None) -> list[Enum]:
+        """Get a list of plugin column types.
+
+        Args:
+            enum_type: The enum type to use for plugin entries.
+            required_resources: If provided, only return plugins with the required resources.
+
+        Returns:
+            A list of plugin column types.
+        """
+        type_list = []
+        if self._plugins_supported:
+            for plugin in self._plugin_registry.get_plugins(PluginType.COLUMN_GENERATOR):
+                if required_resources:
+                    task_required_resources = plugin.task_cls.metadata().required_resources or []
+                    if not all(resource in task_required_resources for resource in required_resources):
+                        continue
+                type_list.append(enum_type(plugin.name))
+        return type_list
+
+    def inject_into_column_config_type_union(self, column_config_type: Type[TypeAlias]) -> Type[TypeAlias]:
+        """Inject plugins into the column config type.
+
+        Args:
+            column_config_type: The column config type to inject plugins into.
+
+        Returns:
+            The column config type with plugins injected.
+        """
+        if self._plugins_supported:
+            column_config_type = self._plugin_registry.add_plugin_types_to_union(
+                column_config_type, PluginType.COLUMN_GENERATOR
+            )
+        return column_config_type

data_designer/plugins/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.plugins.plugin import Plugin, PluginType
+
+__all__ = ["Plugin", "PluginType"]

data_designer/plugins/errors.py

@@ -0,0 +1,10 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from data_designer.errors import DataDesignerError
+
+
+class PluginRegistrationError(DataDesignerError): ...
+
+
+class PluginNotFoundError(DataDesignerError): ...