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

data_designer/engine/column_generators/generators/llm_completion.py

@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import functools
+import logging
+
+from data_designer.config.column_configs import (
+    LLMCodeColumnConfig,
+    LLMJudgeColumnConfig,
+    LLMStructuredColumnConfig,
+    LLMTextColumnConfig,
+)
+from data_designer.config.utils.constants import TRACE_COLUMN_POSTFIX
+from data_designer.engine.column_generators.generators.base import ColumnGeneratorWithModel, GenerationStrategy
+from data_designer.engine.column_generators.utils.prompt_renderer import (
+    PromptType,
+    RecordBasedPromptRenderer,
+    create_response_recipe,
+)
+from data_designer.engine.configurable_task import TaskConfigT
+from data_designer.engine.models.recipes.base import ResponseRecipe
+from data_designer.engine.processing.utils import deserialize_json_values
+
+logger = logging.getLogger(__name__)
+
+
+class ColumnGeneratorWithModelChatCompletion(ColumnGeneratorWithModel[TaskConfigT]):
+    @staticmethod
+    def get_generation_strategy() -> GenerationStrategy:
+        return GenerationStrategy.CELL_BY_CELL
+
+    @functools.cached_property
+    def response_recipe(self) -> ResponseRecipe:
+        return create_response_recipe(self.config, self.model_config)
+
+    @property
+    def max_conversation_correction_steps(self) -> int:
+        return self.resource_provider.run_config.max_conversation_correction_steps
+
+    @property
+    def max_conversation_restarts(self) -> int:
+        return self.resource_provider.run_config.max_conversation_restarts
+
+    @functools.cached_property
+    def prompt_renderer(self) -> RecordBasedPromptRenderer:
+        return RecordBasedPromptRenderer(
+            response_recipe=self.response_recipe,
+            error_message_context={
+                "column_name": self.config.name,
+                "column_type": self.config.column_type,
+                "model_alias": self.config.model_alias,
+            },
+        )
+
+    def generate(self, data: dict) -> dict:
+        # Deserialize input data from previous columns so Jinja2 templates can access nested fields
+        # Example: If prev column stored '{"key": "value"}', templates can use {{ prev_column.key }}
+        # Note: This creates a new dict and doesn't mutate the original `data` argument
+        deserialized_record = deserialize_json_values(data)
+
+        multi_modal_context = None
+        if self.config.multi_modal_context is not None and len(self.config.multi_modal_context) > 0:
+            multi_modal_context = []
+            for context in self.config.multi_modal_context:
+                multi_modal_context.extend(context.get_contexts(deserialized_record))
+
+        response, trace = self.model.generate(
+            prompt=self.prompt_renderer.render(
+                record=deserialized_record,
+                prompt_template=self.config.prompt,
+                prompt_type=PromptType.USER_PROMPT,
+            ),
+            system_prompt=self.prompt_renderer.render(
+                record=deserialized_record,
+                prompt_template=self.config.system_prompt,
+                prompt_type=PromptType.SYSTEM_PROMPT,
+            ),
+            parser=self.response_recipe.parse,
+            multi_modal_context=multi_modal_context,
+            max_correction_steps=self.max_conversation_correction_steps,
+            max_conversation_restarts=self.max_conversation_restarts,
+            purpose=f"running generation for column '{self.config.name}'",
+        )
+
+        serialized_output = self.response_recipe.serialize_output(response)
+        data[self.config.name] = self._process_serialized_output(serialized_output)
+
+        should_save_trace = (
+            self.config.with_trace or self.resource_provider.run_config.debug_override_save_all_column_traces
+        )
+        if should_save_trace:
+            data[self.config.name + TRACE_COLUMN_POSTFIX] = [message.to_dict() for message in trace]
+
+        return data
+
+    def _process_serialized_output(self, serialized_output: str) -> str | dict | list:
+        """Process the serialized output from the model. Subclasses can override to customize deserialization."""
+        return serialized_output
+
+
+class LLMTextCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMTextColumnConfig]): ...
+
+
+class LLMCodeCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMCodeColumnConfig]): ...
+
+
+class LLMStructuredCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMStructuredColumnConfig]):
+    def _process_serialized_output(self, serialized_output: str) -> dict | list:
+        return deserialize_json_values(serialized_output)
+
+
+class LLMJudgeCellGenerator(ColumnGeneratorWithModelChatCompletion[LLMJudgeColumnConfig]):
+    def _process_serialized_output(self, serialized_output: str) -> dict | list:
+        return deserialize_json_values(serialized_output)

data_designer/engine/column_generators/generators/samplers.py

@@ -0,0 +1,69 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+import random
+from functools import partial
+from typing import TYPE_CHECKING, Callable
+
+from data_designer.config.utils.constants import LOCALES_WITH_MANAGED_DATASETS
+from data_designer.engine.column_generators.generators.base import FromScratchColumnGenerator, GenerationStrategy
+from data_designer.engine.dataset_builders.multi_column_configs import SamplerMultiColumnConfig
+from data_designer.engine.processing.utils import concat_datasets
+from data_designer.engine.resources.managed_dataset_generator import ManagedDatasetGenerator
+from data_designer.engine.sampling_gen.data_sources.sources import SamplerType
+from data_designer.engine.sampling_gen.entities.person import load_person_data_sampler
+from data_designer.engine.sampling_gen.generator import DatasetGenerator as SamplingDatasetGenerator
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+logger = logging.getLogger(__name__)
+
+
+class SamplerColumnGenerator(FromScratchColumnGenerator[SamplerMultiColumnConfig]):
+    @staticmethod
+    def get_generation_strategy() -> GenerationStrategy:
+        return GenerationStrategy.FULL_COLUMN
+
+    def generate(self, data: pd.DataFrame) -> pd.DataFrame:
+        df_samplers = self.generate_from_scratch(len(data))
+        return concat_datasets([data, df_samplers])
+
+    def generate_from_scratch(self, num_records: int) -> pd.DataFrame:
+        sampling_generator = self._prepare_for_generation(num_records)
+        return sampling_generator.generate(num_records)
+
+    @property
+    def _needs_person_generator(self) -> bool:
+        columns = [c for c in self.config.columns if c.sampler_type == SamplerType.PERSON]
+        return any(c.params.locale in LOCALES_WITH_MANAGED_DATASETS for c in columns)
+
+    @property
+    def _person_generator_loader(self) -> Callable[[bool], ManagedDatasetGenerator]:
+        return partial(load_person_data_sampler, blob_storage=self.resource_provider.blob_storage)
+
+    def _create_sampling_dataset_generator(self) -> SamplingDatasetGenerator:
+        return SamplingDatasetGenerator(
+            sampler_columns=self.config,
+            person_generator_loader=(self._person_generator_loader if self._needs_person_generator else None),
+        )
+
+    def _log_person_generation_if_needed(self) -> None:
+        if self._needs_person_generator:
+            columns = [c for c in self.config.columns if c.sampler_type == SamplerType.PERSON]
+            emoji = random.choice(["🧑‍🎨", "🙋‍♂️", "🙋‍♀️", "🧑‍🚀", "👩‍🎤", "👨‍🍳", "👩‍🔬", "👨‍💻", "👩‍💼"])
+            log_msg = f"🎲 {emoji} Initializing person generation"
+            if any(c.params.with_synthetic_personas for c in columns):
+                log_msg += " ⚡️ with synthetic personas ⚡️"
+            logger.info(log_msg)
+
+    def _prepare_for_generation(self, num_records: int) -> SamplingDatasetGenerator:
+        logger.info(
+            f"🎲 Preparing samplers to generate {num_records} records across {len(self.config.columns)} columns"
+        )
+        self._log_person_generation_if_needed()
+        return self._create_sampling_dataset_generator()

data_designer/engine/column_generators/generators/seed_dataset.py

@@ -0,0 +1,144 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import functools
+import logging
+from typing import TYPE_CHECKING
+
+from data_designer.config.seed import IndexRange, PartitionBlock, SamplingStrategy
+from data_designer.engine.column_generators.generators.base import FromScratchColumnGenerator, GenerationStrategy
+from data_designer.engine.column_generators.utils.errors import SeedDatasetError
+from data_designer.engine.dataset_builders.multi_column_configs import SeedDatasetMultiColumnConfig
+from data_designer.engine.processing.utils import concat_datasets
+from data_designer.lazy_heavy_imports import duckdb, pd
+
+if TYPE_CHECKING:
+    import duckdb
+    import pandas as pd
+
+MAX_ZERO_RECORD_RESPONSE_FACTOR = 2
+
+logger = logging.getLogger(__name__)
+
+
+class SeedDatasetColumnGenerator(FromScratchColumnGenerator[SeedDatasetMultiColumnConfig]):
+    @staticmethod
+    def get_generation_strategy() -> GenerationStrategy:
+        return GenerationStrategy.FULL_COLUMN
+
+    @property
+    def num_records_sampled(self) -> int:
+        return self._num_records_sampled
+
+    @functools.cached_property
+    def duckdb_conn(self) -> duckdb.DuckDBPyConnection:
+        return self.resource_provider.seed_reader.create_duckdb_connection()
+
+    def generate(self, data: pd.DataFrame) -> pd.DataFrame:
+        return concat_datasets([self.generate_from_scratch(len(data)), data])
+
+    def generate_from_scratch(self, num_records: int) -> pd.DataFrame:
+        if num_records <= 0:
+            raise ValueError("🛑 `num_records` must be positive.")
+
+        if self._batch_reader is None:
+            self._reset_batch_reader(num_records)
+
+        return self._sample_records(num_records)
+
+    def _initialize(self) -> None:
+        self._num_records_sampled = 0
+        self._batch_reader = None
+        self._df_remaining = None
+        self._dataset_uri = self.resource_provider.seed_reader.get_dataset_uri()
+        self._seed_dataset_size = self.duckdb_conn.execute(f"SELECT COUNT(*) FROM '{self._dataset_uri}'").fetchone()[0]
+        self._index_range = self._resolve_index_range()
+
+    def _validate_selection_strategy(self) -> None:
+        err_msg = None
+        if self.config.selection_strategy is not None:
+            if (
+                isinstance(self.config.selection_strategy, IndexRange)
+                and self.config.selection_strategy.end >= self._seed_dataset_size
+            ):
+                err_msg = f"Selection strategy 'end' index {self.config.selection_strategy.end} is out of bounds for dataset size {self._seed_dataset_size}"
+            elif (
+                isinstance(self.config.selection_strategy, PartitionBlock)
+                and self.config.selection_strategy.num_partitions > self._seed_dataset_size
+            ):
+                err_msg = f"Selection strategy 'num_partitions' {self.config.selection_strategy.num_partitions} is out of bounds for dataset size {self._seed_dataset_size}"
+            if err_msg is not None:
+                raise SeedDatasetError(err_msg)
+
+    def _resolve_index_range(self) -> IndexRange | None:
+        self._validate_selection_strategy()
+        index_range = None
+        if self.config.selection_strategy is not None:
+            if isinstance(self.config.selection_strategy, IndexRange):
+                index_range = self.config.selection_strategy
+            elif isinstance(self.config.selection_strategy, PartitionBlock):
+                index_range = self.config.selection_strategy.to_index_range(self._seed_dataset_size)
+        return index_range
+
+    def _reset_batch_reader(self, num_records: int) -> None:
+        shuffle = self.config.sampling_strategy == SamplingStrategy.SHUFFLE
+        shuffle_query = " ORDER BY RANDOM()" if shuffle else ""
+
+        if self._index_range is not None:
+            # Use LIMIT and OFFSET for efficient index range filtering
+            # IndexRange uses 0-based indexing [start, end] inclusive
+            # OFFSET skips the first 'start' rows (0-based)
+            # LIMIT takes 'end - start + 1' rows to include both start and end (inclusive)
+            offset_value = self._index_range.start
+            limit_value = self._index_range.end - self._index_range.start + 1
+            read_query = f"""
+                SELECT * FROM '{self._dataset_uri}'
+                LIMIT {limit_value} OFFSET {offset_value}
+            """
+
+            read_query = f"SELECT * FROM ({read_query}){shuffle_query}"
+        else:
+            read_query = f"SELECT * FROM '{self._dataset_uri}'{shuffle_query}"
+        self._batch_reader = self.duckdb_conn.query(read_query).record_batch(batch_size=num_records)
+
+    def _sample_records(self, num_records: int) -> pd.DataFrame:
+        logger.info(f"🌱 Sampling {num_records} records from seed dataset")
+        logger.info(f"  |-- seed dataset size: {self._seed_dataset_size} records")
+        logger.info(f"  |-- sampling strategy: {self.config.sampling_strategy}")
+        if self._index_range is not None:
+            if isinstance(self.config.selection_strategy, IndexRange):
+                logger.info(f"  |-- selection: rows [{self._index_range.start} to {self._index_range.end}] inclusive")
+            else:
+                logger.info(
+                    f"  |-- selection: partition {self.config.selection_strategy.index + 1} of {self.config.selection_strategy.num_partitions}"
+                )
+            logger.info(f"  |-- seed dataset size after selection: {self._index_range.size} records")
+        df_batch = pd.DataFrame()
+        df_sample = pd.DataFrame() if self._df_remaining is None else self._df_remaining
+        num_zero_record_responses = 0
+
+        while len(df_sample) < num_records:
+            try:
+                df_batch = self._batch_reader.read_next_batch().to_pandas()
+                df_sample = pd.concat([df_sample, df_batch], ignore_index=True)
+            except StopIteration:
+                self._reset_batch_reader(num_records)
+
+            if len(df_batch) == 0:
+                num_zero_record_responses += 1
+                if num_zero_record_responses > MAX_ZERO_RECORD_RESPONSE_FACTOR * num_records:
+                    raise RuntimeError(
+                        "🛑 Something went wrong while reading from the datastore. "
+                        "Please check your connection and try again. "
+                        "If the issue persists, please contact support."
+                    )
+
+        self._df_remaining = None
+        if len(df_sample) > num_records:
+            self._df_remaining = df_sample.iloc[num_records:].reset_index(drop=True)
+        df_sample = df_sample.iloc[:num_records]
+        self._num_records_sampled += len(df_sample)
+
+        return df_sample

data_designer/engine/column_generators/generators/validation.py

@@ -0,0 +1,140 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from data_designer.config.column_configs import ValidationColumnConfig
+from data_designer.config.errors import InvalidConfigError
+from data_designer.config.utils.code_lang import SQL_DIALECTS, CodeLang
+from data_designer.config.validator_params import ValidatorParamsT, ValidatorType
+from data_designer.engine.column_generators.generators.base import ColumnGeneratorFullColumn
+from data_designer.engine.dataset_builders.utils.concurrency import ConcurrentThreadExecutor
+from data_designer.engine.errors import DataDesignerRuntimeError
+from data_designer.engine.validators import (
+    BaseValidator,
+    LocalCallableValidator,
+    PythonValidator,
+    RemoteValidator,
+    SQLValidator,
+    ValidationResult,
+)
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+logger = logging.getLogger(__name__)
+
+
+def get_validator_from_params(validator_type: ValidatorType, validator_params: ValidatorParamsT) -> BaseValidator:
+    if validator_type == ValidatorType.CODE:
+        if validator_params.code_lang == CodeLang.PYTHON:
+            return PythonValidator(validator_params)
+        elif validator_params.code_lang in SQL_DIALECTS:
+            return SQLValidator(validator_params)
+    elif validator_type == ValidatorType.REMOTE:
+        return RemoteValidator(validator_params)
+    else:
+        return LocalCallableValidator(validator_params)
+
+
+class ValidationColumnGenerator(ColumnGeneratorFullColumn[ValidationColumnConfig]):
+    def generate(self, data: pd.DataFrame) -> pd.DataFrame:
+        logger.info(f"🔍 Validating column {self.config.name!r} with {len(data)} records")
+        logger.info(f"  |-- target columns: {self.config.target_columns}")
+        logger.info(f"  |-- validator type: {self.config.validator_type}")
+        logger.info(f"  |-- validator params: {self.config.validator_params}")
+        logger.info(f"  |-- batch size: {self.config.batch_size}")
+
+        validator = get_validator_from_params(self.config.validator_type, self.config.validator_params)
+
+        # Check if the target columns are present in the dataset
+        missing_columns = set(self.config.target_columns) - set(data.columns)
+        if missing_columns:
+            raise InvalidConfigError(
+                f"Target columns {missing_columns} defined in validation column {self.config.name!r} are missing in dataset"
+            )
+
+        # Check whether to pass single columns or multiple columns to the validator
+        validate_columns_separately = False
+        if self.config.validator_type == ValidatorType.CODE and len(self.config.target_columns) > 1:
+            # Code validator expects single column input, so we validate each column separately
+            validate_columns_separately = True
+
+            columns_to_validate = [[col] for col in self.config.target_columns]
+        else:
+            columns_to_validate = [self.config.target_columns]
+
+        outputs_as_dicts = None
+        for cols in columns_to_validate:
+            # Filter the dataset to only include the target columns, and convert to a list of dictionaries
+            records = data[cols].to_dict(orient="records")
+
+            batched_records = [
+                records[batch_start : batch_start + self.config.batch_size]
+                for batch_start in range(0, len(records), self.config.batch_size)
+            ]
+
+            # Run validation in parallel or sequentially, depending on the validator type and parameters
+            if (
+                self.config.validator_type == ValidatorType.REMOTE
+                and self.config.validator_params.max_parallel_requests > 1
+            ):
+                concatenated_outputs = self._validate_in_parallel(validator, batched_records)
+            else:
+                concatenated_outputs = []
+                for batch in batched_records:
+                    concatenated_outputs.extend(self._validate_batch(validator, batch))
+
+            if validate_columns_separately:
+                if outputs_as_dicts is None:
+                    outputs_as_dicts = [{cols[0]: output.model_dump(mode="json")} for output in concatenated_outputs]
+                else:
+                    for dict_output in outputs_as_dicts:
+                        dict_output[cols[0]] = concatenated_outputs[0].model_dump(mode="json")
+            else:
+                outputs_as_dicts = [output.model_dump(mode="json") for output in concatenated_outputs]
+
+        validation_results = pd.DataFrame({self.config.name: outputs_as_dicts})
+        return pd.concat([data, validation_results], axis=1)
+
+    def _validate_in_parallel(self, validator: BaseValidator, batched_records: list[list[dict]]) -> pd.DataFrame:
+        """Run validation in parallel."""
+
+        outputs = [None] * len(batched_records)
+
+        def result_callback(result: ValidationResult, context: dict):
+            outputs[context["index"]] = result
+
+        def error_callback(error: Exception, context: dict):
+            outputs[context["index"]] = ValidationResult.empty(size=len(batched_records[context["index"]]))
+
+        settings = self.resource_provider.run_config
+        with ConcurrentThreadExecutor(
+            max_workers=self.config.validator_params.max_parallel_requests,
+            column_name=self.config.name,
+            result_callback=result_callback,
+            error_callback=error_callback,
+            shutdown_error_rate=settings.shutdown_error_rate,
+            shutdown_error_window=settings.shutdown_error_window,
+            disable_early_shutdown=settings.disable_early_shutdown,
+        ) as executor:
+            for i, batch in enumerate(batched_records):
+                executor.submit(lambda batch: self._validate_batch(validator, batch), batch, context={"index": i})
+
+        if any(output is None for output in outputs):
+            raise DataDesignerRuntimeError("Validation task failed due to an unexpected error in parallel execution")
+
+        # Concatenate the outputs and convert to a DataFrame
+        return sum([output.data for output in outputs], [])
+
+    def _validate_batch(self, validator: BaseValidator, batch: list[dict]) -> ValidationResult:
+        try:
+            return validator.run_validation(batch)
+        except Exception as e:
+            error_to_display = str(e).replace("\n", "\n  ")  # add spaces to improve readability
+            logger.error(f"Batch could not be validated:\n  {error_to_display}")
+            raise e

data_designer/engine/column_generators/registry.py

@@ -0,0 +1,60 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from data_designer.config.base import ConfigBase
+from data_designer.config.column_configs import (
+    EmbeddingColumnConfig,
+    ExpressionColumnConfig,
+    LLMCodeColumnConfig,
+    LLMJudgeColumnConfig,
+    LLMStructuredColumnConfig,
+    LLMTextColumnConfig,
+    ValidationColumnConfig,
+)
+from data_designer.config.column_types import DataDesignerColumnType
+from data_designer.engine.column_generators.generators.base import ColumnGenerator
+from data_designer.engine.column_generators.generators.embedding import EmbeddingCellGenerator
+from data_designer.engine.column_generators.generators.expression import ExpressionColumnGenerator
+from data_designer.engine.column_generators.generators.llm_completion import (
+    LLMCodeCellGenerator,
+    LLMJudgeCellGenerator,
+    LLMStructuredCellGenerator,
+    LLMTextCellGenerator,
+)
+from data_designer.engine.column_generators.generators.samplers import SamplerColumnGenerator
+from data_designer.engine.column_generators.generators.seed_dataset import SeedDatasetColumnGenerator
+from data_designer.engine.column_generators.generators.validation import ValidationColumnGenerator
+from data_designer.engine.dataset_builders.multi_column_configs import (
+    SamplerMultiColumnConfig,
+    SeedDatasetMultiColumnConfig,
+)
+from data_designer.engine.registry.base import TaskRegistry
+from data_designer.plugins.plugin import PluginType
+from data_designer.plugins.registry import PluginRegistry
+
+
+class ColumnGeneratorRegistry(TaskRegistry[DataDesignerColumnType, ColumnGenerator, ConfigBase]): ...
+
+
+def create_default_column_generator_registry(with_plugins: bool = True) -> ColumnGeneratorRegistry:
+    registry = ColumnGeneratorRegistry()
+    registry.register(DataDesignerColumnType.LLM_TEXT, LLMTextCellGenerator, LLMTextColumnConfig)
+    registry.register(DataDesignerColumnType.LLM_CODE, LLMCodeCellGenerator, LLMCodeColumnConfig)
+    registry.register(DataDesignerColumnType.LLM_JUDGE, LLMJudgeCellGenerator, LLMJudgeColumnConfig)
+    registry.register(DataDesignerColumnType.EXPRESSION, ExpressionColumnGenerator, ExpressionColumnConfig)
+    registry.register(DataDesignerColumnType.EMBEDDING, EmbeddingCellGenerator, EmbeddingColumnConfig)
+    registry.register(DataDesignerColumnType.SAMPLER, SamplerColumnGenerator, SamplerMultiColumnConfig)
+    registry.register(DataDesignerColumnType.SEED_DATASET, SeedDatasetColumnGenerator, SeedDatasetMultiColumnConfig)
+    registry.register(DataDesignerColumnType.VALIDATION, ValidationColumnGenerator, ValidationColumnConfig)
+    registry.register(DataDesignerColumnType.LLM_STRUCTURED, LLMStructuredCellGenerator, LLMStructuredColumnConfig)
+    if with_plugins:
+        for plugin in PluginRegistry().get_plugins(PluginType.COLUMN_GENERATOR):
+            registry.register(
+                DataDesignerColumnType(plugin.name),
+                plugin.impl_cls,
+                plugin.config_cls,
+            )
+
+    return registry

data_designer/engine/column_generators/utils/errors.py

@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from data_designer.engine.errors import DataDesignerError
+
+
+class PromptTemplateRenderError(DataDesignerError): ...
+
+
+class ExpressionTemplateRenderError(DataDesignerError): ...
+
+
+class SeedDatasetError(DataDesignerError): ...

data_designer/engine/column_generators/utils/generator_classification.py

@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from data_designer.config.column_types import DataDesignerColumnType
+from data_designer.config.utils.type_helpers import resolve_string_enum
+from data_designer.engine.column_generators.generators.base import ColumnGeneratorWithModelRegistry
+from data_designer.plugin_manager import PluginManager
+
+plugin_manager = PluginManager()
+
+
+def column_type_used_in_execution_dag(column_type: str | DataDesignerColumnType) -> bool:
+    """Return True if the column type is used in the workflow execution DAG."""
+    column_type = resolve_string_enum(column_type, DataDesignerColumnType)
+    dag_column_types = {
+        DataDesignerColumnType.EXPRESSION,
+        DataDesignerColumnType.LLM_CODE,
+        DataDesignerColumnType.LLM_JUDGE,
+        DataDesignerColumnType.LLM_STRUCTURED,
+        DataDesignerColumnType.LLM_TEXT,
+        DataDesignerColumnType.VALIDATION,
+        DataDesignerColumnType.EMBEDDING,
+    }
+    dag_column_types.update(plugin_manager.get_plugin_column_types(DataDesignerColumnType))
+    return column_type in dag_column_types
+
+
+def column_type_is_model_generated(column_type: str | DataDesignerColumnType) -> bool:
+    """Return True if the column type is a model-generated column."""
+    column_type = resolve_string_enum(column_type, DataDesignerColumnType)
+    model_generated_column_types = {
+        DataDesignerColumnType.LLM_TEXT,
+        DataDesignerColumnType.LLM_CODE,
+        DataDesignerColumnType.LLM_STRUCTURED,
+        DataDesignerColumnType.LLM_JUDGE,
+        DataDesignerColumnType.EMBEDDING,
+    }
+    for plugin in plugin_manager.get_column_generator_plugins():
+        if issubclass(plugin.impl_cls, ColumnGeneratorWithModelRegistry):
+            model_generated_column_types.add(plugin.name)
+    return column_type in model_generated_column_types

data_designer/engine/column_generators/utils/judge_score_factory.py

@@ -0,0 +1,58 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from enum import Enum
+
+from pydantic import BaseModel, ConfigDict, Field, create_model
+
+from data_designer.config.column_configs import Score
+
+SCORING_FORMAT = "* {score}: {description}"
+SCORE_FIELD_DESCRIPTION_FORMAT = "Score Descriptions for {enum_name}:\n{scoring}"
+
+
+class BaseJudgeResponse(BaseModel):
+    """Base model for all rubrics."""
+
+    model_config = ConfigDict(use_enum_values=True)
+    reasoning: str = Field(..., description="Reasoning for the assigned score.")
+
+
+def _stringify_scoring(options: dict, enum_type: type[Enum]) -> str:
+    """Convert score descriptions into a single text block."""
+    list_block = "\n".join(
+        [SCORING_FORMAT.format(score=score, description=description) for score, description in options.items()]
+    )
+    return SCORE_FIELD_DESCRIPTION_FORMAT.format(enum_name=enum_type.__name__, scoring=list_block)
+
+
+def create_judge_response_model(score: Score) -> type[BaseJudgeResponse]:
+    """Create a JudgeResponse data type."""
+    enum_members = {}
+    for option in score.options.keys():
+        member_name = f"VALUE_{option}"
+        enum_members[member_name] = option
+
+    DynamicScaleEnum = Enum(f"{score.name}Enum", enum_members)
+    options = _stringify_scoring(score.options, enum_type=DynamicScaleEnum)
+
+    return create_model(
+        score.name,
+        __doc__=score.description if score.description else None,
+        __base__=BaseJudgeResponse,
+        score=(DynamicScaleEnum, Field(..., description=options)),
+    )
+
+
+def create_judge_structured_output_model(
+    judge_responses: list[type[BaseJudgeResponse]],
+) -> type[BaseModel]:
+    """Create a JudgeStructuredOutput class dynamically."""
+    return create_model(
+        "JudgeStructuredOutput",
+        __doc__=f"Response schema for scores with the following names: {[response.__name__ for response in judge_responses]}.",
+        __base__=BaseModel,
+        **{response.__name__: (response, ...) for response in judge_responses},
+    )