data-designer 0.3.8rc1__py3-none-any.whl → 0.4.0__py3-none-any.whl

data_designer/engine/column_generators/utils/prompt_renderer.py

@@ -1,100 +0,0 @@
-# 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

@@ -1,97 +0,0 @@
-# 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

@@ -1,71 +0,0 @@
-# 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

@@ -1,283 +0,0 @@
-# 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)