data-designer 0.1.5__py3-none-any.whl → 0.2.0__py3-none-any.whl

data_designer/engine/dataset_builders/utils/concurrency.py

@@ -8,7 +8,7 @@ import json
 import logging
 from concurrent.futures import Future, ThreadPoolExecutor
 from threading import Lock, Semaphore
-from typing import Any, Optional, Protocol
+from typing import Any, Protocol
 
 from pydantic import BaseModel, Field
 
@@ -46,13 +46,13 @@ class ExecutorResults(BaseModel):
 class CallbackWithContext(Protocol):
     """Executor callback functions must accept a context kw argument."""
 
-    def __call__(self, result: Any, *, context: Optional[dict] = None) -> Any: ...
+    def __call__(self, result: Any, *, context: dict | None = None) -> Any: ...
 
 
 class ErrorCallbackWithContext(Protocol):
     """Error callbacks take the Exception instance and context."""
 
-    def __call__(self, exc: Exception, *, context: Optional[dict] = None) -> Any: ...
+    def __call__(self, exc: Exception, *, context: dict | None = None) -> Any: ...
 
 
 class ConcurrentThreadExecutor:
@@ -92,8 +92,8 @@ class ConcurrentThreadExecutor:
         *,
         max_workers: int,
         column_name: str,
-        result_callback: Optional[CallbackWithContext] = None,
-        error_callback: Optional[ErrorCallbackWithContext] = None,
+        result_callback: CallbackWithContext | None = None,
+        error_callback: ErrorCallbackWithContext | None = None,
         shutdown_error_rate: float = 0.50,
         shutdown_error_window: int = 10,
     ):
@@ -136,7 +136,7 @@ class ConcurrentThreadExecutor:
             )
         )
 
-    def submit(self, fn, *args, context: Optional[dict] = None, **kwargs) -> None:
+    def submit(self, fn, *args, context: dict | None = None, **kwargs) -> None:
         if self._executor is None:
             raise RuntimeError("Executor is not initialized, this class should be used as a context manager.")
 

data_designer/engine/models/facade.py

@@ -9,9 +9,9 @@ from copy import deepcopy
 from typing import Any
 
 from litellm.types.router import DeploymentTypedDict, LiteLLM_Params
-from litellm.types.utils import ModelResponse
+from litellm.types.utils import EmbeddingResponse, ModelResponse
 
-from data_designer.config.models import ModelConfig, ModelProvider
+from data_designer.config.models import GenerationType, ModelConfig, ModelProvider
 from data_designer.engine.model_provider import ModelProviderRegistry
 from data_designer.engine.models.errors import (
     GenerationValidationFailureError,
@@ -49,6 +49,10 @@ class ModelFacade:
     def model_provider(self) -> ModelProvider:
         return self._model_provider_registry.get_provider(self._model_config.provider)
 
+    @property
+    def model_generation_type(self) -> GenerationType:
+        return self._model_config.generation_type
+
     @property
     def model_provider_name(self) -> str:
         return self.model_provider.name
@@ -64,13 +68,12 @@ class ModelFacade:
     def completion(self, messages: list[dict[str, str]], skip_usage_tracking: bool = False, **kwargs) -> ModelResponse:
         logger.debug(
             f"Prompting model {self.model_name!r}...",
-            extra={"model": self.model_name, "messages": messages, "sensitive": True},
+            extra={"model": self.model_name, "messages": messages},
         )
         response = None
-        if self.model_provider.extra_body:
-            kwargs["extra_body"] = {**kwargs.get("extra_body", {}), **self.model_provider.extra_body}
+        kwargs = self.consolidate_kwargs(**kwargs)
         try:
-            response = self._router.completion(self.model_name, messages, **kwargs)
+            response = self._router.completion(model=self.model_name, messages=messages, **kwargs)
             logger.debug(
                 f"Received completion from model {self.model_name!r}",
                 extra={
@@ -84,9 +87,50 @@ class ModelFacade:
         except Exception as e:
             raise e
         finally:
-            if not skip_usage_tracking:
+            if not skip_usage_tracking and response is not None:
                 self._track_usage(response)
 
+    def consolidate_kwargs(self, **kwargs) -> dict[str, Any]:
+        # Remove purpose from kwargs to avoid passing it to the model
+        kwargs.pop("purpose", None)
+        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}
+        return kwargs
+
+    @catch_llm_exceptions
+    def generate_text_embeddings(
+        self, input_texts: list[str], skip_usage_tracking: bool = False, **kwargs
+    ) -> list[list[float]]:
+        logger.debug(
+            f"Generating embeddings with model {self.model_name!r}...",
+            extra={
+                "model": self.model_name,
+                "input_count": len(input_texts),
+            },
+        )
+        kwargs = self.consolidate_kwargs(**kwargs)
+        response = None
+        try:
+            response = self._router.embedding(model=self.model_name, input=input_texts, **kwargs)
+            logger.debug(
+                f"Received embeddings from model {self.model_name!r}",
+                extra={
+                    "model": self.model_name,
+                    "embedding_count": len(response.data) if response.data else 0,
+                    "usage": self._usage_stats.model_dump(),
+                },
+            )
+            if response.data and len(response.data) == len(input_texts):
+                return [data["embedding"] for data in response.data]
+            else:
+                raise ValueError(f"Expected {len(input_texts)} embeddings, but received {len(response.data)}")
+        except Exception as e:
+            raise e
+        finally:
+            if not skip_usage_tracking and response is not None:
+                self._track_usage_from_embedding(response)
+
     @catch_llm_exceptions
     def generate(
         self,
@@ -218,8 +262,21 @@ class ModelFacade:
         ):
             self._usage_stats.extend(
                 token_usage=TokenUsageStats(
-                    prompt_tokens=response.usage.prompt_tokens,
-                    completion_tokens=response.usage.completion_tokens,
+                    input_tokens=response.usage.prompt_tokens,
+                    output_tokens=response.usage.completion_tokens,
+                ),
+                request_usage=RequestUsageStats(successful_requests=1, failed_requests=0),
+            )
+
+    def _track_usage_from_embedding(self, response: EmbeddingResponse | None) -> None:
+        if response is None:
+            self._usage_stats.extend(request_usage=RequestUsageStats(successful_requests=0, failed_requests=1))
+            return
+        if response.usage is not None and response.usage.prompt_tokens is not None:
+            self._usage_stats.extend(
+                token_usage=TokenUsageStats(
+                    input_tokens=response.usage.prompt_tokens,
+                    output_tokens=0,
                 ),
                 request_usage=RequestUsageStats(successful_requests=1, failed_requests=0),
             )

data_designer/engine/models/litellm_overrides.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 
 import random
 import threading
-from typing import Optional, Union
 
 import httpx
 import litellm
@@ -90,7 +89,7 @@ class CustomRouter(Router):
         self._initial_retry_after_s = initial_retry_after_s
         self._jitter_pct = jitter_pct
 
-    def _extract_retry_delay_from_headers(self, e: Exception) -> Optional[Union[int, float]]:
+    def _extract_retry_delay_from_headers(self, e: Exception) -> int | float | None:
         """
         Most of this code logic was extracted directly from the parent
         `Router`'s `_time_to_sleep_before_retry` function. Our override
@@ -99,7 +98,7 @@ class CustomRouter(Router):
         return this info, we'll simply use that retry value returned here.
         """
 
-        response_headers: Optional[httpx.Headers] = None
+        response_headers: httpx.Headers | None = None
         if hasattr(e, "response") and hasattr(e.response, "headers"):  # type: ignore
             response_headers = e.response.headers  # type: ignore
         if hasattr(e, "litellm_response_headers"):
@@ -119,9 +118,9 @@ class CustomRouter(Router):
         e: Exception,
         remaining_retries: int,
         num_retries: int,
-        healthy_deployments: Optional[list] = None,
-        all_deployments: Optional[list] = None,
-    ) -> Union[int, float]:
+        healthy_deployments: list | None = None,
+        all_deployments: list | None = None,
+    ) -> int | float:
         """
         Implements exponential backoff for retries.
 

data_designer/engine/models/parsers/errors.py

@@ -1,8 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Optional
-
 
 class ParserException(Exception):
     """Identifies errors resulting from generic parser errors.
@@ -12,7 +10,7 @@ class ParserException(Exception):
             attempted to parse.
     """
 
-    source: Optional[str]
+    source: str | None
 
     @staticmethod
     def _log_format(source: str) -> str:
@@ -24,7 +22,7 @@ class ParserException(Exception):
         # return f"<source>{source}</source>"
         return ""
 
-    def __init__(self, msg: Optional[str] = None, source: Optional[str] = None):
+    def __init__(self, msg: str | None = None, source: str | None = None):
         msg = "" if msg is None else msg.strip()
 
         if source is not None:

data_designer/engine/models/parsers/parser.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from functools import reduce
-from typing import Optional
 
 import marko
 from lxml import etree
@@ -105,8 +104,8 @@ class LLMResponseParser:
 
     def __init__(
         self,
-        tag_parsers: Optional[dict[str, TagParser]] = None,
-        postprocessors: Optional[list[PostProcessor]] = None,
+        tag_parsers: dict[str, TagParser] | None = None,
+        postprocessors: list[PostProcessor] | None = None,
     ):
         """
         Initializes the LLMResponseParser with optional tag parsers and post-processors.

data_designer/engine/models/parsers/postprocessors.py

@@ -1,7 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Optional, Type
 
 import json_repair
 from pydantic import BaseModel, ValidationError
@@ -60,12 +59,12 @@ def deserialize_json_code(
 
 
 class RealizePydanticTypes:
-    types: list[Type[BaseModel]]
+    types: list[type[BaseModel]]
 
-    def __init__(self, types: list[Type[BaseModel]]):
+    def __init__(self, types: list[type[BaseModel]]):
         self.types = types
 
-    def _fit_types(self, obj: dict) -> Optional[BaseModel]:
+    def _fit_types(self, obj: dict) -> BaseModel | None:
         final_obj = None
 
         for t in self.types:

data_designer/engine/models/parsers/types.py

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any, Optional, Protocol, Type, runtime_checkable
+from typing import Any, Protocol, runtime_checkable
 
 from lxml.etree import _Element
 from pydantic import BaseModel, Field
@@ -30,7 +30,7 @@ class LLMStructuredResponse(BaseModel):
         out.parsed = out.parsed[-n:]
         return out
 
-    def filter(self, block_types: list[Type[BaseModel]]) -> Self:
+    def filter(self, block_types: list[type[BaseModel]]) -> Self:
         out = self.model_copy()
         out.parsed = [b for b in out.parsed if isinstance(b, tuple(block_types))]
         return out
@@ -44,7 +44,7 @@ class TagParser(Protocol):
     element, do some computation, and return some kind of structured
     output, represented as a subclass of Pydantic `BaseModel`.
     This protocol implementation can cover both classes as well
-    as curried fuctions as parsers (e.g. `partial`).
+    as curried functions as parsers (e.g. `partial`).
     """
 
     def __call__(self, element: _Element) -> BaseModel: ...
@@ -69,7 +69,7 @@ class TextBlock(BaseModel):
 
 class CodeBlock(BaseModel):
     code: str
-    code_lang: Optional[str] = None
+    code_lang: str | None = None
 
 
 class StructuredDataBlock(BaseModel):

data_designer/engine/models/registry.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 
 import logging
 
-from data_designer.config.models import ModelConfig
+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
@@ -73,7 +73,7 @@ class ModelRegistry:
         model_config = self.get_model_config(model_alias=model_alias)
         return self._model_provider_registry.get_provider(model_config.provider)
 
-    def run_health_check(self, model_aliases: set[str]) -> None:
+    def run_health_check(self, model_aliases: list[str]) -> None:
         logger.info("🩺 Running health checks for models...")
         for model_alias in model_aliases:
             model = self.get_model(model_alias=model_alias)
@@ -81,15 +81,24 @@ class ModelRegistry:
                 f"  |-- 👀 Checking {model.model_name!r} in provider named {model.model_provider_name!r} for model alias {model.model_alias!r}..."
             )
             try:
-                model.generate(
-                    prompt="Hello!",
-                    parser=lambda x: x,
-                    system_prompt="You are a helpful assistant.",
-                    max_correction_steps=0,
-                    max_conversation_restarts=0,
-                    skip_usage_tracking=True,
-                    purpose="running health checks",
-                )
+                if model.model_generation_type == GenerationType.EMBEDDING:
+                    model.generate_text_embeddings(
+                        input_texts=["Hello!"],
+                        skip_usage_tracking=True,
+                        purpose="running health checks",
+                    )
+                elif model.model_generation_type == GenerationType.CHAT_COMPLETION:
+                    model.generate(
+                        prompt="Hello!",
+                        parser=lambda x: x,
+                        system_prompt="You are a helpful assistant.",
+                        max_correction_steps=0,
+                        max_conversation_restarts=0,
+                        skip_usage_tracking=True,
+                        purpose="running health checks",
+                    )
+                else:
+                    raise ValueError(f"Unsupported generation type: {model.model_generation_type}")
                 logger.info("  |-- ✅ Passed!")
             except Exception as e:
                 logger.error("  |-- ❌ Failed!")

data_designer/engine/models/usage.py

@@ -11,20 +11,20 @@ logger = logging.getLogger(__name__)
 
 
 class TokenUsageStats(BaseModel):
-    prompt_tokens: int = 0
-    completion_tokens: int = 0
+    input_tokens: int = 0
+    output_tokens: int = 0
 
     @computed_field
     def total_tokens(self) -> int:
-        return self.prompt_tokens + self.completion_tokens
+        return self.input_tokens + self.output_tokens
 
     @property
     def has_usage(self) -> bool:
         return self.total_tokens > 0
 
-    def extend(self, *, prompt_tokens: int, completion_tokens: int) -> None:
-        self.prompt_tokens += prompt_tokens
-        self.completion_tokens += completion_tokens
+    def extend(self, *, input_tokens: int, output_tokens: int) -> None:
+        self.input_tokens += input_tokens
+        self.output_tokens += output_tokens
 
 
 class RequestUsageStats(BaseModel):
@@ -56,9 +56,7 @@ class ModelUsageStats(BaseModel):
         self, *, token_usage: TokenUsageStats | None = None, request_usage: RequestUsageStats | None = None
     ) -> None:
         if token_usage is not None:
-            self.token_usage.extend(
-                prompt_tokens=token_usage.prompt_tokens, completion_tokens=token_usage.completion_tokens
-            )
+            self.token_usage.extend(input_tokens=token_usage.input_tokens, output_tokens=token_usage.output_tokens)
         if request_usage is not None:
             self.request_usage.extend(
                 successful_requests=request_usage.successful_requests, failed_requests=request_usage.failed_requests

data_designer/engine/processing/ginja/ast.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from collections import deque
-from typing import Optional, Type
 
 from jinja2 import nodes as j_nodes
 
@@ -33,7 +32,7 @@ def ast_max_depth(node: j_nodes.Node) -> int:
     return max_depth
 
 
-def ast_descendant_count(ast: j_nodes.Node, only_type: Optional[Type[j_nodes.Node]] = None) -> int:
+def ast_descendant_count(ast: j_nodes.Node, only_type: type[j_nodes.Node] | None = None) -> int:
     """Count the number of nodes which descend from the given node.
 
     Args:

data_designer/engine/processing/utils.py

@@ -1,9 +1,11 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import ast
 import json
 import logging
-from typing import Any, TypeVar, Union, overload
+import re
+from typing import Any, TypeVar, overload
 
 import pandas as pd
 
@@ -25,7 +27,7 @@ def concat_datasets(datasets: list[pd.DataFrame]) -> pd.DataFrame:
 # Overloads to help static type checker better understand
 # the input/output types of the deserialize_json_values function.
 @overload
-def deserialize_json_values(data: str) -> Union[dict[str, Any], list[Any], Any]: ...
+def deserialize_json_values(data: str) -> dict[str, Any] | list[Any] | Any: ...
 
 
 @overload
@@ -100,6 +102,42 @@ def deserialize_json_values(data):
         return data
 
 
+def parse_list_string(text: str) -> list[str]:
+    """Parse a list from a string, handling JSON arrays, Python lists, and trailing commas."""
+    text = text.strip()
+
+    # Try JSON first
+    try:
+        list_obj = json.loads(text)
+        if isinstance(list_obj, list):
+            return _clean_whitespace(list_obj)
+    except json.JSONDecodeError:
+        pass
+
+    # Remove trailing commas before closing brackets (common in JSON-like strings)
+    text_cleaned = re.sub(r",\s*]", "]", text)
+    text_cleaned = re.sub(r",\s*}", "}", text_cleaned)
+
+    # Try JSON again with cleaned text
+    try:
+        return _clean_whitespace(json.loads(text_cleaned))
+    except json.JSONDecodeError:
+        pass
+
+    # Try Python literal eval (handles single quotes)
+    try:
+        return _clean_whitespace(ast.literal_eval(text_cleaned))
+    except (ValueError, SyntaxError):
+        pass
+
+    # If all else fails, return the original text
+    return [text.strip()]
+
+
+def _clean_whitespace(texts: list[str]) -> list[str]:
+    return [text.strip() for text in texts]
+
+
 def _verify_columns_are_unique(datasets: list[pd.DataFrame]) -> None:
     joined_columns = set()
     for df in datasets:

data_designer/engine/registry/base.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import threading
-from typing import Any, Generic, Type, TypeVar
+from typing import Any, Generic, TypeVar
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.utils.type_helpers import StrEnum
@@ -16,14 +16,14 @@ TaskConfigT = TypeVar("TaskConfigT", bound=ConfigBase)
 
 class TaskRegistry(Generic[EnumNameT, TaskT, TaskConfigT]):
     # registered type name -> type
-    _registry: dict[EnumNameT, Type[TaskT]] = {}
+    _registry: dict[EnumNameT, type[TaskT]] = {}
     # type -> registered type name
-    _reverse_registry: dict[Type[TaskT], EnumNameT] = {}
+    _reverse_registry: dict[type[TaskT], EnumNameT] = {}
 
     # registered type name -> config type
-    _config_registry: dict[EnumNameT, Type[TaskConfigT]] = {}
+    _config_registry: dict[EnumNameT, type[TaskConfigT]] = {}
     # config type -> registered type name
-    _reverse_config_registry: dict[Type[TaskConfigT], EnumNameT] = {}
+    _reverse_config_registry: dict[type[TaskConfigT], EnumNameT] = {}
 
     # all registries are singletons
     _instance = None
@@ -33,8 +33,8 @@ class TaskRegistry(Generic[EnumNameT, TaskT, TaskConfigT]):
     def register(
         cls,
         name: EnumNameT,
-        task: Type[TaskT],
-        config: Type[TaskConfigT],
+        task: type[TaskT],
+        config: type[TaskConfigT],
         raise_on_collision: bool = False,
     ) -> None:
         if cls._has_been_registered(name):
@@ -52,22 +52,22 @@ class TaskRegistry(Generic[EnumNameT, TaskT, TaskConfigT]):
             cls._reverse_config_registry[config] = name
 
     @classmethod
-    def get_task_type(cls, name: EnumNameT) -> Type[TaskT]:
+    def get_task_type(cls, name: EnumNameT) -> type[TaskT]:
         cls._raise_if_not_registered(name, cls._registry)
         return cls._registry[name]
 
     @classmethod
-    def get_config_type(cls, name: EnumNameT) -> Type[TaskConfigT]:
+    def get_config_type(cls, name: EnumNameT) -> type[TaskConfigT]:
         cls._raise_if_not_registered(name, cls._config_registry)
         return cls._config_registry[name]
 
     @classmethod
-    def get_registered_name(cls, task: Type[TaskT]) -> EnumNameT:
+    def get_registered_name(cls, task: type[TaskT]) -> EnumNameT:
         cls._raise_if_not_registered(task, cls._reverse_registry)
         return cls._reverse_registry[task]
 
     @classmethod
-    def get_for_config_type(cls, config: Type[TaskConfigT]) -> Type[TaskT]:
+    def get_for_config_type(cls, config: type[TaskConfigT]) -> type[TaskT]:
         cls._raise_if_not_registered(config, cls._reverse_config_registry)
         name = cls._reverse_config_registry[config]
         return cls.get_task_type(name)
@@ -77,7 +77,7 @@ class TaskRegistry(Generic[EnumNameT, TaskT, TaskConfigT]):
         return name in cls._registry
 
     @classmethod
-    def _raise_if_not_registered(cls, key: EnumNameT | Type[TaskT] | Type[TaskConfigT], mapping: dict) -> None:
+    def _raise_if_not_registered(cls, key: EnumNameT | type[TaskT] | type[TaskConfigT], mapping: dict) -> None:
         if not (isinstance(key, StrEnum) or isinstance(key, str)):
             cls._raise_if_not_type(key)
         if key not in mapping:

data_designer/engine/sampling_gen/constraints.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC, abstractmethod
-from typing import Type
 
 import numpy as np
 import pandas as pd
@@ -91,5 +90,5 @@ CONSTRAINT_TYPE_TO_CHECKER = {
 }
 
 
-def get_constraint_checker(constraint_type: ConstraintType) -> Type[ConstraintChecker]:
+def get_constraint_checker(constraint_type: ConstraintType) -> type[ConstraintChecker]:
     return CONSTRAINT_TYPE_TO_CHECKER[ConstraintType(constraint_type)]