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

data_designer/engine/dataset_builders/artifact_storage.py

@@ -1,23 +1,30 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import json
 import logging
 import shutil
 from datetime import datetime
 from functools import cached_property
 from pathlib import Path
+from typing import TYPE_CHECKING
 
-import pandas as pd
 from pydantic import BaseModel, field_validator, model_validator
 
 from data_designer.config.utils.io_helpers import read_parquet_dataset
 from data_designer.config.utils.type_helpers import StrEnum, resolve_string_enum
 from data_designer.engine.dataset_builders.errors import ArtifactStorageError
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd
 
 logger = logging.getLogger(__name__)
 
 BATCH_FILE_NAME_FORMAT = "batch_{batch_number:05d}.parquet"
+SDG_CONFIG_FILENAME = "sdg.json"
 
 
 class BatchStage(StrEnum):
@@ -164,12 +171,6 @@ class ArtifactStorage(BaseModel):
         shutil.move(partial_result_path, final_file_path)
         return final_file_path
 
-    def write_configs(self, json_file_name: str, configs: list[dict]) -> Path:
-        self.mkdir_if_needed(self.base_dataset_path)
-        with open(self.base_dataset_path / json_file_name, "w") as file:
-            json.dump([c.model_dump(mode="json") for c in configs], file, indent=4)
-        return self.base_dataset_path / json_file_name
-
     def write_batch_to_parquet_file(
         self,
         batch_number: int,
@@ -194,11 +195,89 @@ class ArtifactStorage(BaseModel):
         dataframe.to_parquet(file_path, index=False)
         return file_path
 
+    def get_parquet_file_paths(self) -> list[str]:
+        """Get list of parquet file paths relative to base_dataset_path.
+
+        Returns:
+            List of relative paths to parquet files in the final dataset folder.
+        """
+        return [str(f.relative_to(self.base_dataset_path)) for f in sorted(self.final_dataset_path.glob("*.parquet"))]
+
+    def get_processor_file_paths(self) -> dict[str, list[str]]:
+        """Get processor output files organized by processor name.
+
+        Returns:
+            Dictionary mapping processor names to lists of relative file paths.
+        """
+        processor_files: dict[str, list[str]] = {}
+        if self.processors_outputs_path.exists():
+            for processor_dir in sorted(self.processors_outputs_path.iterdir()):
+                if processor_dir.is_dir():
+                    processor_name = processor_dir.name
+                    processor_files[processor_name] = [
+                        str(f.relative_to(self.base_dataset_path))
+                        for f in sorted(processor_dir.rglob("*"))
+                        if f.is_file()
+                    ]
+        return processor_files
+
+    def get_file_paths(self) -> dict[str, list[str] | dict[str, list[str]]]:
+        """Get all file paths organized by type.
+
+        Returns:
+            Dictionary with 'parquet-files' and 'processor-files' keys.
+        """
+        file_paths = {
+            "parquet-files": self.get_parquet_file_paths(),
+        }
+        processor_file_paths = self.get_processor_file_paths()
+        if processor_file_paths:
+            file_paths["processor-files"] = processor_file_paths
+
+        return file_paths
+
+    def read_metadata(self) -> dict:
+        """Read metadata from the metadata.json file.
+
+        Returns:
+            Dictionary containing the metadata.
+
+        Raises:
+            FileNotFoundError: If metadata file doesn't exist.
+        """
+        with open(self.metadata_file_path, "r") as file:
+            return json.load(file)
+
     def write_metadata(self, metadata: dict) -> Path:
+        """Write metadata to the metadata.json file.
+
+        Args:
+            metadata: Dictionary containing metadata to write.
+
+        Returns:
+            Path to the written metadata file.
+        """
         self.mkdir_if_needed(self.base_dataset_path)
         with open(self.metadata_file_path, "w") as file:
-            json.dump(metadata, file)
+            json.dump(metadata, file, indent=4, sort_keys=True)
         return self.metadata_file_path
 
+    def update_metadata(self, updates: dict) -> Path:
+        """Update existing metadata with new fields.
+
+        Args:
+            updates: Dictionary of fields to add/update in metadata.
+
+        Returns:
+            Path to the updated metadata file.
+        """
+        try:
+            existing_metadata = self.read_metadata()
+        except FileNotFoundError:
+            existing_metadata = {}
+
+        existing_metadata.update(updates)
+        return self.write_metadata(existing_metadata)
+
     def _get_stage_path(self, stage: BatchStage) -> Path:
         return getattr(self, resolve_string_enum(stage, BatchStage).value)

data_designer/engine/dataset_builders/column_wise_builder.py

@@ -12,9 +12,9 @@ import uuid
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable
 
-import pandas as pd
-
 from data_designer.config.column_types import ColumnConfigT
+from data_designer.config.config_builder import BuilderConfig
+from data_designer.config.data_designer_config import DataDesignerConfig
 from data_designer.config.dataset_builders import BuildStage
 from data_designer.config.processors import (
     DropColumnsProcessorConfig,
@@ -27,40 +27,38 @@ from data_designer.engine.column_generators.generators.base import (
     GenerationStrategy,
 )
 from data_designer.engine.column_generators.utils.generator_classification import column_type_is_model_generated
-from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
+from data_designer.engine.compiler import compile_data_designer_config
+from data_designer.engine.dataset_builders.artifact_storage import SDG_CONFIG_FILENAME, ArtifactStorage
 from data_designer.engine.dataset_builders.errors import DatasetGenerationError, DatasetProcessingError
-from data_designer.engine.dataset_builders.multi_column_configs import (
-    DatasetBuilderColumnConfigT,
-    MultiColumnConfig,
-)
+from data_designer.engine.dataset_builders.multi_column_configs import MultiColumnConfig
 from data_designer.engine.dataset_builders.utils.concurrency import (
     MAX_CONCURRENCY_PER_NON_LLM_GENERATOR,
     ConcurrentThreadExecutor,
 )
-from data_designer.engine.dataset_builders.utils.dataset_batch_manager import (
-    DatasetBatchManager,
-)
+from data_designer.engine.dataset_builders.utils.config_compiler import compile_dataset_builder_column_configs
+from data_designer.engine.dataset_builders.utils.dataset_batch_manager import DatasetBatchManager
 from data_designer.engine.models.telemetry import InferenceEvent, NemoSourceEnum, TaskStatusEnum, TelemetryHandler
 from data_designer.engine.processing.processors.base import Processor
 from data_designer.engine.processing.processors.drop_columns import DropColumnsProcessor
 from data_designer.engine.registry.data_designer_registry import DataDesignerRegistry
 from data_designer.engine.resources.resource_provider import ResourceProvider
+from data_designer.lazy_heavy_imports import pd
 
 if TYPE_CHECKING:
+    import pandas as pd
+
     from data_designer.engine.column_generators.generators.base import ColumnGeneratorWithModelRegistry
     from data_designer.engine.models.usage import ModelUsageStats
 
 logger = logging.getLogger(__name__)
 
-
 _CLIENT_VERSION: str = importlib.metadata.version("data_designer")
 
 
 class ColumnWiseDatasetBuilder:
     def __init__(
         self,
-        column_configs: list[DatasetBuilderColumnConfigT],
-        processor_configs: list[ProcessorConfig],
+        data_designer_config: DataDesignerConfig,
         resource_provider: ResourceProvider,
         registry: DataDesignerRegistry | None = None,
     ):
@@ -68,8 +66,12 @@ class ColumnWiseDatasetBuilder:
         self._resource_provider = resource_provider
         self._records_to_drop: set[int] = set()
         self._registry = registry or DataDesignerRegistry()
-        self._column_configs = column_configs
-        self._processors: dict[BuildStage, list[Processor]] = self._initialize_processors(processor_configs)
+
+        self._data_designer_config = compile_data_designer_config(data_designer_config, resource_provider)
+        self._column_configs = compile_dataset_builder_column_configs(self._data_designer_config)
+        self._processors: dict[BuildStage, list[Processor]] = self._initialize_processors(
+            self._data_designer_config.processors or []
+        )
         self._validate_column_configs()
 
     @property
@@ -96,9 +98,8 @@ class ColumnWiseDatasetBuilder:
         num_records: int,
         on_batch_complete: Callable[[Path], None] | None = None,
     ) -> Path:
-        self._write_configs()
         self._run_model_health_check_if_needed()
-
+        self._write_builder_config()
         generators = self._initialize_generators()
         start_time = time.perf_counter()
         group_id = uuid.uuid4().hex
@@ -157,6 +158,12 @@ class ColumnWiseDatasetBuilder:
             for config in self._column_configs
         ]
 
+    def _write_builder_config(self) -> None:
+        self.artifact_storage.mkdir_if_needed(self.artifact_storage.base_dataset_path)
+        BuilderConfig(data_designer=self._data_designer_config).to_json(
+            self.artifact_storage.base_dataset_path / SDG_CONFIG_FILENAME
+        )
+
     def _run_batch(
         self, generators: list[ColumnGenerator], *, batch_mode: str, save_partial_results: bool = True, group_id: str
     ) -> None:
@@ -164,15 +171,16 @@ class ColumnWiseDatasetBuilder:
         for generator in generators:
             generator.log_pre_generation()
             try:
+                generation_strategy = generator.get_generation_strategy()
                 if generator.can_generate_from_scratch and self.batch_manager.buffer_is_empty:
                     self._run_from_scratch_column_generator(generator)
-                elif generator.generation_strategy == GenerationStrategy.CELL_BY_CELL:
+                elif generation_strategy == GenerationStrategy.CELL_BY_CELL:
                     self._run_cell_by_cell_generator(generator)
-                elif generator.generation_strategy == GenerationStrategy.FULL_COLUMN:
+                elif generation_strategy == GenerationStrategy.FULL_COLUMN:
                     self._run_full_column_generator(generator)
                 else:
-                    logger.error(f"❌ Unknown generation strategy: {generator.generation_strategy}")
-                    raise DatasetGenerationError(f"🛑 Unknown generation strategy: {generator.generation_strategy}")
+                    logger.error(f"❌ Unknown generation strategy: {generation_strategy}")
+                    raise DatasetGenerationError(f"🛑 Unknown generation strategy: {generation_strategy}")
                 if save_partial_results:
                     self.batch_manager.write()
             except Exception as e:
@@ -210,9 +218,9 @@ class ColumnWiseDatasetBuilder:
             )
 
     def _fan_out_with_threads(self, generator: ColumnGeneratorWithModelRegistry, max_workers: int) -> None:
-        if generator.generation_strategy != GenerationStrategy.CELL_BY_CELL:
+        if generator.get_generation_strategy() != GenerationStrategy.CELL_BY_CELL:
             raise DatasetGenerationError(
-                f"Generator {generator.metadata().name} is not a {GenerationStrategy.CELL_BY_CELL} "
+                f"Generator {generator.name} is not a {GenerationStrategy.CELL_BY_CELL} "
                 "generator so concurrency through threads is not supported."
             )
 
@@ -292,7 +300,7 @@ class ColumnWiseDatasetBuilder:
                 dataframe = processor.process(dataframe, current_batch_number=current_batch_number)
             except Exception as e:
                 raise DatasetProcessingError(
-                    f"🛑 Failed to process dataset with processor {processor.metadata().name} in stage {stage}: {e}"
+                    f"🛑 Failed to process dataset with processor {processor.name} in stage {stage}: {e}"
                 ) from e
         return dataframe
 
@@ -307,16 +315,6 @@ class ColumnWiseDatasetBuilder:
     def _worker_result_callback(self, result: dict, *, context: dict | None = None) -> None:
         self.batch_manager.update_record(context["index"], result)
 
-    def _write_configs(self) -> None:
-        self.artifact_storage.write_configs(
-            json_file_name="column_configs.json",
-            configs=self._column_configs,
-        )
-        self.artifact_storage.write_configs(
-            json_file_name="model_configs.json",
-            configs=self._resource_provider.model_registry.model_configs.values(),
-        )
-
     def _emit_batch_inference_events(
         self, batch_mode: str, usage_deltas: dict[str, ModelUsageStats], group_id: str
     ) -> None:

data_designer/engine/dataset_builders/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/dataset_builders/multi_column_configs.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 abc import ABC
 from typing import TypeAlias
 

data_designer/engine/dataset_builders/utils/config_compiler.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.column_types import DataDesignerColumnType
 from data_designer.config.data_designer_config import DataDesignerConfig
 from data_designer.config.processors import ProcessorConfig

data_designer/engine/dataset_builders/utils/dag.py

@@ -1,13 +1,18 @@
 # 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 networkx as nx
+import logging
+from typing import TYPE_CHECKING
 
 from data_designer.config.column_types import ColumnConfigT
 from data_designer.engine.column_generators.utils.generator_classification import column_type_used_in_execution_dag
 from data_designer.engine.dataset_builders.utils.errors import DAGCircularDependencyError
+from data_designer.lazy_heavy_imports import nx
+
+if TYPE_CHECKING:
+    import networkx as nx
 
 logger = logging.getLogger(__name__)
 

data_designer/engine/dataset_builders/utils/dataset_batch_manager.py

@@ -1,16 +1,20 @@
 # 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 shutil
 from pathlib import Path
-from typing import Callable, Container, Iterator
-
-import pandas as pd
-import pyarrow.parquet as pq
+from typing import TYPE_CHECKING, Callable, Container, Iterator
 
 from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage, BatchStage
 from data_designer.engine.dataset_builders.utils.errors import DatasetBatchManagementError
+from data_designer.lazy_heavy_imports import pd, pq
+
+if TYPE_CHECKING:
+    import pandas as pd
+    import pyarrow.parquet as pq
 
 logger = logging.getLogger(__name__)
 
@@ -87,8 +91,7 @@ class DatasetBatchManager:
                     "total_num_batches": self.num_batches,
                     "buffer_size": self._buffer_size,
                     "schema": {field.name: str(field.type) for field in pq.read_schema(final_file_path)},
-                    "file_paths": [str(f) for f in sorted(self.artifact_storage.final_dataset_path.glob("*.parquet"))],
-                    "num_records": self.num_records_list[: self._current_batch_number + 1],
+                    "file_paths": self.artifact_storage.get_file_paths(),
                     "num_completed_batches": self._current_batch_number + 1,
                     "dataset_name": self.artifact_storage.dataset_name,
                 }

data_designer/engine/dataset_builders/utils/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/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 pydantic import BaseModel, Field
 
 from data_designer.errors import DataDesignerError

data_designer/engine/model_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 functools import cached_property
 
 from pydantic import BaseModel, field_validator, model_validator

data_designer/engine/models/errors.py

@@ -6,25 +6,15 @@ from __future__ import annotations
 import logging
 from collections.abc import Callable
 from functools import wraps
-from typing import Any
-
-from litellm.exceptions import (
-    APIConnectionError,
-    APIError,
-    AuthenticationError,
-    BadRequestError,
-    ContextWindowExceededError,
-    InternalServerError,
-    NotFoundError,
-    PermissionDeniedError,
-    RateLimitError,
-    Timeout,
-    UnprocessableEntityError,
-    UnsupportedParamsError,
-)
+from typing import TYPE_CHECKING, Any
+
 from pydantic import BaseModel
 
 from data_designer.engine.errors import DataDesignerError
+from data_designer.lazy_heavy_imports import litellm
+
+if TYPE_CHECKING:
+    import litellm
 
 logger = logging.getLogger(__name__)
 
@@ -132,10 +122,10 @@ def handle_llm_exceptions(
     err_msg_parser = DownstreamLLMExceptionMessageParser(model_name, model_provider_name, purpose)
     match exception:
         # Common errors that can come from LiteLLM
-        case APIError():
+        case litellm.exceptions.APIError():
             raise err_msg_parser.parse_api_error(exception, authentication_error) from None
 
-        case APIConnectionError():
+        case litellm.exceptions.APIConnectionError():
             raise ModelAPIConnectionError(
                 FormattedLLMErrorMessage(
                     cause=f"Connection to model {model_name!r} hosted on model provider {model_provider_name!r} failed while {purpose}.",
@@ -143,13 +133,13 @@ def handle_llm_exceptions(
                 )
             ) from None
 
-        case AuthenticationError():
+        case litellm.exceptions.AuthenticationError():
             raise ModelAuthenticationError(authentication_error) from None
 
-        case ContextWindowExceededError():
+        case litellm.exceptions.ContextWindowExceededError():
             raise err_msg_parser.parse_context_window_exceeded_error(exception) from None
 
-        case UnsupportedParamsError():
+        case litellm.exceptions.UnsupportedParamsError():
             raise ModelUnsupportedParamsError(
                 FormattedLLMErrorMessage(
                     cause=f"One or more of the parameters you provided were found to be unsupported by model {model_name!r} while {purpose}.",
@@ -157,10 +147,10 @@ def handle_llm_exceptions(
                 )
             ) from None
 
-        case BadRequestError():
+        case litellm.exceptions.BadRequestError():
             raise err_msg_parser.parse_bad_request_error(exception) from None
 
-        case InternalServerError():
+        case litellm.exceptions.InternalServerError():
             raise ModelInternalServerError(
                 FormattedLLMErrorMessage(
                     cause=f"Model {model_name!r} is currently experiencing internal server issues while {purpose}.",
@@ -168,7 +158,7 @@ def handle_llm_exceptions(
                 )
             ) from None
 
-        case NotFoundError():
+        case litellm.exceptions.NotFoundError():
             raise ModelNotFoundError(
                 FormattedLLMErrorMessage(
                     cause=f"The specified model {model_name!r} could not be found while {purpose}.",
@@ -176,7 +166,7 @@ def handle_llm_exceptions(
                 )
             ) from None
 
-        case PermissionDeniedError():
+        case litellm.exceptions.PermissionDeniedError():
             raise ModelPermissionDeniedError(
                 FormattedLLMErrorMessage(
                     cause=f"Your API key was found to lack the necessary permissions to use model {model_name!r} while {purpose}.",
@@ -184,7 +174,7 @@ def handle_llm_exceptions(
                 )
             ) from None
 
-        case RateLimitError():
+        case litellm.exceptions.RateLimitError():
             raise ModelRateLimitError(
                 FormattedLLMErrorMessage(
                     cause=f"You have exceeded the rate limit for model {model_name!r} while {purpose}.",
@@ -192,7 +182,7 @@ def handle_llm_exceptions(
                 )
             ) from None
 
-        case Timeout():
+        case litellm.exceptions.Timeout():
             raise ModelTimeoutError(
                 FormattedLLMErrorMessage(
                     cause=f"The request to model {model_name!r} timed out while {purpose}.",
@@ -200,7 +190,7 @@ def handle_llm_exceptions(
                 )
             ) from None
 
-        case UnprocessableEntityError():
+        case litellm.exceptions.UnprocessableEntityError():
             raise ModelUnprocessableEntityError(
                 FormattedLLMErrorMessage(
                     cause=f"The request to model {model_name!r} failed despite correct request format while {purpose}.",
@@ -264,7 +254,7 @@ class DownstreamLLMExceptionMessageParser:
         self.model_provider_name = model_provider_name
         self.purpose = purpose
 
-    def parse_bad_request_error(self, exception: BadRequestError) -> DataDesignerError:
+    def parse_bad_request_error(self, exception: litellm.exceptions.BadRequestError) -> DataDesignerError:
         err_msg = FormattedLLMErrorMessage(
             cause=f"The request for model {self.model_name!r} was found to be malformed or missing required parameters while {self.purpose}.",
             solution="Check your request parameters and try again.",
@@ -276,7 +266,9 @@ class DownstreamLLMExceptionMessageParser:
             )
         return ModelBadRequestError(err_msg)
 
-    def parse_context_window_exceeded_error(self, exception: ContextWindowExceededError) -> DataDesignerError:
+    def parse_context_window_exceeded_error(
+        self, exception: litellm.exceptions.ContextWindowExceededError
+    ) -> DataDesignerError:
         cause = f"The input data for model '{self.model_name}' was found to exceed its supported context width while {self.purpose}."
         try:
             if "OpenAIException - This model's maximum context length is " in str(exception):
@@ -295,7 +287,7 @@ class DownstreamLLMExceptionMessageParser:
             )
 
     def parse_api_error(
-        self, exception: InternalServerError, auth_error_msg: FormattedLLMErrorMessage
+        self, exception: litellm.exceptions.InternalServerError, auth_error_msg: FormattedLLMErrorMessage
     ) -> DataDesignerError:
         if "Error code: 403" in str(exception):
             return ModelAuthenticationError(auth_error_msg)

data_designer/engine/models/facade.py

@@ -6,10 +6,7 @@ from __future__ import annotations
 import logging
 from collections.abc import Callable
 from copy import deepcopy
-from typing import Any
-
-from litellm.types.router import DeploymentTypedDict, LiteLLM_Params
-from litellm.types.utils import EmbeddingResponse, ModelResponse
+from typing import TYPE_CHECKING, Any
 
 from data_designer.config.models import GenerationType, ModelConfig, ModelProvider
 from data_designer.engine.model_provider import ModelProviderRegistry
@@ -23,6 +20,10 @@ from data_designer.engine.models.parsers.errors import ParserException
 from data_designer.engine.models.usage import ModelUsageStats, RequestUsageStats, TokenUsageStats
 from data_designer.engine.models.utils import prompt_to_messages, str_to_message
 from data_designer.engine.secret_resolver import SecretResolver
+from data_designer.lazy_heavy_imports import litellm
+
+if TYPE_CHECKING:
+    import litellm
 
 logger = logging.getLogger(__name__)
 
@@ -65,7 +66,9 @@ class ModelFacade:
     def usage_stats(self) -> ModelUsageStats:
         return self._usage_stats
 
-    def completion(self, messages: list[dict[str, str]], skip_usage_tracking: bool = False, **kwargs) -> ModelResponse:
+    def completion(
+        self, messages: list[dict[str, str]], skip_usage_tracking: bool = False, **kwargs
+    ) -> litellm.ModelResponse:
         logger.debug(
             f"Prompting model {self.model_name!r}...",
             extra={"model": self.model_name, "messages": messages},
@@ -236,14 +239,14 @@ class ModelFacade:
                     ) from exc
         return output_obj, reasoning_trace
 
-    def _get_litellm_deployment(self, model_config: ModelConfig) -> DeploymentTypedDict:
+    def _get_litellm_deployment(self, model_config: ModelConfig) -> litellm.DeploymentTypedDict:
         provider = self._model_provider_registry.get_provider(model_config.provider)
         api_key = None
         if provider.api_key:
             api_key = self._secret_resolver.resolve(provider.api_key)
         api_key = api_key or "not-used-but-required"
 
-        litellm_params = LiteLLM_Params(
+        litellm_params = litellm.LiteLLM_Params(
             model=f"{provider.provider_type}/{model_config.model}",
             api_base=provider.endpoint,
             api_key=api_key,
@@ -253,7 +256,7 @@ class ModelFacade:
             "litellm_params": litellm_params.model_dump(),
         }
 
-    def _track_usage(self, response: ModelResponse | None) -> None:
+    def _track_usage(self, response: litellm.types.utils.ModelResponse | None) -> None:
         if response is None:
             self._usage_stats.extend(request_usage=RequestUsageStats(successful_requests=0, failed_requests=1))
             return
@@ -270,7 +273,7 @@ class ModelFacade:
                 request_usage=RequestUsageStats(successful_requests=1, failed_requests=0),
             )
 
-    def _track_usage_from_embedding(self, response: EmbeddingResponse | None) -> None:
+    def _track_usage_from_embedding(self, response: litellm.types.utils.EmbeddingResponse | None) -> None:
         if response is None:
             self._usage_stats.extend(request_usage=RequestUsageStats(successful_requests=0, failed_requests=1))
             return

data_designer/engine/models/factory.py

@@ -0,0 +1,42 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from data_designer.config.models import ModelConfig
+from data_designer.engine.model_provider import ModelProviderRegistry
+from data_designer.engine.secret_resolver import SecretResolver
+
+if TYPE_CHECKING:
+    from data_designer.engine.models.registry import ModelRegistry
+
+
+def create_model_registry(
+    *,
+    model_configs: list[ModelConfig] | None = None,
+    secret_resolver: SecretResolver,
+    model_provider_registry: ModelProviderRegistry,
+) -> ModelRegistry:
+    """Factory function for creating a ModelRegistry instance.
+
+    Heavy dependencies (litellm, httpx) are deferred until this function is called.
+    This is a factory function pattern - imports inside factories are idiomatic Python
+    for lazy initialization.
+    """
+    from data_designer.engine.models.facade import ModelFacade
+    from data_designer.engine.models.litellm_overrides import apply_litellm_patches
+    from data_designer.engine.models.registry import ModelRegistry
+
+    apply_litellm_patches()
+
+    def model_facade_factory(model_config, secret_resolver, model_provider_registry):
+        return ModelFacade(model_config, secret_resolver, model_provider_registry)
+
+    return ModelRegistry(
+        model_configs=model_configs,
+        secret_resolver=secret_resolver,
+        model_provider_registry=model_provider_registry,
+        model_facade_factory=model_facade_factory,
+    )