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

data_designer/engine/column_generators/generators/{llm_generators.py → llm_completion.py}

@@ -10,43 +10,41 @@ from data_designer.config.column_configs import (
     LLMStructuredColumnConfig,
     LLMTextColumnConfig,
 )
-from data_designer.config.column_types import COLUMN_TYPE_EMOJI_MAP
-from data_designer.config.models import InferenceParameters, ModelConfig
 from data_designer.config.utils.constants import REASONING_TRACE_COLUMN_POSTFIX
 from data_designer.engine.column_generators.generators.base import (
     ColumnGenerator,
     GenerationStrategy,
     GeneratorMetadata,
+    WithModelGeneration,
 )
 from data_designer.engine.column_generators.utils.prompt_renderer import (
     PromptType,
     RecordBasedPromptRenderer,
     create_response_recipe,
 )
-from data_designer.engine.models.facade import ModelFacade
 from data_designer.engine.models.recipes.base import ResponseRecipe
 from data_designer.engine.processing.utils import deserialize_json_values
 from data_designer.engine.resources.resource_provider import ResourceType
 
-DEFAULT_MAX_CONVERSATION_RESTARTS = 5
-DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS = 0
+logger = logging.getLogger(__name__)
 
 
-logger = logging.getLogger(__name__)
+DEFAULT_MAX_CONVERSATION_RESTARTS = 5
+DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS = 0
 
 
-class WithLLMGeneration:
+class WithChatCompletionGeneration(WithModelGeneration):
     @functools.cached_property
-    def model(self) -> ModelFacade:
-        return self.resource_provider.model_registry.get_model(model_alias=self.config.model_alias)
+    def response_recipe(self) -> ResponseRecipe:
+        return create_response_recipe(self.config, self.model_config)
 
-    @functools.cached_property
-    def model_config(self) -> ModelConfig:
-        return self.resource_provider.model_registry.get_model_config(model_alias=self.config.model_alias)
+    @property
+    def max_conversation_correction_steps(self) -> int:
+        return DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS
 
-    @functools.cached_property
-    def inference_parameters(self) -> InferenceParameters:
-        return self.model_config.inference_parameters
+    @property
+    def max_conversation_restarts(self) -> int:
+        return DEFAULT_MAX_CONVERSATION_RESTARTS
 
     @functools.cached_property
     def prompt_renderer(self) -> RecordBasedPromptRenderer:
@@ -59,18 +57,6 @@ class WithLLMGeneration:
             },
         )
 
-    @functools.cached_property
-    def response_recipe(self) -> ResponseRecipe:
-        return create_response_recipe(self.config, self.model_config)
-
-    @property
-    def max_conversation_correction_steps(self) -> int:
-        return DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS
-
-    @property
-    def max_conversation_restarts(self) -> int:
-        return DEFAULT_MAX_CONVERSATION_RESTARTS
-
     def generate(self, data: dict) -> dict:
         deserialized_record = deserialize_json_values(data)
 
@@ -96,7 +82,6 @@ class WithLLMGeneration:
             max_correction_steps=self.max_conversation_correction_steps,
             max_conversation_restarts=self.max_conversation_restarts,
             purpose=f"running generation for column '{self.config.name}'",
-            **self.inference_parameters.generate_kwargs,
         )
 
         data[self.config.name] = deserialize_json_values(self.response_recipe.serialize_output(response))
@@ -106,21 +91,8 @@ class WithLLMGeneration:
 
         return data
 
-    def log_pre_generation(self) -> None:
-        emoji = COLUMN_TYPE_EMOJI_MAP[self.config.column_type]
-        logger.info(f"{emoji} Preparing {self.config.column_type} column generation")
-        logger.info(f"  |-- column name: {self.config.name!r}")
-        logger.info(f"  |-- model config:\n{self.model_config.model_dump_json(indent=4)}")
-        if self.model_config.provider is None:
-            logger.info(f"  |-- default model provider: {self._get_provider_name()!r}")
-
-    def _get_provider_name(self) -> str:
-        model_alias = self.model_config.alias
-        provider = self.resource_provider.model_registry.get_model_provider(model_alias=model_alias)
-        return provider.name
-
 
-class LLMTextCellGenerator(WithLLMGeneration, ColumnGenerator[LLMTextColumnConfig]):
+class LLMTextCellGenerator(WithChatCompletionGeneration, ColumnGenerator[LLMTextColumnConfig]):
     @staticmethod
     def metadata() -> GeneratorMetadata:
         return GeneratorMetadata(
@@ -131,7 +103,7 @@ class LLMTextCellGenerator(WithLLMGeneration, ColumnGenerator[LLMTextColumnConfi
         )
 
 
-class LLMCodeCellGenerator(WithLLMGeneration, ColumnGenerator[LLMCodeColumnConfig]):
+class LLMCodeCellGenerator(WithChatCompletionGeneration, ColumnGenerator[LLMCodeColumnConfig]):
     @staticmethod
     def metadata() -> GeneratorMetadata:
         return GeneratorMetadata(
@@ -142,7 +114,7 @@ class LLMCodeCellGenerator(WithLLMGeneration, ColumnGenerator[LLMCodeColumnConfi
         )
 
 
-class LLMStructuredCellGenerator(WithLLMGeneration, ColumnGenerator[LLMStructuredColumnConfig]):
+class LLMStructuredCellGenerator(WithChatCompletionGeneration, ColumnGenerator[LLMStructuredColumnConfig]):
     @staticmethod
     def metadata() -> GeneratorMetadata:
         return GeneratorMetadata(
@@ -153,7 +125,7 @@ class LLMStructuredCellGenerator(WithLLMGeneration, ColumnGenerator[LLMStructure
         )
 
 
-class LLMJudgeCellGenerator(WithLLMGeneration, ColumnGenerator[LLMJudgeColumnConfig]):
+class LLMJudgeCellGenerator(WithChatCompletionGeneration, ColumnGenerator[LLMJudgeColumnConfig]):
     @staticmethod
     def metadata() -> GeneratorMetadata:
         return GeneratorMetadata(
@@ -163,10 +135,6 @@ class LLMJudgeCellGenerator(WithLLMGeneration, ColumnGenerator[LLMJudgeColumnCon
             required_resources=[ResourceType.MODEL_REGISTRY],
         )
 
-    @property
-    def max_conversation_correction_steps(self) -> int:
-        return DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS
-
     @property
     def max_conversation_restarts(self) -> int:
         return 2 * DEFAULT_MAX_CONVERSATION_RESTARTS

data_designer/engine/column_generators/registry.py

@@ -3,6 +3,7 @@
 
 from data_designer.config.base import ConfigBase
 from data_designer.config.column_configs import (
+    EmbeddingColumnConfig,
     ExpressionColumnConfig,
     LLMCodeColumnConfig,
     LLMJudgeColumnConfig,
@@ -12,8 +13,9 @@ from data_designer.config.column_configs import (
 )
 from data_designer.config.column_types import DataDesignerColumnType
 from data_designer.engine.column_generators.generators.base import ColumnGenerator
+from data_designer.engine.column_generators.generators.embedding import EmbeddingCellGenerator
 from data_designer.engine.column_generators.generators.expression import ExpressionColumnGenerator
-from data_designer.engine.column_generators.generators.llm_generators import (
+from data_designer.engine.column_generators.generators.llm_completion import (
     LLMCodeCellGenerator,
     LLMJudgeCellGenerator,
     LLMStructuredCellGenerator,
@@ -40,11 +42,11 @@ def create_default_column_generator_registry(with_plugins: bool = True) -> Colum
     registry.register(DataDesignerColumnType.LLM_CODE, LLMCodeCellGenerator, LLMCodeColumnConfig)
     registry.register(DataDesignerColumnType.LLM_JUDGE, LLMJudgeCellGenerator, LLMJudgeColumnConfig)
     registry.register(DataDesignerColumnType.EXPRESSION, ExpressionColumnGenerator, ExpressionColumnConfig)
+    registry.register(DataDesignerColumnType.EMBEDDING, EmbeddingCellGenerator, EmbeddingColumnConfig)
     registry.register(DataDesignerColumnType.SAMPLER, SamplerColumnGenerator, SamplerMultiColumnConfig)
     registry.register(DataDesignerColumnType.SEED_DATASET, SeedDatasetColumnGenerator, SeedDatasetMultiColumnConfig)
     registry.register(DataDesignerColumnType.VALIDATION, ValidationColumnGenerator, ValidationColumnConfig)
     registry.register(DataDesignerColumnType.LLM_STRUCTURED, LLMStructuredCellGenerator, LLMStructuredColumnConfig)
-
     if with_plugins:
         for plugin in PluginRegistry().get_plugins(PluginType.COLUMN_GENERATOR):
             registry.register(

data_designer/engine/column_generators/utils/judge_score_factory.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from enum import Enum
-from typing import Type
 
 from pydantic import BaseModel, ConfigDict, Field, create_model
 
@@ -19,7 +18,7 @@ class BaseJudgeResponse(BaseModel):
     reasoning: str = Field(..., description="Reasoning for the assigned score.")
 
 
-def _stringify_scoring(options: dict, enum_type: Type[Enum]) -> str:
+def _stringify_scoring(options: dict, enum_type: type[Enum]) -> str:
     """Convert score descriptions into a single text block."""
     list_block = "\n".join(
         [SCORING_FORMAT.format(score=score, description=description) for score, description in options.items()]
@@ -27,7 +26,7 @@ def _stringify_scoring(options: dict, enum_type: Type[Enum]) -> str:
     return SCORE_FIELD_DESCRIPTION_FORMAT.format(enum_name=enum_type.__name__, scoring=list_block)
 
 
-def create_judge_response_model(score: Score) -> Type[BaseJudgeResponse]:
+def create_judge_response_model(score: Score) -> type[BaseJudgeResponse]:
     """Create a JudgeResponse data type."""
     enum_members = {}
     for option in score.options.keys():
@@ -46,12 +45,12 @@ def create_judge_response_model(score: Score) -> Type[BaseJudgeResponse]:
 
 
 def create_judge_structured_output_model(
-    judge_responses: list[Type[BaseJudgeResponse]],
-) -> Type[BaseModel]:
+    judge_responses: list[type[BaseJudgeResponse]],
+) -> type[BaseModel]:
     """Create a JudgeStructuredOutput class dynamically."""
     return create_model(
         "JudgeStructuredOutput",
         __doc__=f"Response schema for scores with the following names: {[response.__name__ for response in judge_responses]}.",
         __base__=BaseModel,
-        **{response.__name__.lower(): (response, ...) for response in judge_responses},
+        **{response.__name__: (response, ...) for response in judge_responses},
     )

data_designer/engine/configurable_task.py

@@ -3,7 +3,7 @@
 
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Generic, Type, TypeVar, get_origin
+from typing import Generic, TypeVar, get_origin
 
 import pandas as pd
 
@@ -30,7 +30,7 @@ class ConfigurableTask(ABC, Generic[TaskConfigT]):
         self._initialize()
 
     @classmethod
-    def get_config_type(cls) -> Type[TaskConfigT]:
+    def get_config_type(cls) -> type[TaskConfigT]:
         for base in cls.__orig_bases__:
             if hasattr(base, "__args__") and len(base.__args__) == 1:
                 arg = base.__args__[0]

data_designer/engine/dataset_builders/artifact_storage.py

@@ -7,7 +7,6 @@ import shutil
 from datetime import datetime
 from functools import cached_property
 from pathlib import Path
-from typing import Union
 
 import pandas as pd
 from pydantic import BaseModel, field_validator, model_validator
@@ -25,6 +24,7 @@ class BatchStage(StrEnum):
     PARTIAL_RESULT = "partial_results_path"
     FINAL_RESULT = "final_dataset_path"
     DROPPED_COLUMNS = "dropped_columns_dataset_path"
+    PROCESSORS_OUTPUTS = "processors_outputs_path"
 
 
 class ArtifactStorage(BaseModel):
@@ -33,6 +33,7 @@ class ArtifactStorage(BaseModel):
     final_dataset_folder_name: str = "parquet-files"
     partial_results_folder_name: str = "tmp-partial-parquet-files"
     dropped_columns_folder_name: str = "dropped-columns-parquet-files"
+    processors_outputs_folder_name: str = "processors-files"
 
     @property
     def artifact_path_exists(self) -> bool:
@@ -70,8 +71,12 @@ class ArtifactStorage(BaseModel):
     def partial_results_path(self) -> Path:
         return self.base_dataset_path / self.partial_results_folder_name
 
+    @property
+    def processors_outputs_path(self) -> Path:
+        return self.base_dataset_path / self.processors_outputs_folder_name
+
     @field_validator("artifact_path")
-    def validate_artifact_path(cls, v: Union[Path, str]) -> Path:
+    def validate_artifact_path(cls, v: Path | str) -> Path:
         v = Path(v)
         if not v.is_dir():
             raise ArtifactStorageError("Artifact path must exist and be a directory")
@@ -84,6 +89,7 @@ class ArtifactStorage(BaseModel):
             self.final_dataset_folder_name,
             self.partial_results_folder_name,
             self.dropped_columns_folder_name,
+            self.processors_outputs_folder_name,
         ]
 
         for name in folder_names:
@@ -169,9 +175,10 @@ class ArtifactStorage(BaseModel):
         batch_number: int,
         dataframe: pd.DataFrame,
         batch_stage: BatchStage,
+        subfolder: str | None = None,
     ) -> Path:
         file_path = self.create_batch_file_path(batch_number, batch_stage=batch_stage)
-        self.write_parquet_file(file_path.name, dataframe, batch_stage)
+        self.write_parquet_file(file_path.name, dataframe, batch_stage, subfolder=subfolder)
         return file_path
 
     def write_parquet_file(
@@ -179,9 +186,11 @@ class ArtifactStorage(BaseModel):
         parquet_file_name: str,
         dataframe: pd.DataFrame,
         batch_stage: BatchStage,
+        subfolder: str | None = None,
     ) -> Path:
-        self.mkdir_if_needed(self._get_stage_path(batch_stage))
-        file_path = self._get_stage_path(batch_stage) / parquet_file_name
+        subfolder = subfolder or ""
+        self.mkdir_if_needed(self._get_stage_path(batch_stage) / subfolder)
+        file_path = self._get_stage_path(batch_stage) / subfolder / parquet_file_name
         dataframe.to_parquet(file_path, index=False)
         return file_path
 

data_designer/engine/dataset_builders/column_wise_builder.py

@@ -10,15 +10,18 @@ from typing import Callable
 
 import pandas as pd
 
-from data_designer.config.column_types import ColumnConfigT, column_type_is_llm_generated
+from data_designer.config.column_types import ColumnConfigT, column_type_is_model_generated
 from data_designer.config.dataset_builders import BuildStage
 from data_designer.config.processors import (
     DropColumnsProcessorConfig,
     ProcessorConfig,
     ProcessorType,
 )
-from data_designer.engine.column_generators.generators.base import ColumnGenerator, GenerationStrategy
-from data_designer.engine.column_generators.generators.llm_generators import WithLLMGeneration
+from data_designer.engine.column_generators.generators.base import (
+    ColumnGenerator,
+    GenerationStrategy,
+    WithModelGeneration,
+)
 from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
 from data_designer.engine.dataset_builders.errors import DatasetGenerationError, DatasetProcessingError
 from data_designer.engine.dataset_builders.multi_column_configs import (
@@ -72,7 +75,7 @@ class ColumnWiseDatasetBuilder:
 
     @functools.cached_property
     def llm_generated_column_configs(self) -> list[ColumnConfigT]:
-        return [config for config in self.single_column_configs if column_type_is_llm_generated(config.column_type)]
+        return [config for config in self.single_column_configs if column_type_is_model_generated(config.column_type)]
 
     def build(
         self,
@@ -169,7 +172,7 @@ class ColumnWiseDatasetBuilder:
 
     def _run_cell_by_cell_generator(self, generator: ColumnGenerator) -> None:
         max_workers = MAX_CONCURRENCY_PER_NON_LLM_GENERATOR
-        if isinstance(generator, WithLLMGeneration):
+        if isinstance(generator, WithModelGeneration):
             max_workers = generator.inference_parameters.max_parallel_requests
         self._fan_out_with_threads(generator, max_workers=max_workers)
 
@@ -178,12 +181,12 @@ class ColumnWiseDatasetBuilder:
         self.batch_manager.update_records(df.to_dict(orient="records"))
 
     def _run_model_health_check_if_needed(self) -> bool:
-        if any(column_type_is_llm_generated(config.column_type) for config in self.single_column_configs):
+        if any(column_type_is_model_generated(config.column_type) for config in self.single_column_configs):
             self._resource_provider.model_registry.run_health_check(
-                set(config.model_alias for config in self.llm_generated_column_configs)
+                list(set(config.model_alias for config in self.llm_generated_column_configs))
             )
 
-    def _fan_out_with_threads(self, generator: WithLLMGeneration, max_workers: int) -> None:
+    def _fan_out_with_threads(self, generator: WithModelGeneration, max_workers: int) -> None:
         if generator.generation_strategy != GenerationStrategy.CELL_BY_CELL:
             raise DatasetGenerationError(
                 f"Generator {generator.metadata().name} is not a {GenerationStrategy.CELL_BY_CELL} "
@@ -244,6 +247,7 @@ class ColumnWiseDatasetBuilder:
             processors[BuildStage.POST_BATCH].append(  # as post-batch by default
                 DropColumnsProcessor(
                     config=DropColumnsProcessorConfig(
+                        name="default_drop_columns_processor",
                         column_names=columns_to_drop,
                         build_stage=BuildStage.POST_BATCH,
                     ),

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: