data-designer 0.2.3__py3-none-any.whl → 0.3.1__py3-none-any.whl

data_designer/config/utils/constants.py

@@ -282,6 +282,10 @@ OPENAI_PROVIDER_NAME = "openai"
 
 OPENAI_API_KEY_ENV_VAR_NAME = "OPENAI_API_KEY"
 
+OPENROUTER_PROVIDER_NAME = "openrouter"
+
+OPENROUTER_API_KEY_ENV_VAR_NAME = "OPENROUTER_API_KEY"
+
 PREDEFINED_PROVIDERS = [
     {
         "name": NVIDIA_PROVIDER_NAME,
@@ -295,6 +299,12 @@ PREDEFINED_PROVIDERS = [
         "provider_type": "openai",
         "api_key": OPENAI_API_KEY_ENV_VAR_NAME,
     },
+    {
+        "name": OPENROUTER_PROVIDER_NAME,
+        "endpoint": "https://openrouter.ai/api/v1",
+        "provider_type": "openai",
+        "api_key": OPENROUTER_API_KEY_ENV_VAR_NAME,
+    },
 ]
 
 
@@ -302,11 +312,14 @@ DEFAULT_TEXT_INFERENCE_PARAMS = {"temperature": 0.85, "top_p": 0.95}
 DEFAULT_REASONING_INFERENCE_PARAMS = {"temperature": 0.35, "top_p": 0.95}
 DEFAULT_VISION_INFERENCE_PARAMS = {"temperature": 0.85, "top_p": 0.95}
 DEFAULT_EMBEDDING_INFERENCE_PARAMS = {"encoding_format": "float"}
-
+NEMOTRON_3_NANO_30B_A3B_INFERENCE_PARAMS = {"temperature": 1.0, "top_p": 1.0}
 
 PREDEFINED_PROVIDERS_MODEL_MAP = {
     NVIDIA_PROVIDER_NAME: {
-        "text": {"model": "nvidia/nemotron-3-nano-30b-a3b", "inference_parameters": {"temperature": 1.0, "top_p": 1.0}},
+        "text": {
+            "model": "nvidia/nemotron-3-nano-30b-a3b",
+            "inference_parameters": NEMOTRON_3_NANO_30B_A3B_INFERENCE_PARAMS,
+        },
         "reasoning": {"model": "openai/gpt-oss-20b", "inference_parameters": DEFAULT_REASONING_INFERENCE_PARAMS},
         "vision": {"model": "nvidia/nemotron-nano-12b-v2-vl", "inference_parameters": DEFAULT_VISION_INFERENCE_PARAMS},
         "embedding": {
@@ -320,6 +333,18 @@ PREDEFINED_PROVIDERS_MODEL_MAP = {
         "vision": {"model": "gpt-5", "inference_parameters": DEFAULT_VISION_INFERENCE_PARAMS},
         "embedding": {"model": "text-embedding-3-large", "inference_parameters": DEFAULT_EMBEDDING_INFERENCE_PARAMS},
     },
+    OPENROUTER_PROVIDER_NAME: {
+        "text": {
+            "model": "nvidia/nemotron-3-nano-30b-a3b",
+            "inference_parameters": NEMOTRON_3_NANO_30B_A3B_INFERENCE_PARAMS,
+        },
+        "reasoning": {"model": "openai/gpt-oss-20b", "inference_parameters": DEFAULT_REASONING_INFERENCE_PARAMS},
+        "vision": {"model": "nvidia/nemotron-nano-12b-v2-vl", "inference_parameters": DEFAULT_VISION_INFERENCE_PARAMS},
+        "embedding": {
+            "model": "openai/text-embedding-3-large",
+            "inference_parameters": DEFAULT_EMBEDDING_INFERENCE_PARAMS,
+        },
+    },
 }
 
 # Persona locale metadata - used by the CLI and the person sampler.

data_designer/config/utils/io_helpers.py

@@ -108,26 +108,6 @@ def read_parquet_dataset(path: Path) -> pd.DataFrame:
             raise e
 
 
-def write_seed_dataset(dataframe: pd.DataFrame, file_path: Path) -> None:
-    """Write a seed dataset to a file in the specified format.
-
-    Supported file extensions: .parquet, .csv, .json, .jsonl
-
-    Args:
-        dataframe: The pandas DataFrame to write.
-        file_path: The path where the dataset should be saved.
-            Format is inferred from the file extension.
-    """
-    file_path = validate_dataset_file_path(file_path, should_exist=False)
-    logger.info(f"💾 Saving seed dataset to {file_path}")
-    if file_path.suffix.lower() == ".parquet":
-        dataframe.to_parquet(file_path, index=False)
-    elif file_path.suffix.lower() == ".csv":
-        dataframe.to_csv(file_path, index=False)
-    elif file_path.suffix.lower() in {".json", ".jsonl"}:
-        dataframe.to_json(file_path, orient="records", lines=True)
-
-
 def validate_dataset_file_path(file_path: str | Path, should_exist: bool = True) -> Path:
     """Validate that a dataset file path has a valid extension and optionally exists.
 

data_designer/engine/column_generators/generators/seed_dataset.py

@@ -30,7 +30,7 @@ class SeedDatasetColumnGenerator(FromScratchColumnGenerator[SeedDatasetMultiColu
             name="seed_dataset_column_generator",
             description="Sample columns from a seed dataset.",
             generation_strategy=GenerationStrategy.FULL_COLUMN,
-            required_resources=[ResourceType.DATASTORE],
+            required_resources=[ResourceType.SEED_READER],
         )
 
     @property
@@ -39,10 +39,10 @@ class SeedDatasetColumnGenerator(FromScratchColumnGenerator[SeedDatasetMultiColu
 
     @functools.cached_property
     def duckdb_conn(self) -> duckdb.DuckDBPyConnection:
-        return self.resource_provider.datastore.create_duckdb_connection()
+        return self.resource_provider.seed_reader.create_duckdb_connection()
 
-    def generate(self, dataset: pd.DataFrame) -> pd.DataFrame:
-        return concat_datasets([self.generate_from_scratch(len(dataset)), dataset])
+    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:
@@ -57,7 +57,7 @@ class SeedDatasetColumnGenerator(FromScratchColumnGenerator[SeedDatasetMultiColu
         self._num_records_sampled = 0
         self._batch_reader = None
         self._df_remaining = None
-        self._dataset_uri = self.resource_provider.datastore.get_dataset_uri(self.config.dataset)
+        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()
 

data_designer/engine/column_generators/generators/validation.py

@@ -123,11 +123,14 @@ class ValidationColumnGenerator(ColumnGenerator[ValidationColumnConfig]):
         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,
         ) as executor:
             for i, batch in enumerate(batched_records):
                 executor.submit(lambda batch: self._validate_batch(validator, batch), batch, context={"index": i})

data_designer/engine/column_generators/registry.py

@@ -51,7 +51,7 @@ def create_default_column_generator_registry(with_plugins: bool = True) -> Colum
         for plugin in PluginRegistry().get_plugins(PluginType.COLUMN_GENERATOR):
             registry.register(
                 DataDesignerColumnType(plugin.name),
-                plugin.task_cls,
+                plugin.impl_cls,
                 plugin.config_cls,
             )
 

data_designer/engine/compiler.py

@@ -0,0 +1,69 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+from data_designer.config.column_configs import SeedDatasetColumnConfig
+from data_designer.config.config_builder import DataDesignerConfigBuilder
+from data_designer.config.data_designer_config import DataDesignerConfig
+from data_designer.config.errors import InvalidConfigError
+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_builder: DataDesignerConfigBuilder, resource_provider: ResourceProvider
+) -> DataDesignerConfig:
+    config = config_builder.build()
+    _resolve_and_add_seed_columns(config, resource_provider.seed_reader)
+    _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 _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/dataset_builders/column_wise_builder.py

@@ -217,11 +217,14 @@ class ColumnWiseDatasetBuilder:
             f"🐙 Processing {generator.config.column_type} column '{generator.config.name}' "
             f"with {max_workers} concurrent workers"
         )
+        settings = self._resource_provider.run_config
         with ConcurrentThreadExecutor(
             max_workers=max_workers,
             column_name=generator.config.name,
             result_callback=self._worker_result_callback,
             error_callback=self._worker_error_callback,
+            shutdown_error_rate=settings.shutdown_error_rate,
+            shutdown_error_window=settings.shutdown_error_window,
         ) as executor:
             for i, record in self.batch_manager.iter_current_batch():
                 executor.submit(lambda record: generator.generate(record), record, context={"index": i})

data_designer/engine/dataset_builders/utils/config_compiler.py

@@ -34,7 +34,7 @@ def compile_dataset_builder_column_configs(config: DataDesignerConfig) -> list[D
         compiled_column_configs.append(
             SeedDatasetMultiColumnConfig(
                 columns=seed_column_configs,
-                dataset=config.seed_config.dataset,
+                source=config.seed_config.source,
                 sampling_strategy=config.seed_config.sampling_strategy,
                 selection_strategy=config.seed_config.selection_strategy,
             )

data_designer/engine/models/facade.py

@@ -96,6 +96,8 @@ class ModelFacade:
         kwargs = {**self._model_config.inference_parameters.generate_kwargs, **kwargs}
         if self.model_provider.extra_body:
             kwargs["extra_body"] = {**kwargs.get("extra_body", {}), **self.model_provider.extra_body}
+        if self.model_provider.extra_headers:
+            kwargs["extra_headers"] = self.model_provider.extra_headers
         return kwargs
 
     @catch_llm_exceptions

data_designer/engine/processing/gsonschema/validators.py

@@ -2,7 +2,9 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
+import re
 from copy import deepcopy
+from decimal import ROUND_HALF_UP, Decimal
 from typing import Any, overload
 
 from jsonschema import Draft202012Validator, ValidationError, validators
@@ -70,6 +72,57 @@ def extend_jsonschema_validator_with_pruning(validator):
     return validators.extend(validator, {"additionalProperties": prune_additional_properties})
 
 
+def _get_decimal_info_from_anyof(schema: dict) -> tuple[bool, int | None]:
+    """Check if schema is a Decimal anyOf and extract decimal places.
+
+    Returns (is_decimal, decimal_places) where decimal_places is None if no constraint.
+    """
+    any_of = schema.get("anyOf")
+    if not isinstance(any_of, list):
+        return False, None
+
+    has_number = any(item.get("type") == "number" for item in any_of)
+    if not has_number:
+        return False, None
+
+    for item in any_of:
+        if item.get("type") == "string" and "pattern" in item:
+            match = re.search(r"\\d\{0,(\d+)\}", item["pattern"])
+            if match:
+                return True, int(match.group(1))
+            return True, None  # Decimal without precision constraint
+    return False, None
+
+
+def normalize_decimal_fields(obj: DataObjectT, schema: JSONSchemaT) -> DataObjectT:
+    """Normalize Decimal-like anyOf fields to floats with proper precision."""
+    if not isinstance(obj, dict):
+        return obj
+
+    defs = schema.get("$defs", {})
+    obj_schema = defs.get(schema.get("$ref", "")[len("#/$defs/") :], schema)
+    props = obj_schema.get("properties", {})
+
+    for key, value in obj.items():
+        field_schema = props.get(key, {})
+        if "$ref" in field_schema:
+            field_schema = defs.get(field_schema["$ref"][len("#/$defs/") :], {})
+
+        if isinstance(value, dict):
+            obj[key] = normalize_decimal_fields(value, schema)
+        elif isinstance(value, list):
+            obj[key] = [normalize_decimal_fields(v, schema) if isinstance(v, dict) else v for v in value]
+        elif isinstance(value, (int, float, str)) and not isinstance(value, bool):
+            is_decimal, decimal_places = _get_decimal_info_from_anyof(field_schema)
+            if is_decimal:
+                d = Decimal(str(value))
+                if decimal_places is not None:
+                    d = d.quantize(Decimal(f"0.{'0' * decimal_places}"), rounding=ROUND_HALF_UP)
+                obj[key] = float(d)
+
+    return obj
+
+
 ## We don't expect the outer data type (e.g. dict, list, or const) to be
 ## modified by the pruning action.
 @overload
@@ -140,4 +193,6 @@ def validate(
     except ValidationError as exc:
         raise JSONSchemaValidationError(str(exc)) from exc
 
+    final_object = normalize_decimal_fields(final_object, schema)
+
     return final_object

data_designer/engine/resources/resource_provider.py

@@ -3,26 +3,29 @@
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.models import ModelConfig
+from data_designer.config.run_config import RunConfig
+from data_designer.config.seed_source import SeedSource
 from data_designer.config.utils.type_helpers import StrEnum
 from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
 from data_designer.engine.model_provider import ModelProviderRegistry
 from data_designer.engine.models.registry import ModelRegistry, create_model_registry
 from data_designer.engine.resources.managed_storage import ManagedBlobStorage, init_managed_blob_storage
-from data_designer.engine.resources.seed_dataset_data_store import SeedDatasetDataStore
+from data_designer.engine.resources.seed_reader import SeedReader, SeedReaderRegistry
 from data_designer.engine.secret_resolver import SecretResolver
 
 
 class ResourceType(StrEnum):
     BLOB_STORAGE = "blob_storage"
-    DATASTORE = "datastore"
     MODEL_REGISTRY = "model_registry"
+    SEED_READER = "seed_reader"
 
 
 class ResourceProvider(ConfigBase):
     artifact_storage: ArtifactStorage
     blob_storage: ManagedBlobStorage | None = None
-    datastore: SeedDatasetDataStore | None = None
     model_registry: ModelRegistry | None = None
+    run_config: RunConfig = RunConfig()
+    seed_reader: SeedReader | None = None
 
 
 def create_resource_provider(
@@ -31,16 +34,25 @@ def create_resource_provider(
     model_configs: list[ModelConfig],
     secret_resolver: SecretResolver,
     model_provider_registry: ModelProviderRegistry,
-    datastore: SeedDatasetDataStore | None = None,
+    seed_reader_registry: SeedReaderRegistry,
     blob_storage: ManagedBlobStorage | None = None,
+    seed_dataset_source: SeedSource | None = None,
+    run_config: RunConfig | None = None,
 ) -> ResourceProvider:
+    seed_reader = None
+    if seed_dataset_source:
+        seed_reader = seed_reader_registry.get_reader(
+            seed_dataset_source,
+            secret_resolver,
+        )
     return ResourceProvider(
         artifact_storage=artifact_storage,
-        datastore=datastore,
         model_registry=create_model_registry(
             model_configs=model_configs,
             secret_resolver=secret_resolver,
             model_provider_registry=model_provider_registry,
         ),
         blob_storage=blob_storage or init_managed_blob_storage(),
+        seed_reader=seed_reader,
+        run_config=run_config or RunConfig(),
     )

data_designer/engine/resources/seed_reader.py

@@ -0,0 +1,149 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import Generic, TypeVar, get_args, get_origin
+
+import duckdb
+from huggingface_hub import HfFileSystem
+from typing_extensions import Self
+
+from data_designer.config.seed_source import (
+    DataFrameSeedSource,
+    HuggingFaceSeedSource,
+    LocalFileSeedSource,
+    SeedSource,
+)
+from data_designer.engine.secret_resolver import SecretResolver
+from data_designer.errors import DataDesignerError
+
+
+class SeedReaderError(DataDesignerError): ...
+
+
+SourceT = TypeVar("ConfigT", bound=SeedSource)
+
+
+class SeedReader(ABC, Generic[SourceT]):
+    """Base class for reading a seed dataset.
+
+    Seeds are read using duckdb. Reader implementations define duckdb connection setup details
+    and how to get a URI that can be queried with duckdb (i.e. "... FROM <uri> ...").
+
+    The Data Designer engine automatically supplies the appropriate SeedSource
+    and a SecretResolver to use for any secret fields in the config.
+    """
+
+    source: SourceT
+    secret_resolver: SecretResolver
+
+    @abstractmethod
+    def get_dataset_uri(self) -> str: ...
+
+    @abstractmethod
+    def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection: ...
+
+    def attach(self, source: SourceT, secret_resolver: SecretResolver):
+        """Attach a source and secret resolver to the instance.
+
+        This is called internally by the engine so that these objects do not
+        need to be provided in the reader's constructor.
+        """
+        self.source = source
+        self.secret_resolver = secret_resolver
+
+    def get_column_names(self) -> list[str]:
+        """Returns the seed dataset's column names"""
+        conn = self.create_duckdb_connection()
+        describe_query = f"DESCRIBE SELECT * FROM '{self.get_dataset_uri()}'"
+        column_descriptions = conn.execute(describe_query).fetchall()
+        return [col[0] for col in column_descriptions]
+
+    def get_seed_type(self) -> str:
+        """Return the seed_type of the source class this reader is generic over."""
+        # Get the generic type arguments from the reader class
+        # Check __orig_bases__ for the generic base class
+        for base in getattr(type(self), "__orig_bases__", []):
+            origin = get_origin(base)
+            if origin is SeedReader:
+                args = get_args(base)
+                if args:
+                    source_cls = args[0]
+                    # Extract seed_type from the source class
+                    if hasattr(source_cls, "model_fields") and "seed_type" in source_cls.model_fields:
+                        field = source_cls.model_fields["seed_type"]
+                        default_value = field.default
+                        if isinstance(default_value, str):
+                            return default_value
+
+        raise SeedReaderError("Reader does not have a valid generic source type with seed_type")
+
+
+class LocalFileSeedReader(SeedReader[LocalFileSeedSource]):
+    def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection:
+        return duckdb.connect()
+
+    def get_dataset_uri(self) -> str:
+        return self.source.path
+
+
+class HuggingFaceSeedReader(SeedReader[HuggingFaceSeedSource]):
+    def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection:
+        token = self.secret_resolver.resolve(self.source.token) if self.source.token else None
+
+        # Use skip_instance_cache to avoid fsspec-level caching
+        hffs = HfFileSystem(endpoint=self.source.endpoint, token=token, skip_instance_cache=True)
+
+        # Clear all internal caches to avoid stale metadata issues
+        # HfFileSystem caches file metadata (size, etc.) which can become stale when files are re-uploaded
+        if hasattr(hffs, "dircache"):
+            hffs.dircache.clear()
+
+        conn = duckdb.connect()
+        conn.register_filesystem(hffs)
+        return conn
+
+    def get_dataset_uri(self) -> str:
+        return f"hf://{self.source.path}"
+
+
+class DataFrameSeedReader(SeedReader[DataFrameSeedSource]):
+    # This is a "magic string" that gets registered in the duckdb connection to make the dataframe directly queryable.
+    _table_name = "df"
+
+    def create_duckdb_connection(self) -> duckdb.DuckDBPyConnection:
+        conn = duckdb.connect()
+        conn.register(self._table_name, self.source.df)
+        return conn
+
+    def get_dataset_uri(self) -> str:
+        return self._table_name
+
+
+class SeedReaderRegistry:
+    def __init__(self, readers: Sequence[SeedReader]):
+        self._readers: dict[str, SeedReader] = {}
+        for reader in readers:
+            self.add_reader(reader)
+
+    def add_reader(self, reader: SeedReader) -> Self:
+        seed_type = reader.get_seed_type()
+
+        if seed_type in self._readers:
+            raise SeedReaderError(f"A reader for seed_type {seed_type!r} already exists")
+
+        self._readers[seed_type] = reader
+        return self
+
+    def get_reader(self, seed_dataset_source: SeedSource, secret_resolver: SecretResolver) -> SeedReader:
+        reader = self._get_reader_for_source(seed_dataset_source)
+        reader.attach(seed_dataset_source, secret_resolver)
+        return reader
+
+    def _get_reader_for_source(self, seed_dataset_source: SeedSource) -> SeedReader:
+        seed_type = seed_dataset_source.seed_type
+        try:
+            return self._readers[seed_type]
+        except KeyError:
+            raise SeedReaderError(f"No reader found for seed_type {seed_type!r}")

data_designer/essentials/__init__.py

@@ -3,6 +3,7 @@
 
 from data_designer.config.default_model_settings import resolve_seed_default_model_settings
 from data_designer.config.exports import *  # noqa: F403
+from data_designer.config.run_config import RunConfig
 from data_designer.config.validator_params import LocalCallableValidatorParams
 from data_designer.interface.data_designer import DataDesigner
 from data_designer.logging import LoggingConfig, configure_logging
@@ -21,6 +22,7 @@ def get_essentials_exports() -> list[str]:
     local = [
         DataDesigner.__name__,
         LocalCallableValidatorParams.__name__,
+        RunConfig.__name__,
     ]
 
     return logging + local + get_config_exports()  # noqa: F405