data-designer-engine 0.4.0__py3-none-any.whl

data_designer/engine/column_generators/utils/prompt_renderer.py

@@ -0,0 +1,100 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import json
+import logging
+
+from data_designer.config.column_configs 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
+from data_designer.config.utils.misc import extract_keywords_from_jinja2_template
+from data_designer.config.utils.type_helpers import StrEnum
+from data_designer.engine.column_generators.utils.errors import PromptTemplateRenderError
+from data_designer.engine.column_generators.utils.judge_score_factory import (
+    create_judge_response_model,
+    create_judge_structured_output_model,
+)
+from data_designer.engine.models.recipes.base import ResponseRecipe
+from data_designer.engine.models.recipes.response_recipes import (
+    CodeResponseRecipe,
+    PydanticResponseRecipe,
+    StructuredResponseRecipe,
+    TextResponseRecipe,
+)
+from data_designer.engine.processing.ginja.environment import WithJinja2UserTemplateRendering
+from data_designer.engine.processing.ginja.exceptions import UserTemplateError, UserTemplateUnsupportedFiltersError
+
+logger = logging.getLogger(__name__)
+
+
+class PromptType(StrEnum):
+    SYSTEM_PROMPT = "system_prompt"
+    USER_PROMPT = "user_prompt"
+
+
+class RecordBasedPromptRenderer(WithJinja2UserTemplateRendering):
+    def __init__(self, response_recipe: ResponseRecipe, *, error_message_context: dict[str, str] | None = None):
+        self.response_recipe = response_recipe
+        self._error_message_context = error_message_context
+
+    def render(self, *, prompt_template: str | None, record: dict, prompt_type: PromptType) -> str | None:
+        self._prepare_environment(prompt_template=prompt_template, record=record, prompt_type=prompt_type)
+        rendered_prompt = self.render_multi_template(prompt_type, record) if prompt_template else ""
+        recipe_applicator = (
+            self.response_recipe.apply_recipe_to_user_prompt
+            if prompt_type == PromptType.USER_PROMPT
+            else self.response_recipe.apply_recipe_to_system_prompt
+        )
+        return recipe_applicator(rendered_prompt)
+
+    def _prepare_environment(self, *, prompt_template: str | None, record: dict, prompt_type: PromptType) -> None:
+        try:
+            self.prepare_jinja2_multi_template_renderer(
+                template_name=prompt_type.value,
+                prompt_template=prompt_template,
+                dataset_variables=list(record.keys()),
+            )
+        except (UserTemplateUnsupportedFiltersError, UserTemplateError) as exc:
+            template_variables = extract_keywords_from_jinja2_template(prompt_template)
+            missing_columns = list(set(template_variables) - set(record.keys()))
+
+            error_msg = (
+                f"There was an error preparing the {prompt_type.value.replace('_', ' ')} "
+                "template. Please double check that the template is valid Jinja2 syntax, that all "
+                "referenced variables are defined, and that any filters you are using are supported."
+            )
+            if len(missing_columns) > 0:
+                error_msg += f"\nThe following {missing_columns} columns are missing!"
+            if self._error_message_context is not None:
+                error_msg += f"\n{json.dumps(self._error_message_context, indent=2)}"
+            logger.error(f"🛑 {error_msg}")
+            raise PromptTemplateRenderError(f"{exc!s} {error_msg}")
+
+
+def create_response_recipe(
+    column_config: SingleColumnConfig, model_config: ModelConfig | None = None
+) -> ResponseRecipe:
+    if model_config and column_config.model_alias != model_config.alias:
+        raise ValueError(
+            f"Column config model alias {column_config.model_alias} does not match model config alias {model_config.alias}"
+        )
+    if column_config.column_type == DataDesignerColumnType.LLM_TEXT:
+        return TextResponseRecipe()
+    if column_config.column_type == DataDesignerColumnType.LLM_CODE:
+        return CodeResponseRecipe(
+            syntax=CodeLang.parse_lang(column_config.code_lang),
+        )
+    if column_config.column_type == DataDesignerColumnType.LLM_STRUCTURED:
+        return StructuredResponseRecipe(
+            json_schema=column_config.output_format,
+        )
+    if column_config.column_type == DataDesignerColumnType.LLM_JUDGE:
+        return PydanticResponseRecipe(
+            data_type=create_judge_structured_output_model(
+                [create_judge_response_model(s) for s in column_config.scores]
+            ),
+        )
+    raise ValueError(f"No response recipe found for column type: {column_config.column_type}")

data_designer/engine/compiler.py

@@ -0,0 +1,97 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+
+from data_designer.config.column_configs import SamplerColumnConfig, SeedDatasetColumnConfig
+from data_designer.config.data_designer_config import DataDesignerConfig
+from data_designer.config.errors import InvalidConfigError
+from data_designer.config.sampler_params import UUIDSamplerParams
+from data_designer.engine.resources.resource_provider import ResourceProvider
+from data_designer.engine.resources.seed_reader import SeedReader
+from data_designer.engine.validation import ViolationLevel, rich_print_violations, validate_data_designer_config
+
+logger = logging.getLogger(__name__)
+
+
+def compile_data_designer_config(config: DataDesignerConfig, resource_provider: ResourceProvider) -> DataDesignerConfig:
+    _resolve_and_add_seed_columns(config, resource_provider.seed_reader)
+    _add_internal_row_id_column_if_needed(config)
+    _validate(config)
+    return config
+
+
+def _resolve_and_add_seed_columns(config: DataDesignerConfig, seed_reader: SeedReader | None) -> None:
+    """Fetches the seed dataset column names, ensures there are no conflicts
+    with other columns, and adds seed column configs to the DataDesignerConfig.
+    """
+
+    if not seed_reader:
+        return
+
+    seed_col_names = seed_reader.get_column_names()
+    existing_columns = {column.name for column in config.columns}
+    colliding_columns = {name for name in seed_col_names if name in existing_columns}
+    if colliding_columns:
+        raise InvalidConfigError(
+            f"🛑 Seed dataset column(s) {colliding_columns} collide with existing column(s). "
+            "Please remove the conflicting columns or use a seed dataset with different column names."
+        )
+
+    config.columns.extend([SeedDatasetColumnConfig(name=col_name) for col_name in seed_col_names])
+
+
+def _add_internal_row_id_column_if_needed(config: DataDesignerConfig) -> None:
+    """Adds a UUID sampler column named '_internal_row_id' (set to drop) if needed to enable generation.
+
+    Generation requires either:
+    - At least one sampler column (which can generate data from scratch), OR
+    - A seed dataset (which provides initial data rows)
+
+    If neither exists, a UUID sampler column '_internal_row_id' is automatically added and marked for drop
+    to enable the generation process to start.
+
+    Args:
+        config: The DataDesigner configuration to potentially modify.
+    """
+    has_sampler_column = any(isinstance(col, SamplerColumnConfig) for col in config.columns)
+    has_seed_dataset_column = any(isinstance(col, SeedDatasetColumnConfig) for col in config.columns)
+
+    if not has_sampler_column and not has_seed_dataset_column:
+        logger.warning(
+            "🔔 No sampler column or seed dataset detected. Adding UUID column '_internal_row_id' (marked for drop) to enable generation."
+        )
+        id_column = SamplerColumnConfig(
+            name="_internal_row_id",
+            sampler_type="uuid",
+            params=UUIDSamplerParams(),
+            drop=True,
+        )
+        config.columns.insert(0, id_column)
+
+
+def _validate(config: DataDesignerConfig) -> None:
+    allowed_references = _get_allowed_references(config)
+    violations = validate_data_designer_config(
+        columns=config.columns,
+        processor_configs=config.processors or [],
+        allowed_references=allowed_references,
+    )
+    rich_print_violations(violations)
+    if len([v for v in violations if v.level == ViolationLevel.ERROR]) > 0:
+        raise InvalidConfigError(
+            "🛑 Your configuration contains validation errors. Please address the indicated issues and try again."
+        )
+    if len(violations) == 0:
+        logger.info("✅ Validation passed")
+
+
+def _get_allowed_references(config: DataDesignerConfig) -> list[str]:
+    refs = set[str]()
+    for column_config in config.columns:
+        refs.add(column_config.name)
+        for side_effect_column in column_config.side_effect_columns:
+            refs.add(side_effect_column)
+    return list(refs)

data_designer/engine/configurable_task.py

@@ -0,0 +1,71 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from abc import ABC
+from pathlib import Path
+from typing import TYPE_CHECKING, Generic, TypeVar, get_origin
+
+from data_designer.config.base import ConfigBase
+from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
+from data_designer.engine.resources.resource_provider import ResourceProvider
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+DataT = TypeVar("DataT", dict, pd.DataFrame)
+TaskConfigT = TypeVar("ConfigT", bound=ConfigBase)
+
+
+class ConfigurableTask(ABC, Generic[TaskConfigT]):
+    def __init__(self, config: TaskConfigT, resource_provider: ResourceProvider):
+        self._config = self.get_config_type().model_validate(config)
+        self._resource_provider = resource_provider
+        self._validate()
+        self._initialize()
+
+    @classmethod
+    def get_config_type(cls) -> type[TaskConfigT]:
+        for base in cls.__orig_bases__:
+            if hasattr(base, "__args__") and len(base.__args__) == 1:
+                arg = base.__args__[0]
+                origin = get_origin(arg) or arg
+                if isinstance(origin, type) and issubclass(origin, ConfigBase):
+                    return base.__args__[0]
+        raise TypeError(
+            f"Could not determine config type for `{cls.__name__}`. Please ensure that the "
+            "`ConfigurableTask` is defined with a generic type argument, where the type argument "
+            "is a subclass of `ConfigBase`."
+        )
+
+    @property
+    def artifact_path(self) -> Path:
+        return self.artifact_storage.artifact_path
+
+    @property
+    def artifact_storage(self) -> ArtifactStorage:
+        return self.resource_provider.artifact_storage
+
+    @property
+    def base_dataset_path(self) -> Path:
+        return self.artifact_storage.base_dataset_path
+
+    @property
+    def config(self) -> TaskConfigT:
+        return self._config
+
+    @property
+    def name(self) -> str:
+        return self.__class__.__name__
+
+    @property
+    def resource_provider(self) -> ResourceProvider:
+        return self._resource_provider
+
+    def _initialize(self) -> None:
+        """An internal method for custom initialization logic, which will be called in the constructor."""
+
+    def _validate(self) -> None:
+        """An internal method for custom validation logic, which will be called in the constructor."""

data_designer/engine/dataset_builders/artifact_storage.py

@@ -0,0 +1,283 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import json
+import logging
+import shutil
+from datetime import datetime
+from functools import cached_property
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, field_validator, model_validator
+
+from data_designer.config.utils.io_helpers import read_parquet_dataset
+from data_designer.config.utils.type_helpers import StrEnum, resolve_string_enum
+from data_designer.engine.dataset_builders.errors import ArtifactStorageError
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+logger = logging.getLogger(__name__)
+
+BATCH_FILE_NAME_FORMAT = "batch_{batch_number:05d}.parquet"
+SDG_CONFIG_FILENAME = "sdg.json"
+
+
+class BatchStage(StrEnum):
+    PARTIAL_RESULT = "partial_results_path"
+    FINAL_RESULT = "final_dataset_path"
+    DROPPED_COLUMNS = "dropped_columns_dataset_path"
+    PROCESSORS_OUTPUTS = "processors_outputs_path"
+
+
+class ArtifactStorage(BaseModel):
+    artifact_path: Path | str
+    dataset_name: str = "dataset"
+    final_dataset_folder_name: str = "parquet-files"
+    partial_results_folder_name: str = "tmp-partial-parquet-files"
+    dropped_columns_folder_name: str = "dropped-columns-parquet-files"
+    processors_outputs_folder_name: str = "processors-files"
+
+    @property
+    def artifact_path_exists(self) -> bool:
+        return self.artifact_path.exists()
+
+    @cached_property
+    def resolved_dataset_name(self) -> str:
+        dataset_path = self.artifact_path / self.dataset_name
+        if dataset_path.exists() and len(list(dataset_path.iterdir())) > 0:
+            new_dataset_name = f"{self.dataset_name}_{datetime.now().strftime('%m-%d-%Y_%H%M%S')}"
+            logger.info(
+                f"📂 Dataset path {str(dataset_path)!r} already exists. Dataset from this session"
+                f"\n\t\t     will be saved to {str(self.artifact_path / new_dataset_name)!r} instead."
+            )
+            return new_dataset_name
+        return self.dataset_name
+
+    @property
+    def base_dataset_path(self) -> Path:
+        return self.artifact_path / self.resolved_dataset_name
+
+    @property
+    def dropped_columns_dataset_path(self) -> Path:
+        return self.base_dataset_path / self.dropped_columns_folder_name
+
+    @property
+    def final_dataset_path(self) -> Path:
+        return self.base_dataset_path / self.final_dataset_folder_name
+
+    @property
+    def metadata_file_path(self) -> Path:
+        return self.base_dataset_path / "metadata.json"
+
+    @property
+    def partial_results_path(self) -> Path:
+        return self.base_dataset_path / self.partial_results_folder_name
+
+    @property
+    def processors_outputs_path(self) -> Path:
+        return self.base_dataset_path / self.processors_outputs_folder_name
+
+    @field_validator("artifact_path")
+    def validate_artifact_path(cls, v: Path | str) -> Path:
+        v = Path(v)
+        if not v.is_dir():
+            raise ArtifactStorageError("Artifact path must exist and be a directory")
+        return v
+
+    @model_validator(mode="after")
+    def validate_folder_names(self):
+        folder_names = [
+            self.dataset_name,
+            self.final_dataset_folder_name,
+            self.partial_results_folder_name,
+            self.dropped_columns_folder_name,
+            self.processors_outputs_folder_name,
+        ]
+
+        for name in folder_names:
+            if len(name) == 0:
+                raise ArtifactStorageError("🛑 Directory names must be non-empty strings.")
+
+        if len(set(folder_names)) != len(folder_names):
+            raise ArtifactStorageError("🛑 Folder names must be unique (no collisions allowed).")
+
+        invalid_chars = {"<", ">", ":", '"', "/", "\\", "|", "?", "*"}
+        for name in folder_names:
+            if any(char in invalid_chars for char in name):
+                raise ArtifactStorageError(f"🛑 Directory name '{name}' contains invalid characters.")
+
+        return self
+
+    @staticmethod
+    def mkdir_if_needed(path: Path | str) -> Path:
+        """Create the directory if it does not exist."""
+        path = Path(path)
+        if not path.exists():
+            logger.debug(f"📁 Creating directory: {path}")
+            path.mkdir(parents=True, exist_ok=True)
+        return path
+
+    @staticmethod
+    def read_parquet_files(path: Path) -> pd.DataFrame:
+        return read_parquet_dataset(path)
+
+    def create_batch_file_path(
+        self,
+        batch_number: int,
+        batch_stage: BatchStage,
+    ) -> Path:
+        if batch_number < 0:
+            raise ArtifactStorageError("🛑 Batch number must be non-negative.")
+        return self._get_stage_path(batch_stage) / BATCH_FILE_NAME_FORMAT.format(batch_number=batch_number)
+
+    def load_dataset(self, batch_stage: BatchStage = BatchStage.FINAL_RESULT) -> pd.DataFrame:
+        return read_parquet_dataset(self._get_stage_path(batch_stage))
+
+    def load_dataset_with_dropped_columns(self) -> pd.DataFrame:
+        # The pyarrow backend has better support for nested data types.
+        df = self.load_dataset()
+        if (
+            self.dropped_columns_dataset_path.exists()
+            and self.create_batch_file_path(0, BatchStage.DROPPED_COLUMNS).is_file()
+        ):
+            logger.debug("Concatenating dropped columns to the final dataset.")
+            df_dropped = self.load_dataset(batch_stage=BatchStage.DROPPED_COLUMNS)
+            if len(df_dropped) != len(df):
+                raise ArtifactStorageError(
+                    "🛑 The dropped-columns dataset has a different number of rows than the main dataset. "
+                    "Something unexpected must have happened to the dataset builder's artifacts."
+                )
+            # To ensure indexes are aligned and avoid silent misalignment (which would introduce NaNs),
+            # check that the indexes are identical before concatenation.
+            if not df.index.equals(df_dropped.index):
+                raise ArtifactStorageError(
+                    "🛑 The indexes of the main and dropped columns DataFrames are not aligned. "
+                    "Something unexpected must have happened to the dataset builder's artifacts."
+                )
+            df = pd.concat([df, df_dropped], axis=1)
+        return df
+
+    def move_partial_result_to_final_file_path(self, batch_number: int) -> Path:
+        partial_result_path = self.create_batch_file_path(batch_number, batch_stage=BatchStage.PARTIAL_RESULT)
+        if not partial_result_path.exists():
+            raise ArtifactStorageError("🛑 Partial result file not found.")
+        self.mkdir_if_needed(self._get_stage_path(BatchStage.FINAL_RESULT))
+        final_file_path = self.create_batch_file_path(batch_number, batch_stage=BatchStage.FINAL_RESULT)
+        shutil.move(partial_result_path, final_file_path)
+        return final_file_path
+
+    def write_batch_to_parquet_file(
+        self,
+        batch_number: int,
+        dataframe: pd.DataFrame,
+        batch_stage: BatchStage,
+        subfolder: str | None = None,
+    ) -> Path:
+        file_path = self.create_batch_file_path(batch_number, batch_stage=batch_stage)
+        self.write_parquet_file(file_path.name, dataframe, batch_stage, subfolder=subfolder)
+        return file_path
+
+    def write_parquet_file(
+        self,
+        parquet_file_name: str,
+        dataframe: pd.DataFrame,
+        batch_stage: BatchStage,
+        subfolder: str | None = None,
+    ) -> Path:
+        subfolder = subfolder or ""
+        self.mkdir_if_needed(self._get_stage_path(batch_stage) / subfolder)
+        file_path = self._get_stage_path(batch_stage) / subfolder / parquet_file_name
+        dataframe.to_parquet(file_path, index=False)
+        return file_path
+
+    def get_parquet_file_paths(self) -> list[str]:
+        """Get list of parquet file paths relative to base_dataset_path.
+
+        Returns:
+            List of relative paths to parquet files in the final dataset folder.
+        """
+        return [str(f.relative_to(self.base_dataset_path)) for f in sorted(self.final_dataset_path.glob("*.parquet"))]
+
+    def get_processor_file_paths(self) -> dict[str, list[str]]:
+        """Get processor output files organized by processor name.
+
+        Returns:
+            Dictionary mapping processor names to lists of relative file paths.
+        """
+        processor_files: dict[str, list[str]] = {}
+        if self.processors_outputs_path.exists():
+            for processor_dir in sorted(self.processors_outputs_path.iterdir()):
+                if processor_dir.is_dir():
+                    processor_name = processor_dir.name
+                    processor_files[processor_name] = [
+                        str(f.relative_to(self.base_dataset_path))
+                        for f in sorted(processor_dir.rglob("*"))
+                        if f.is_file()
+                    ]
+        return processor_files
+
+    def get_file_paths(self) -> dict[str, list[str] | dict[str, list[str]]]:
+        """Get all file paths organized by type.
+
+        Returns:
+            Dictionary with 'parquet-files' and 'processor-files' keys.
+        """
+        file_paths = {
+            "parquet-files": self.get_parquet_file_paths(),
+        }
+        processor_file_paths = self.get_processor_file_paths()
+        if processor_file_paths:
+            file_paths["processor-files"] = processor_file_paths
+
+        return file_paths
+
+    def read_metadata(self) -> dict:
+        """Read metadata from the metadata.json file.
+
+        Returns:
+            Dictionary containing the metadata.
+
+        Raises:
+            FileNotFoundError: If metadata file doesn't exist.
+        """
+        with open(self.metadata_file_path, "r") as file:
+            return json.load(file)
+
+    def write_metadata(self, metadata: dict) -> Path:
+        """Write metadata to the metadata.json file.
+
+        Args:
+            metadata: Dictionary containing metadata to write.
+
+        Returns:
+            Path to the written metadata file.
+        """
+        self.mkdir_if_needed(self.base_dataset_path)
+        with open(self.metadata_file_path, "w") as file:
+            json.dump(metadata, file, indent=4, sort_keys=True)
+        return self.metadata_file_path
+
+    def update_metadata(self, updates: dict) -> Path:
+        """Update existing metadata with new fields.
+
+        Args:
+            updates: Dictionary of fields to add/update in metadata.
+
+        Returns:
+            Path to the updated metadata file.
+        """
+        try:
+            existing_metadata = self.read_metadata()
+        except FileNotFoundError:
+            existing_metadata = {}
+
+        existing_metadata.update(updates)
+        return self.write_metadata(existing_metadata)
+
+    def _get_stage_path(self, stage: BatchStage) -> Path:
+        return getattr(self, resolve_string_enum(stage, BatchStage).value)