data-designer 0.3.3__py3-none-any.whl → 0.3.5__py3-none-any.whl

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,
+    )

data_designer/engine/models/litellm_overrides.py

@@ -5,21 +5,26 @@ 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
 from pydantic import BaseModel, Field
 from typing_extensions import override
 
+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 +40,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.caching.in_memory_cache.InMemoryCache):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -78,7 +83,7 @@ class ThreadSafeCache(InMemoryCache):
             super().flush_cache()
 
 
-class CustomRouter(Router):
+class CustomRouter(litellm.router.Router):
     def __init__(
         self,
         *args,
@@ -155,7 +160,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,