data-designer 0.3.6__tar.gz → 0.3.7__tar.gz

{data_designer-0.3.6 → data_designer-0.3.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-designer
-Version: 0.3.6
+Version: 0.3.7
 Summary: General framework for synthetic data generation
 License-Expression: Apache-2.0
 License-File: LICENSE

{data_designer-0.3.6 → data_designer-0.3.7}/src/data_designer/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.6'
-__version_tuple__ = version_tuple = (0, 3, 6)
+__version__ = version = '0.3.7'
+__version_tuple__ = version_tuple = (0, 3, 7)
 
 __commit_id__ = commit_id = None

{data_designer-0.3.6 → data_designer-0.3.7}/src/data_designer/engine/models/litellm_overrides.py

@@ -1,36 +1,42 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+
+"""
+LiteLLM overrides and customizations.
+
+Note on imports: This module uses direct (eager) imports for litellm rather than lazy loading.
+This is intentional because:
+
+1. Class inheritance requires base classes to be resolved at class definition time,
+   making lazy imports incompatible with our ThreadSafeCache and CustomRouter classes.
+
+2. This module is already lazily loaded at the application level - it's only imported
+   by facade.py, which itself is imported inside the create_model_registry() factory
+   function. So litellm is only loaded when models are actually needed.
+
+3. Attempting to use lazy imports here causes intermittent ImportErrors.
+"""
+
 from __future__ import annotations
 
 import random
 import threading
-from typing import TYPE_CHECKING
 
-# 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
+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
 
-# 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.
@@ -46,15 +52,15 @@ class LiteLLMRouterDefaultKwargs(BaseModel):
 
     ## Sets the default retry policy, including the number
     ## of retries to use in particular scenarios.
-    retry_policy: litellm.RetryPolicy = Field(
-        default_factory=lambda: litellm.RetryPolicy(
+    retry_policy: RetryPolicy = Field(
+        default_factory=lambda: RetryPolicy(
             RateLimitErrorRetries=3,
             TimeoutErrorRetries=3,
         )
     )
 
 
-class ThreadSafeCache(_litellm_cache.InMemoryCache):
+class ThreadSafeCache(InMemoryCache):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -89,7 +95,7 @@ class ThreadSafeCache(_litellm_cache.InMemoryCache):
             super().flush_cache()
 
 
-class CustomRouter(_litellm_router.Router):
+class CustomRouter(Router):
     def __init__(
         self,
         *args,
@@ -166,7 +172,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
-    _get_logging_callback_manager().MAX_CALLBACKS = DEFAULT_MAX_CALLBACKS
+    LoggingCallbackManager.MAX_CALLBACKS = DEFAULT_MAX_CALLBACKS
 
     quiet_noisy_logger("httpx")
     quiet_noisy_logger("LiteLLM")

{data_designer-0.3.6 → data_designer-0.3.7}/tests/engine/models/test_litellm_overrides.py

@@ -3,9 +3,9 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
 from unittest.mock import patch
 
+import litellm
 import pytest
 
 from data_designer.engine.models.litellm_overrides import (
@@ -14,10 +14,6 @@ from data_designer.engine.models.litellm_overrides import (
     ThreadSafeCache,
     apply_litellm_patches,
 )
-from data_designer.lazy_heavy_imports import litellm
-
-if TYPE_CHECKING:
-    import litellm
 
 
 @pytest.fixture