data-designer 0.3.4__py3-none-any.whl → 0.3.6__py3-none-any.whl

data_designer/engine/models/litellm_overrides.py

@@ -5,21 +5,32 @@ from __future__ import annotations
 
 import random
 import threading
+from typing import TYPE_CHECKING
 
-import httpx
-import litellm
-from litellm import RetryPolicy
-from litellm.caching.in_memory_cache import InMemoryCache
-from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager
-from litellm.router import Router
+# Import specific litellm submodules needed for class inheritance
+# Note: Class inheritance requires base classes at definition time, so we import these directly.
+# Runtime litellm usage below still benefits from lazy loading via the litellm alias.
+import litellm.caching.in_memory_cache as _litellm_cache
+import litellm.router as _litellm_router
 from pydantic import BaseModel, Field
 from typing_extensions import override
 
+# Use lazy loading for runtime litellm usage (RetryPolicy, utils, etc.)
+from data_designer.lazy_heavy_imports import httpx, litellm
 from data_designer.logging import quiet_noisy_logger
 
+if TYPE_CHECKING:
+    import httpx
+    import litellm
+
 DEFAULT_MAX_CALLBACKS = 1000
 
 
+def _get_logging_callback_manager():
+    """Lazy accessor for LoggingCallbackManager to avoid loading litellm at import time."""
+    return litellm.litellm_core_utils.logging_callback_manager.LoggingCallbackManager
+
+
 class LiteLLMRouterDefaultKwargs(BaseModel):
     ## Number of seconds to wait initially after a connection
     ## failure.
@@ -35,15 +46,15 @@ class LiteLLMRouterDefaultKwargs(BaseModel):
 
     ## Sets the default retry policy, including the number
     ## of retries to use in particular scenarios.
-    retry_policy: RetryPolicy = Field(
-        default_factory=lambda: RetryPolicy(
+    retry_policy: litellm.RetryPolicy = Field(
+        default_factory=lambda: litellm.RetryPolicy(
             RateLimitErrorRetries=3,
             TimeoutErrorRetries=3,
         )
     )
 
 
-class ThreadSafeCache(InMemoryCache):
+class ThreadSafeCache(_litellm_cache.InMemoryCache):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -78,7 +89,7 @@ class ThreadSafeCache(InMemoryCache):
             super().flush_cache()
 
 
-class CustomRouter(Router):
+class CustomRouter(_litellm_router.Router):
     def __init__(
         self,
         *args,
@@ -155,7 +166,7 @@ def apply_litellm_patches():
     litellm.in_memory_llm_clients_cache = ThreadSafeCache()
 
     # Workaround for the litellm issue described in https://github.com/BerriAI/litellm/issues/9792
-    LoggingCallbackManager.MAX_CALLBACKS = DEFAULT_MAX_CALLBACKS
+    _get_logging_callback_manager().MAX_CALLBACKS = DEFAULT_MAX_CALLBACKS
 
     quiet_noisy_logger("httpx")
     quiet_noisy_logger("LiteLLM")

data_designer/engine/models/parsers/errors.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 
 class ParserException(Exception):
     """Identifies errors resulting from generic parser errors.

data_designer/engine/models/parsers/parser.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from functools import reduce
 
 import marko
@@ -80,13 +82,11 @@ class LLMResponseParser:
             code: str
             syntax: Optional[str] = None
 
-
         class CodeBlockParser:
             def __call__(self, element: _Element) -> CodeBlock:
                 # Implementation details...
                 return CodeBlock(code=element.text, syntax=element.get("class"))
 
-
         parser = LLMResponseParser(
             tag_parsers={
                 "pre.code": CodeBlockParser(),

data_designer/engine/models/parsers/postprocessors.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
 
 import json_repair
 from pydantic import BaseModel, ValidationError

data_designer/engine/models/parsers/tag_parsers.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from lxml.etree import _Element
 
 from data_designer.engine.models.parsers.types import CodeBlock, TextBlock

data_designer/engine/models/parsers/types.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from typing import Any, Protocol, runtime_checkable
 
 from lxml.etree import _Element

data_designer/engine/models/recipes/base.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import abc
 from collections.abc import Callable
 from typing import Generic, TypeVar

data_designer/engine/models/recipes/response_recipes.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import json
 from collections.abc import Callable
 

data_designer/engine/models/registry.py

@@ -4,14 +4,17 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Callable
+from typing import TYPE_CHECKING
 
 from data_designer.config.models import GenerationType, ModelConfig
 from data_designer.engine.model_provider import ModelProvider, ModelProviderRegistry
-from data_designer.engine.models.facade import ModelFacade
-from data_designer.engine.models.litellm_overrides import apply_litellm_patches
 from data_designer.engine.models.usage import ModelUsageStats, RequestUsageStats, TokenUsageStats
 from data_designer.engine.secret_resolver import SecretResolver
 
+if TYPE_CHECKING:
+    from data_designer.engine.models.facade import ModelFacade
+
 logger = logging.getLogger(__name__)
 
 
@@ -22,10 +25,12 @@ class ModelRegistry:
         secret_resolver: SecretResolver,
         model_provider_registry: ModelProviderRegistry,
         model_configs: list[ModelConfig] | None = None,
+        model_facade_factory: Callable[[ModelConfig, SecretResolver, ModelProviderRegistry], ModelFacade] | None = None,
     ):
         self._secret_resolver = secret_resolver
         self._model_provider_registry = model_provider_registry
-        self._model_configs = {}
+        self._model_facade_factory = model_facade_factory
+        self._model_configs: dict[str, ModelConfig] = {}
         self._models: dict[str, ModelFacade] = {}
         self._set_model_configs(model_configs)
 
@@ -136,18 +141,6 @@ class ModelRegistry:
         # Models are now lazily initialized in get_model() when first requested
 
     def _get_model(self, model_config: ModelConfig) -> ModelFacade:
-        return ModelFacade(model_config, self._secret_resolver, self._model_provider_registry)
-
-
-def create_model_registry(
-    *,
-    model_configs: list[ModelConfig] | None = None,
-    secret_resolver: SecretResolver,
-    model_provider_registry: ModelProviderRegistry,
-) -> ModelRegistry:
-    apply_litellm_patches()
-    return ModelRegistry(
-        model_configs=model_configs,
-        secret_resolver=secret_resolver,
-        model_provider_registry=model_provider_registry,
-    )
+        if self._model_facade_factory is None:
+            raise RuntimeError("ModelRegistry was not initialized with a model_facade_factory")
+        return self._model_facade_factory(model_config, self._secret_resolver, self._model_provider_registry)

data_designer/engine/models/telemetry.py

@@ -18,11 +18,15 @@ import platform
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from enum import Enum
-from typing import Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
-import httpx
 from pydantic import BaseModel, Field
 
+from data_designer.lazy_heavy_imports import httpx
+
+if TYPE_CHECKING:
+    import httpx
+
 TELEMETRY_ENABLED = os.getenv("NEMO_TELEMETRY_ENABLED", "true").lower() in ("1", "true", "yes")
 CLIENT_ID = "184482118588404"
 NEMO_TELEMETRY_VERSION = "nemo-telemetry/1.0"

data_designer/engine/processing/ginja/ast.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from collections import deque
 
 from jinja2 import nodes as j_nodes

data_designer/engine/processing/ginja/environment.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import re
 from collections.abc import Callable
 from functools import partial, wraps

data_designer/engine/processing/ginja/exceptions.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import re
 
 from jinja2 import TemplateAssertionError

data_designer/engine/processing/ginja/record.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import json
 
 from data_designer.config.utils.io_helpers import serialize_data

data_designer/engine/processing/gsonschema/exceptions.py

@@ -1,8 +1,15 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from jsonschema import ValidationError
+from __future__ import annotations
 
+from typing import TYPE_CHECKING
 
-class JSONSchemaValidationError(ValidationError):
+from data_designer.lazy_heavy_imports import jsonschema
+
+if TYPE_CHECKING:
+    import jsonschema
+
+
+class JSONSchemaValidationError(jsonschema.ValidationError):
     """Alias of ValidationError to ease imports."""

data_designer/engine/processing/gsonschema/schema_transformers.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from copy import deepcopy
 from typing import Any
 

data_designer/engine/processing/gsonschema/types.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from typing import Any, TypeVar
 
 T_primitive = TypeVar("T_primitive", str, int, float, bool)

data_designer/engine/processing/gsonschema/validators.py

@@ -1,19 +1,23 @@
 # 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 re
 from copy import deepcopy
 from decimal import ROUND_HALF_UP, Decimal
-from typing import Any, overload
-
-from jsonschema import Draft202012Validator, ValidationError, validators
+from typing import TYPE_CHECKING, Any, overload
 
 from data_designer.engine.processing.gsonschema.exceptions import JSONSchemaValidationError
 from data_designer.engine.processing.gsonschema.schema_transformers import forbid_additional_properties
 from data_designer.engine.processing.gsonschema.types import DataObjectT, JSONSchemaT, T_primitive
+from data_designer.lazy_heavy_imports import jsonschema
+
+if TYPE_CHECKING:
+    import jsonschema
 
-DEFAULT_JSONSCHEMA_VALIDATOR = Draft202012Validator
+DEFAULT_JSONSCHEMA_VALIDATOR = jsonschema.Draft202012Validator
 
 logger = logging.getLogger(__name__)
 
@@ -69,7 +73,7 @@ def extend_jsonschema_validator_with_pruning(validator):
         Type[jsonschema.Validator]: A validator class that will
             prune extra fields.
     """
-    return validators.extend(validator, {"additionalProperties": prune_additional_properties})
+    return jsonschema.validators.extend(validator, {"additionalProperties": prune_additional_properties})
 
 
 def _get_decimal_info_from_anyof(schema: dict) -> tuple[bool, int | None]:
@@ -190,7 +194,7 @@ def validate(
 
     try:
         validator(schema).validate(final_object)
-    except ValidationError as exc:
+    except jsonschema.ValidationError as exc:
         raise JSONSchemaValidationError(str(exc)) from exc
 
     final_object = normalize_decimal_fields(final_object, schema)

data_designer/engine/processing/processors/base.py

@@ -5,13 +5,9 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 
-from data_designer.engine.configurable_task import ConfigurableTask, ConfigurableTaskMetadata, DataT, TaskConfigT
+from data_designer.engine.configurable_task import ConfigurableTask, DataT, TaskConfigT
 
 
 class Processor(ConfigurableTask[TaskConfigT], ABC):
-    @staticmethod
-    @abstractmethod
-    def metadata() -> ConfigurableTaskMetadata: ...
-
     @abstractmethod
     def process(self, data: DataT, *, current_batch_number: int | None = None) -> DataT: ...

data_designer/engine/processing/processors/drop_columns.py

@@ -1,26 +1,23 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-import logging
+from __future__ import annotations
 
-import pandas as pd
+import logging
+from typing import TYPE_CHECKING
 
 from data_designer.config.processors import DropColumnsProcessorConfig
-from data_designer.engine.configurable_task import ConfigurableTaskMetadata
 from data_designer.engine.dataset_builders.artifact_storage import BatchStage
 from data_designer.engine.processing.processors.base import Processor
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 logger = logging.getLogger(__name__)
 
 
 class DropColumnsProcessor(Processor[DropColumnsProcessorConfig]):
-    @staticmethod
-    def metadata() -> ConfigurableTaskMetadata:
-        return ConfigurableTaskMetadata(
-            name="drop_columns_processor",
-            description="Drop columns from the input dataset.",
-        )
-
     def process(self, data: pd.DataFrame, *, current_batch_number: int | None = None) -> pd.DataFrame:
         logger.info(f"🙈 Dropping columns: {self.config.column_names}")
         if current_batch_number is not None:  # not in preview mode

data_designer/engine/processing/processors/registry.py

@@ -1,6 +1,8 @@
 # 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.processors import (
     DropColumnsProcessorConfig,

data_designer/engine/processing/processors/schema_transform.py

@@ -1,29 +1,26 @@
 # 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 pandas as pd
+from typing import TYPE_CHECKING
 
 from data_designer.config.processors import SchemaTransformProcessorConfig
-from data_designer.engine.configurable_task import ConfigurableTaskMetadata
 from data_designer.engine.dataset_builders.artifact_storage import BatchStage
 from data_designer.engine.processing.ginja.environment import WithJinja2UserTemplateRendering
 from data_designer.engine.processing.processors.base import Processor
 from data_designer.engine.processing.utils import deserialize_json_values
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 logger = logging.getLogger(__name__)
 
 
 class SchemaTransformProcessor(WithJinja2UserTemplateRendering, Processor[SchemaTransformProcessorConfig]):
-    @staticmethod
-    def metadata() -> ConfigurableTaskMetadata:
-        return ConfigurableTaskMetadata(
-            name="schema_transform_processor",
-            description="Generate dataset with transformed schema using a Jinja2 template.",
-        )
-
     @property
     def template_as_str(self) -> str:
         return json.dumps(self.config.template)

data_designer/engine/processing/utils.py

@@ -1,13 +1,18 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import ast
 import json
 import logging
 import re
-from typing import Any, TypeVar, overload
+from typing import TYPE_CHECKING, Any, TypeVar, overload
+
+from data_designer.lazy_heavy_imports import pd
 
-import pandas as pd
+if TYPE_CHECKING:
+    import pandas as pd
 
 logger = logging.getLogger(__name__)
 
@@ -52,7 +57,6 @@ def deserialize_json_values(data):
             - Dictionary (potentially with nested JSON strings to deserialize)
             - Some other object that can't be deserialized.
 
-
     Returns:
         Deserialized data in the corresponding format:
             - Dictionary (when input is a single string)

data_designer/engine/registry/base.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import threading
 from typing import Any, Generic, TypeVar
 

data_designer/engine/registry/data_designer_registry.py

@@ -1,6 +1,8 @@
 # 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.analysis.column_profilers.registry import (
     ColumnProfilerRegistry,
     create_default_column_profiler_registry,

data_designer/engine/registry/errors.py

@@ -1,6 +1,8 @@
 # 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
 
 

data_designer/engine/resources/managed_dataset_generator.py

@@ -1,11 +1,15 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any
+from __future__ import annotations
 
-import pandas as pd
+from typing import TYPE_CHECKING, Any
 
 from data_designer.engine.resources.managed_dataset_repository import ManagedDatasetRepository
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 
 class ManagedDatasetGenerator:

data_designer/engine/resources/managed_dataset_repository.py

@@ -1,6 +1,8 @@
 # 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 tempfile
 import threading
@@ -9,13 +11,15 @@ from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from functools import cached_property
 from pathlib import Path
-from typing import Any
-
-import duckdb
-import pandas as pd
+from typing import TYPE_CHECKING, Any
 
 from data_designer.config.utils.constants import LOCALES_WITH_MANAGED_DATASETS
 from data_designer.engine.resources.managed_storage import LocalBlobStorageProvider, ManagedBlobStorage
+from data_designer.lazy_heavy_imports import duckdb, pd
+
+if TYPE_CHECKING:
+    import duckdb
+    import pandas as pd
 
 logger = logging.getLogger(__name__)
 
@@ -52,7 +56,6 @@ class Table:
 
 DataCatalog = list[Table]
 
-
 # For now we hardcode the remote data catalog in code. This make it easier
 # initialize the data catalog. Eventually we can make this work more
 # dynamically once this data catalog pattern becomes more widely adopted.

data_designer/engine/resources/managed_storage.py

@@ -1,6 +1,8 @@
 # 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 abc import ABC, abstractmethod
 from collections.abc import Iterator

data_designer/engine/resources/resource_provider.py

@@ -1,6 +1,8 @@
 # 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.dataset_metadata import DatasetMetadata
 from data_designer.config.models import ModelConfig
@@ -9,7 +11,8 @@ 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.models.factory import create_model_registry
+from data_designer.engine.models.registry import ModelRegistry
 from data_designer.engine.resources.managed_storage import ManagedBlobStorage, init_managed_blob_storage
 from data_designer.engine.resources.seed_reader import SeedReader, SeedReaderRegistry
 from data_designer.engine.secret_resolver import SecretResolver
@@ -51,12 +54,16 @@ def create_resource_provider(
     seed_dataset_source: SeedSource | None = None,
     run_config: RunConfig | None = None,
 ) -> ResourceProvider:
+    """Factory function for creating a ResourceProvider instance.
+    This function triggers lazy loading of heavy dependencies like litellm.
+    """
     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,
         model_registry=create_model_registry(

data_designer/engine/resources/seed_reader.py

@@ -1,11 +1,12 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
-from typing import Generic, TypeVar, get_args, get_origin
+from typing import TYPE_CHECKING, Generic, TypeVar, get_args, get_origin
 
-import duckdb
 from huggingface_hub import HfFileSystem
 from typing_extensions import Self
 
@@ -17,6 +18,10 @@ from data_designer.config.seed_source import (
 )
 from data_designer.engine.secret_resolver import SecretResolver
 from data_designer.errors import DataDesignerError
+from data_designer.lazy_heavy_imports import duckdb
+
+if TYPE_CHECKING:
+    import duckdb
 
 
 class SeedReaderError(DataDesignerError): ...

data_designer/engine/sampling_gen/column.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from typing import Any
 
 from pydantic import field_serializer, model_validator