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

data_designer/config/column_configs.py

@@ -1,470 +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, abstractmethod
-from typing import Annotated, Literal
-
-from pydantic import BaseModel, Discriminator, Field, model_validator
-from typing_extensions import Self
-
-from data_designer.config.base import ConfigBase
-from data_designer.config.errors import InvalidConfigError
-from data_designer.config.models import ImageContext
-from data_designer.config.sampler_params import SamplerParamsT, SamplerType
-from data_designer.config.utils.code_lang import CodeLang
-from data_designer.config.utils.constants import REASONING_TRACE_COLUMN_POSTFIX
-from data_designer.config.utils.misc import assert_valid_jinja2_template, extract_keywords_from_jinja2_template
-from data_designer.config.validator_params import ValidatorParamsT, ValidatorType
-
-
-class SingleColumnConfig(ConfigBase, ABC):
-    """Abstract base class for all single-column configuration types.
-
-    This class serves as the foundation for all column configurations in DataDesigner,
-    defining shared fields and properties across all column types.
-
-    Attributes:
-        name: Unique name of the column to be generated.
-        drop: If True, the column will be generated but removed from the final dataset.
-            Useful for intermediate columns that are dependencies for other columns.
-        column_type: Discriminator field that identifies the specific column type.
-            Subclasses must override this field to specify the column type with a `Literal` value.
-    """
-
-    name: str
-    drop: bool = False
-    column_type: str
-
-    @staticmethod
-    def get_column_emoji() -> str:
-        return "🎨"
-
-    @property
-    @abstractmethod
-    def required_columns(self) -> list[str]:
-        """Returns a list of column names that must exist before this column can be generated.
-
-        Returns:
-            List of column names that this column depends on. Empty list indicates
-            no dependencies. Override in subclasses to specify dependencies.
-        """
-
-    @property
-    @abstractmethod
-    def side_effect_columns(self) -> list[str]:
-        """Returns a list of additional columns that this column will create as a side effect.
-
-        Some column types generate additional metadata or auxiliary columns alongside
-        the primary column (e.g., reasoning traces for LLM columns).
-
-        Returns:
-            List of column names that this column will create as a side effect. Empty list
-            indicates no side effect columns. Override in subclasses to specify side effects.
-        """
-
-
-class SamplerColumnConfig(SingleColumnConfig):
-    """Configuration for columns generated using numerical samplers.
-
-    Sampler columns provide efficient data generation using numerical samplers for
-    common data types and distributions. Supported samplers include UUID generation,
-    datetime/timedelta sampling, person generation, category / subcategory sampling,
-    and various statistical distributions (uniform, gaussian, binomial, poisson, scipy).
-
-    Attributes:
-        sampler_type: Type of sampler to use. Available types include:
-            "uuid", "category", "subcategory", "uniform", "gaussian", "bernoulli",
-            "bernoulli_mixture", "binomial", "poisson", "scipy", "person", "datetime", "timedelta".
-        params: Parameters specific to the chosen sampler type. Type varies based on the `sampler_type`
-            (e.g., `CategorySamplerParams`, `UniformSamplerParams`, `PersonSamplerParams`).
-        conditional_params: Optional dictionary for conditional parameters. The dict keys
-            are the conditions that must be met (e.g., "age > 21") for the conditional parameters
-            to be used. The values of dict are the parameters to use when the condition is met.
-        convert_to: Optional type conversion to apply after sampling. Must be one of "float", "int", or "str".
-            Useful for converting numerical samples to strings or other types.
-        column_type: Discriminator field, always "sampler" for this configuration type.
-
-    !!! tip "Displaying available samplers and their parameters"
-        The config builder has an `info` attribute that can be used to display the
-        available samplers and their parameters:
-        ```python
-        config_builder.info.display("samplers")
-        ```
-    """
-
-    sampler_type: SamplerType
-    params: Annotated[SamplerParamsT, Discriminator("sampler_type")]
-    conditional_params: dict[str, Annotated[SamplerParamsT, Discriminator("sampler_type")]] = {}
-    convert_to: str | None = None
-    column_type: Literal["sampler"] = "sampler"
-
-    @staticmethod
-    def get_column_emoji() -> str:
-        return "🎲"
-
-    @property
-    def required_columns(self) -> list[str]:
-        return []
-
-    @property
-    def side_effect_columns(self) -> list[str]:
-        return []
-
-    @model_validator(mode="before")
-    @classmethod
-    def inject_sampler_type_into_params(cls, data: dict) -> dict:
-        """Inject sampler_type into params dict to enable discriminated union resolution.
-
-        This allows users to pass params as a simple dict without the sampler_type field,
-        which will be automatically added based on the outer sampler_type field.
-        """
-        if isinstance(data, dict):
-            sampler_type = data.get("sampler_type")
-            params = data.get("params")
-
-            # If params is a dict and doesn't have sampler_type, inject it
-            if sampler_type and isinstance(params, dict) and "sampler_type" not in params:
-                data["params"] = {"sampler_type": sampler_type, **params}
-
-            # Handle conditional_params similarly
-            conditional_params = data.get("conditional_params")
-            if conditional_params and isinstance(conditional_params, dict):
-                for condition, cond_params in conditional_params.items():
-                    if isinstance(cond_params, dict) and "sampler_type" not in cond_params:
-                        data["conditional_params"][condition] = {"sampler_type": sampler_type, **cond_params}
-
-        return data
-
-
-class LLMTextColumnConfig(SingleColumnConfig):
-    """Configuration for text generation columns using Large Language Models.
-
-    LLM text columns generate free-form text content using language models via LiteLLM.
-    Prompts support Jinja2 templating to reference values from other columns, enabling
-    context-aware generation. The generated text can optionally include reasoning traces
-    when models support extended thinking.
-
-    Attributes:
-        prompt: Prompt template for text generation. Supports Jinja2 syntax to
-            reference other columns (e.g., "Write a story about {{ character_name }}").
-            Must be a valid Jinja2 template.
-        model_alias: Alias of the model configuration to use for generation.
-            Must match a model alias defined when initializing the DataDesignerConfigBuilder.
-        system_prompt: Optional system prompt to set model behavior and constraints.
-            Also supports Jinja2 templating. If provided, must be a valid Jinja2 template.
-            Do not put any output parsing instructions in the system prompt. Instead,
-            use the appropriate column type for the output you want to generate - e.g.,
-            `LLMStructuredColumnConfig` for structured output, `LLMCodeColumnConfig` for code.
-        multi_modal_context: Optional list of image contexts for multi-modal generation.
-            Enables vision-capable models to generate text based on image inputs.
-        column_type: Discriminator field, always "llm-text" for this configuration type.
-    """
-
-    prompt: str
-    model_alias: str
-    system_prompt: str | None = None
-    multi_modal_context: list[ImageContext] | None = None
-    column_type: Literal["llm-text"] = "llm-text"
-
-    @staticmethod
-    def get_column_emoji() -> str:
-        return "📝"
-
-    @property
-    def required_columns(self) -> list[str]:
-        """Get columns referenced in the prompt and system_prompt templates.
-
-        Returns:
-            List of unique column names referenced in Jinja2 templates.
-        """
-        required_cols = list(extract_keywords_from_jinja2_template(self.prompt))
-        if self.system_prompt:
-            required_cols.extend(list(extract_keywords_from_jinja2_template(self.system_prompt)))
-        return list(set(required_cols))
-
-    @property
-    def side_effect_columns(self) -> list[str]:
-        """Returns the reasoning trace column, which may be generated alongside the main column.
-
-        Reasoning traces are only returned if the served model parses and returns reasoning content.
-
-        Returns:
-            List containing the reasoning trace column name.
-        """
-        return [f"{self.name}{REASONING_TRACE_COLUMN_POSTFIX}"]
-
-    @model_validator(mode="after")
-    def assert_prompt_valid_jinja(self) -> Self:
-        """Validate that prompt and system_prompt are valid Jinja2 templates.
-
-        Returns:
-            The validated instance.
-
-        Raises:
-            InvalidConfigError: If prompt or system_prompt contains invalid Jinja2 syntax.
-        """
-        assert_valid_jinja2_template(self.prompt)
-        if self.system_prompt:
-            assert_valid_jinja2_template(self.system_prompt)
-        return self
-
-
-class LLMCodeColumnConfig(LLMTextColumnConfig):
-    """Configuration for code generation columns using Large Language Models.
-
-    Extends LLMTextColumnConfig to generate code snippets in specific programming languages
-    or SQL dialects. The generated code is automatically extracted from markdown code blocks
-    for the specified language. Inherits all prompt templating capabilities.
-
-    Attributes:
-        code_lang: Programming language or SQL dialect for code generation. Supported
-            values include: "python", "javascript", "typescript", "java", "kotlin", "go",
-            "rust", "ruby", "scala", "swift", "sql:sqlite", "sql:postgres", "sql:mysql",
-            "sql:tsql", "sql:bigquery", "sql:ansi". See CodeLang enum for complete list.
-        column_type: Discriminator field, always "llm-code" for this configuration type.
-    """
-
-    code_lang: CodeLang
-    column_type: Literal["llm-code"] = "llm-code"
-
-    @staticmethod
-    def get_column_emoji() -> str:
-        return "💻"
-
-
-class LLMStructuredColumnConfig(LLMTextColumnConfig):
-    """Configuration for structured JSON generation columns using Large Language Models.
-
-    Extends LLMTextColumnConfig to generate structured data conforming to a specified schema.
-    Uses JSON schema or Pydantic models to define the expected output structure, enabling
-    type-safe and validated structured output generation. Inherits prompt templating capabilities.
-
-    Attributes:
-        output_format: The schema defining the expected output structure. Can be either:
-            - A Pydantic BaseModel class (recommended)
-            - A JSON schema dictionary
-        column_type: Discriminator field, always "llm-structured" for this configuration type.
-    """
-
-    output_format: dict | type[BaseModel]
-    column_type: Literal["llm-structured"] = "llm-structured"
-
-    @staticmethod
-    def get_column_emoji() -> str:
-        return "🗂️"
-
-    @model_validator(mode="after")
-    def validate_output_format(self) -> Self:
-        """Convert Pydantic model to JSON schema if needed.
-
-        Returns:
-            The validated instance with output_format as a JSON schema dict.
-        """
-        if not isinstance(self.output_format, dict) and issubclass(self.output_format, BaseModel):
-            self.output_format = self.output_format.model_json_schema()
-        return self
-
-
-class Score(ConfigBase):
-    """Configuration for a "score" in an LLM judge evaluation.
-
-    Defines a single scoring criterion with its possible values and descriptions. Multiple
-    Score objects can be combined in an LLMJudgeColumnConfig to create multi-dimensional
-    quality assessments.
-
-    Attributes:
-        name: A clear, concise name for this scoring dimension (e.g., "Relevance", "Fluency").
-        description: An informative and detailed assessment guide explaining how to evaluate
-            this dimension. Should provide clear criteria for scoring.
-        options: Dictionary mapping score values to their descriptions. Keys can be integers
-            (e.g., 1-5 scale) or strings (e.g., "Poor", "Good", "Excellent"). Values are
-            descriptions explaining what each score level means.
-    """
-
-    name: str = Field(..., description="A clear name for this score.")
-    description: str = Field(..., description="An informative and detailed assessment guide for using this score.")
-    options: dict[int | str, str] = Field(..., description="Score options in the format of {score: description}.")
-
-
-class LLMJudgeColumnConfig(LLMTextColumnConfig):
-    """Configuration for LLM-as-a-judge quality assessment and scoring columns.
-
-    Extends LLMTextColumnConfig to create judge columns that evaluate and score other
-    generated content based on the defined criteria. Useful for quality assessment, preference
-    ranking, and multi-dimensional evaluation of generated data.
-
-    Attributes:
-        scores: List of Score objects defining the evaluation dimensions. Each score
-            represents a different aspect to evaluate (e.g., accuracy, relevance, fluency).
-            Must contain at least one score.
-        column_type: Discriminator field, always "llm-judge" for this configuration type.
-    """
-
-    scores: list[Score] = Field(..., min_length=1)
-    column_type: Literal["llm-judge"] = "llm-judge"
-
-    @staticmethod
-    def get_column_emoji() -> str:
-        return "⚖️"
-
-
-class ExpressionColumnConfig(SingleColumnConfig):
-    """Configuration for derived columns using Jinja2 expressions.
-
-    Expression columns compute values by evaluating Jinja2 templates that reference other
-    columns. Useful for transformations, concatenations, conditional logic, and derived
-    features without requiring LLM generation. The expression is evaluated row-by-row.
-
-    Attributes:
-        expr: Jinja2 expression to evaluate. Can reference other column values using
-            {{ column_name }} syntax. Supports filters, conditionals, and arithmetic.
-            Must be a valid, non-empty Jinja2 template.
-        dtype: Data type to cast the result to. Must be one of "int", "float", "str", or "bool".
-            Defaults to "str". Type conversion is applied after expression evaluation.
-        column_type: Discriminator field, always "expression" for this configuration type.
-    """
-
-    name: str
-    expr: str
-    dtype: Literal["int", "float", "str", "bool"] = "str"
-    column_type: Literal["expression"] = "expression"
-
-    @staticmethod
-    def get_column_emoji() -> str:
-        return "🧩"
-
-    @property
-    def required_columns(self) -> list[str]:
-        """Returns the columns referenced in the expression template."""
-        return list(extract_keywords_from_jinja2_template(self.expr))
-
-    @property
-    def side_effect_columns(self) -> list[str]:
-        return []
-
-    @model_validator(mode="after")
-    def assert_expression_valid_jinja(self) -> Self:
-        """Validate that the expression is a valid, non-empty Jinja2 template.
-
-        Returns:
-            The validated instance.
-
-        Raises:
-            InvalidConfigError: If expression is empty or contains invalid Jinja2 syntax.
-        """
-        if not self.expr.strip():
-            raise InvalidConfigError(
-                f"🛑 Expression column '{self.name}' has an empty or whitespace-only expression. "
-                f"Please provide a valid Jinja2 expression (e.g., '{{ column_name }}' or '{{ col1 }} + {{ col2 }}') "
-                "or remove this column if not needed."
-            )
-        assert_valid_jinja2_template(self.expr)
-        return self
-
-
-class ValidationColumnConfig(SingleColumnConfig):
-    """Configuration for validation columns that validate existing columns.
-
-    Validation columns execute validation logic against specified target columns and return
-    structured results indicating pass/fail status with validation details. Supports multiple
-    validation strategies: code execution (Python/SQL), local callable functions (library only),
-    and remote HTTP endpoints.
-
-    Attributes:
-        target_columns: List of column names to validate. These columns are passed to the
-            validator for validation. All target columns must exist in the dataset
-            before validation runs.
-        validator_type: The type of validator to use. Options:
-            - "code": Execute code (Python or SQL) for validation. The code receives a
-              DataFrame with target columns and must return a DataFrame with validation results.
-            - "local_callable": Call a local Python function with the data. Only supported
-              when running DataDesigner locally.
-            - "remote": Send data to a remote HTTP endpoint for validation. Useful for
-        validator_params: Parameters specific to the validator type. Type varies by validator:
-            - CodeValidatorParams: Specifies code language (python or SQL dialect like
-              "sql:postgres", "sql:mysql").
-            - LocalCallableValidatorParams: Provides validation function (Callable[[pd.DataFrame],
-              pd.DataFrame]) and optional output schema for validation results.
-            - RemoteValidatorParams: Configures endpoint URL, HTTP timeout, retry behavior
-              (max_retries, retry_backoff), and parallel request limits (max_parallel_requests).
-        batch_size: Number of records to process in each validation batch. Defaults to 10.
-            Larger batches are more efficient but use more memory. Adjust based on validator
-            complexity and available resources.
-        column_type: Discriminator field, always "validation" for this configuration type.
-    """
-
-    target_columns: list[str]
-    validator_type: ValidatorType
-    validator_params: ValidatorParamsT
-    batch_size: int = Field(default=10, ge=1, description="Number of records to process in each batch")
-    column_type: Literal["validation"] = "validation"
-
-    @staticmethod
-    def get_column_emoji() -> str:
-        return "🔍"
-
-    @property
-    def required_columns(self) -> list[str]:
-        """Returns the columns that need to be validated."""
-        return self.target_columns
-
-    @property
-    def side_effect_columns(self) -> list[str]:
-        return []
-
-
-class SeedDatasetColumnConfig(SingleColumnConfig):
-    """Configuration for columns sourced from seed datasets.
-
-    This config marks columns that come from seed data. It is typically created
-    automatically when calling `with_seed_dataset()` on the builder, rather than
-    being instantiated directly by users.
-
-    Attributes:
-        column_type: Discriminator field, always "seed-dataset" for this configuration type.
-    """
-
-    column_type: Literal["seed-dataset"] = "seed-dataset"
-
-    @staticmethod
-    def get_column_emoji() -> str:
-        return "🌱"
-
-    @property
-    def required_columns(self) -> list[str]:
-        return []
-
-    @property
-    def side_effect_columns(self) -> list[str]:
-        return []
-
-
-class EmbeddingColumnConfig(SingleColumnConfig):
-    """Configuration for embedding generation columns.
-
-    Embedding columns generate embeddings for text input using a specified model.
-
-    Attributes:
-        target_column: The column to generate embeddings for. The column could be a single text string or a list of text strings in stringified JSON format.
-            If it is a list of text strings in stringified JSON format, the embeddings will be generated for each text string.
-        model_alias: The model to use for embedding generation.
-        column_type: Discriminator field, always "embedding" for this configuration type.
-    """
-
-    target_column: str
-    model_alias: str
-    column_type: Literal["embedding"] = "embedding"
-
-    @staticmethod
-    def get_column_emoji() -> str:
-        return "🧬"
-
-    @property
-    def required_columns(self) -> list[str]:
-        return [self.target_column]
-
-    @property
-    def side_effect_columns(self) -> list[str]:
-        return []

data_designer/config/column_types.py

@@ -1,141 +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 typing_extensions import TypeAlias
-
-from data_designer.config.column_configs import (
-    EmbeddingColumnConfig,
-    ExpressionColumnConfig,
-    LLMCodeColumnConfig,
-    LLMJudgeColumnConfig,
-    LLMStructuredColumnConfig,
-    LLMTextColumnConfig,
-    SamplerColumnConfig,
-    SeedDatasetColumnConfig,
-    ValidationColumnConfig,
-)
-from data_designer.config.errors import InvalidConfigError
-from data_designer.config.sampler_params import SamplerType
-from data_designer.config.utils.type_helpers import (
-    SAMPLER_PARAMS,
-    create_str_enum_from_discriminated_type_union,
-    resolve_string_enum,
-)
-from data_designer.plugin_manager import PluginManager
-
-plugin_manager = PluginManager()
-
-ColumnConfigT: TypeAlias = (
-    ExpressionColumnConfig
-    | LLMCodeColumnConfig
-    | LLMJudgeColumnConfig
-    | LLMStructuredColumnConfig
-    | LLMTextColumnConfig
-    | SamplerColumnConfig
-    | SeedDatasetColumnConfig
-    | ValidationColumnConfig
-    | EmbeddingColumnConfig
-)
-ColumnConfigT = plugin_manager.inject_into_column_config_type_union(ColumnConfigT)
-
-DataDesignerColumnType = create_str_enum_from_discriminated_type_union(
-    enum_name="DataDesignerColumnType",
-    type_union=ColumnConfigT,
-    discriminator_field_name="column_type",
-)
-
-
-def get_column_config_from_kwargs(name: str, column_type: DataDesignerColumnType, **kwargs) -> ColumnConfigT:
-    """Create a Data Designer column config object from kwargs.
-
-    Args:
-        name: Name of the column.
-        column_type: Type of the column.
-        **kwargs: Keyword arguments to pass to the column constructor.
-
-    Returns:
-        Data Designer column object of the appropriate type.
-    """
-    column_type = resolve_string_enum(column_type, DataDesignerColumnType)
-    config_cls = get_column_config_cls_from_type(column_type)
-    if column_type == DataDesignerColumnType.SAMPLER:
-        kwargs = _resolve_sampler_kwargs(name, kwargs)
-    return config_cls(name=name, **kwargs)
-
-
-def get_column_config_cls_from_type(column_type: DataDesignerColumnType) -> type[ColumnConfigT]:
-    """Get the column config class for a column type."""
-    column_type = resolve_string_enum(column_type, DataDesignerColumnType)
-    if column_type in _COLUMN_TYPE_CONFIG_CLS_MAP:
-        return _COLUMN_TYPE_CONFIG_CLS_MAP[column_type]
-    if plugin := plugin_manager.get_column_generator_plugin_if_exists(column_type.value):
-        return plugin.config_cls
-    raise InvalidConfigError(f"🛑 {column_type} is not a valid column type.")
-
-
-def get_column_display_order() -> list[DataDesignerColumnType]:
-    """Return the preferred display order of the column types."""
-    display_order = [
-        DataDesignerColumnType.SEED_DATASET,
-        DataDesignerColumnType.SAMPLER,
-        DataDesignerColumnType.LLM_TEXT,
-        DataDesignerColumnType.LLM_CODE,
-        DataDesignerColumnType.LLM_STRUCTURED,
-        DataDesignerColumnType.LLM_JUDGE,
-        DataDesignerColumnType.EMBEDDING,
-        DataDesignerColumnType.VALIDATION,
-        DataDesignerColumnType.EXPRESSION,
-    ]
-    display_order.extend(plugin_manager.get_plugin_column_types(DataDesignerColumnType))
-    return display_order
-
-
-def get_column_emoji_from_type(column_type: DataDesignerColumnType) -> str:
-    """Get the emoji for a column type."""
-    config_cls = get_column_config_cls_from_type(resolve_string_enum(column_type, DataDesignerColumnType))
-    return config_cls.get_column_emoji()
-
-
-def _resolve_sampler_kwargs(name: str, kwargs: dict) -> dict:
-    if "sampler_type" not in kwargs:
-        raise InvalidConfigError(f"🛑 `sampler_type` is required for sampler column '{name}'.")
-    sampler_type = resolve_string_enum(kwargs["sampler_type"], SamplerType)
-
-    # Handle params - it could be a dict or already a concrete object
-    params_value = kwargs.get("params", {})
-    expected_params_class = SAMPLER_PARAMS[sampler_type.value]
-
-    if isinstance(params_value, expected_params_class):
-        # params is already a concrete object of the right type
-        params = params_value
-    elif isinstance(params_value, dict):
-        # params is a dictionary, create new instance
-        params = expected_params_class(**params_value)
-    else:
-        # params is neither dict nor expected type
-        raise InvalidConfigError(
-            f"🛑 Invalid params for sampler column '{name}'. "
-            f"Expected a dictionary or an instance of {expected_params_class.__name__}. "
-            f"You provided {params_value=}."
-        )
-
-    return {
-        "sampler_type": sampler_type,
-        "params": params,
-        **{k: v for k, v in kwargs.items() if k not in ["sampler_type", "params"]},
-    }
-
-
-_COLUMN_TYPE_CONFIG_CLS_MAP = {
-    DataDesignerColumnType.LLM_TEXT: LLMTextColumnConfig,
-    DataDesignerColumnType.LLM_CODE: LLMCodeColumnConfig,
-    DataDesignerColumnType.LLM_STRUCTURED: LLMStructuredColumnConfig,
-    DataDesignerColumnType.LLM_JUDGE: LLMJudgeColumnConfig,
-    DataDesignerColumnType.VALIDATION: ValidationColumnConfig,
-    DataDesignerColumnType.EXPRESSION: ExpressionColumnConfig,
-    DataDesignerColumnType.SAMPLER: SamplerColumnConfig,
-    DataDesignerColumnType.SEED_DATASET: SeedDatasetColumnConfig,
-    DataDesignerColumnType.EMBEDDING: EmbeddingColumnConfig,
-}