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

data_designer/engine/models/registry.py

@@ -5,10 +5,11 @@ 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
+from data_designer.engine.models.usage import ModelUsageStats, RequestUsageStats, TokenUsageStats
 from data_designer.engine.secret_resolver import SecretResolver
 
 logger = logging.getLogger(__name__)
@@ -25,7 +26,7 @@ class ModelRegistry:
         self._secret_resolver = secret_resolver
         self._model_provider_registry = model_provider_registry
         self._model_configs = {}
-        self._models = {}
+        self._models: dict[str, ModelFacade] = {}
         self._set_model_configs(model_configs)
 
     @property
@@ -69,11 +70,36 @@ class ModelRegistry:
             if model.usage_stats.has_usage
         }
 
+    def get_model_usage_snapshot(self) -> dict[str, ModelUsageStats]:
+        return {
+            model.model_name: model.usage_stats.model_copy(deep=True)
+            for model in self._models.values()
+            if model.usage_stats.has_usage
+        }
+
+    def get_usage_deltas(self, snapshot: dict[str, ModelUsageStats]) -> dict[str, ModelUsageStats]:
+        deltas = {}
+        for model_name, current in self.get_model_usage_snapshot().items():
+            prev = snapshot.get(model_name)
+            delta_input = current.token_usage.input_tokens - (prev.token_usage.input_tokens if prev else 0)
+            delta_output = current.token_usage.output_tokens - (prev.token_usage.output_tokens if prev else 0)
+            delta_successful = current.request_usage.successful_requests - (
+                prev.request_usage.successful_requests if prev else 0
+            )
+            delta_failed = current.request_usage.failed_requests - (prev.request_usage.failed_requests if prev else 0)
+
+            if delta_input > 0 or delta_output > 0 or delta_successful > 0 or delta_failed > 0:
+                deltas[model_name] = ModelUsageStats(
+                    token_usage=TokenUsageStats(input_tokens=delta_input, output_tokens=delta_output),
+                    request_usage=RequestUsageStats(successful_requests=delta_successful, failed_requests=delta_failed),
+                )
+        return deltas
+
     def get_model_provider(self, *, model_alias: str) -> ModelProvider:
         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 +107,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/telemetry.py

@@ -0,0 +1,355 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Telemetry handler for NeMo products.
+
+Environment variables:
+- NEMO_TELEMETRY_ENABLED: Whether telemetry is enabled.
+- NEMO_DEPLOYMENT_TYPE: The deployment type the event came from.
+- NEMO_TELEMETRY_ENDPOINT: The endpoint to send the telemetry events to.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import platform
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any, ClassVar
+
+import httpx
+from pydantic import BaseModel, Field
+
+TELEMETRY_ENABLED = os.getenv("NEMO_TELEMETRY_ENABLED", "true").lower() in ("1", "true", "yes")
+CLIENT_ID = "184482118588404"
+NEMO_TELEMETRY_VERSION = "nemo-telemetry/1.0"
+MAX_RETRIES = 3
+NEMO_TELEMETRY_ENDPOINT = os.getenv(
+    "NEMO_TELEMETRY_ENDPOINT", "https://events.telemetry.data.nvidia.com/v1.1/events/json"
+).lower()
+CPU_ARCHITECTURE = platform.uname().machine
+
+
+class NemoSourceEnum(str, Enum):
+    INFERENCE = "inference"
+    AUDITOR = "auditor"
+    DATADESIGNER = "datadesigner"
+    EVALUATOR = "evaluator"
+    GUARDRAILS = "guardrails"
+    UNDEFINED = "undefined"
+
+
+class DeploymentTypeEnum(str, Enum):
+    LIBRARY = "library"
+    API = "api"
+    UNDEFINED = "undefined"
+
+
+_deployment_type_raw = os.getenv("NEMO_DEPLOYMENT_TYPE", "library").lower()
+try:
+    DEPLOYMENT_TYPE = DeploymentTypeEnum(_deployment_type_raw)
+except ValueError:
+    valid_values = [e.value for e in DeploymentTypeEnum]
+    raise ValueError(
+        f"Invalid NEMO_DEPLOYMENT_TYPE: {_deployment_type_raw!r}. Must be one of: {valid_values}"
+    ) from None
+
+
+class TaskStatusEnum(str, Enum):
+    SUCCESS = "success"
+    FAILURE = "failure"
+    UNDEFINED = "undefined"
+
+
+class TelemetryEvent(BaseModel):
+    _event_name: ClassVar[str]  # Subclasses must define this
+    _schema_version: ClassVar[str] = "1.3"
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        super().__init_subclass__(**kwargs)
+        if "_event_name" not in cls.__dict__:
+            raise TypeError(f"{cls.__name__} must define '_event_name' class variable")
+
+
+class InferenceEvent(TelemetryEvent):
+    _event_name: ClassVar[str] = "inference_event"
+
+    nemo_source: NemoSourceEnum = Field(
+        ...,
+        alias="nemoSource",
+        description="The NeMo product that created the event (i.e. data-designer).",
+    )
+    task: str = Field(
+        ...,
+        description="The type of task that was performed that generated the inference event (i.e. preview-job, batch-job).",
+    )
+    task_status: TaskStatusEnum = Field(
+        ...,
+        alias="taskStatus",
+        description="The status of the task.",
+    )
+    deployment_type: DeploymentTypeEnum = Field(
+        default=DEPLOYMENT_TYPE,
+        alias="deploymentType",
+        description="The deployment type the event came from.",
+    )
+    model: str = Field(
+        ...,
+        description="The name of the model that was used.",
+    )
+    model_group: str = Field(
+        default="undefined",
+        alias="modelGroup",
+        description="An optional identifier to group models together.",
+    )
+    input_bytes: int = Field(
+        default=-1,
+        alias="inputBytes",
+        description="Number of bytes provided as input to the model. -1 if not available.",
+        ge=-9223372036854775808,
+        le=9223372036854775807,
+    )
+    input_tokens: int = Field(
+        default=-1,
+        alias="inputTokens",
+        description="Number of tokens provided as input to the model. -1 if not available.",
+        ge=-9223372036854775808,
+        le=9223372036854775807,
+    )
+    output_bytes: int = Field(
+        default=-1,
+        alias="outputBytes",
+        description="Number of bytes returned by the model. -1 if not available.",
+        ge=-9223372036854775808,
+        le=9223372036854775807,
+    )
+    output_tokens: int = Field(
+        default=-1,
+        alias="outputTokens",
+        description="Number of tokens returned by the model. -1 if not available.",
+        ge=-9223372036854775808,
+        le=9223372036854775807,
+    )
+
+    model_config = {"populate_by_name": True}
+
+
+@dataclass
+class QueuedEvent:
+    event: TelemetryEvent
+    timestamp: datetime
+    retry_count: int = 0
+
+
+def _get_iso_timestamp(dt: datetime | None = None) -> str:
+    if dt is None:
+        dt = datetime.now(timezone.utc)
+    return dt.strftime("%Y-%m-%dT%H:%M:%S.") + f"{dt.microsecond // 1000:03d}Z"
+
+
+def build_payload(
+    events: list[QueuedEvent], *, source_client_version: str, session_id: str = "undefined"
+) -> dict[str, Any]:
+    return {
+        "browserType": "undefined",  # do not change
+        "clientId": CLIENT_ID,
+        "clientType": "Native",  # do not change
+        "clientVariant": "Release",  # do not change
+        "clientVer": source_client_version,
+        "cpuArchitecture": CPU_ARCHITECTURE,
+        "deviceGdprBehOptIn": "None",  # do not change
+        "deviceGdprFuncOptIn": "None",  # do not change
+        "deviceGdprTechOptIn": "None",  # do not change
+        "deviceId": "undefined",  # do not change
+        "deviceMake": "undefined",  # do not change
+        "deviceModel": "undefined",  # do not change
+        "deviceOS": "undefined",  # do not change
+        "deviceOSVersion": "undefined",  # do not change
+        "deviceType": "undefined",  # do not change
+        "eventProtocol": "1.6",  # do not change
+        "eventSchemaVer": events[0].event._schema_version,
+        "eventSysVer": NEMO_TELEMETRY_VERSION,
+        "externalUserId": "undefined",  # do not change
+        "gdprBehOptIn": "None",  # do not change
+        "gdprFuncOptIn": "None",  # do not change
+        "gdprTechOptIn": "None",  # do not change
+        "idpId": "undefined",  # do not change
+        "integrationId": "undefined",  # do not change
+        "productName": "undefined",  # do not change
+        "productVersion": "undefined",  # do not change
+        "sentTs": _get_iso_timestamp(),
+        "sessionId": session_id,
+        "userId": "undefined",  # do not change
+        "events": [
+            {
+                "ts": _get_iso_timestamp(queued.timestamp),
+                "parameters": queued.event.model_dump(by_alias=True),
+                "name": queued.event._event_name,
+            }
+            for queued in events
+        ],
+    }
+
+
+class TelemetryHandler:
+    """
+    Handles telemetry event batching, flushing, and retry logic for NeMo products.
+
+    Args:
+        flush_interval_seconds (float): The interval in seconds to flush the events.
+        max_queue_size (int): The maximum number of events to queue before flushing.
+        max_retries (int): The maximum number of times to retry sending an event.
+        source_client_version (str): The version of the source client. This should be the version of
+            the actual NeMo product that is sending the events, typically the same as the version of
+            a PyPi package that a user would install.
+        session_id (str): An optional session ID to associate with the events.
+            This should be a unique identifier for the session, such as a UUID.
+            It is used to group events together.
+    """
+
+    def __init__(
+        self,
+        flush_interval_seconds: float = 120.0,
+        max_queue_size: int = 50,
+        max_retries: int = MAX_RETRIES,
+        source_client_version: str = "undefined",
+        session_id: str = "undefined",
+    ):
+        self._flush_interval = flush_interval_seconds
+        self._max_queue_size = max_queue_size
+        self._max_retries = max_retries
+        self._events: list[QueuedEvent] = []
+        self._dlq: list[QueuedEvent] = []  # Dead letter queue for retry
+        self._flush_signal = asyncio.Event()
+        self._timer_task: asyncio.Task | None = None
+        self._running = False
+        self._source_client_version = source_client_version
+        self._session_id = session_id
+
+    async def astart(self) -> None:
+        if self._running:
+            return
+        self._running = True
+        self._timer_task = asyncio.create_task(self._timer_loop())
+
+    async def astop(self) -> None:
+        self._running = False
+        self._flush_signal.set()
+        if self._timer_task:
+            self._timer_task.cancel()
+            try:
+                await self._timer_task
+            except asyncio.CancelledError:
+                pass
+            self._timer_task = None
+        await self._flush_events()
+
+    async def aflush(self) -> None:
+        self._flush_signal.set()
+
+    def start(self) -> None:
+        self._run_sync(self.astart())
+
+    def stop(self) -> None:
+        self._run_sync(self.astop())
+
+    def flush(self) -> None:
+        self._flush_signal.set()
+
+    def enqueue(self, event: TelemetryEvent) -> None:
+        if not TELEMETRY_ENABLED:
+            return
+        if not isinstance(event, TelemetryEvent):
+            # Silently fail as we prioritize not disrupting upstream call sites and telemetry is best effort
+            return
+        queued = QueuedEvent(event=event, timestamp=datetime.now(timezone.utc))
+        self._events.append(queued)
+        if len(self._events) >= self._max_queue_size:
+            self._flush_signal.set()
+
+    def _run_sync(self, coro: Any) -> Any:
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = None
+
+        if loop and loop.is_running():
+            import concurrent.futures
+
+            with concurrent.futures.ThreadPoolExecutor() as pool:
+                future = pool.submit(asyncio.run, coro)
+                return future.result()
+        else:
+            return asyncio.run(coro)
+
+    def __enter__(self) -> TelemetryHandler:
+        self.start()
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.stop()
+
+    async def __aenter__(self) -> TelemetryHandler:
+        await self.astart()
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        await self.astop()
+
+    async def _timer_loop(self) -> None:
+        while self._running:
+            try:
+                await asyncio.wait_for(
+                    self._flush_signal.wait(),
+                    timeout=self._flush_interval,
+                )
+            except asyncio.TimeoutError:
+                pass
+            self._flush_signal.clear()
+            await self._flush_events()
+
+    async def _flush_events(self) -> None:
+        dlq_events, self._dlq = self._dlq, []
+        new_events, self._events = self._events, []
+        events_to_send = dlq_events + new_events
+        if events_to_send:
+            await self._send_events(events_to_send)
+
+    async def _send_events(self, events: list[QueuedEvent]) -> None:
+        async with httpx.AsyncClient() as client:
+            await self._send_events_with_client(client, events)
+
+    async def _send_events_with_client(self, client: httpx.AsyncClient, events: list[QueuedEvent]) -> None:
+        if not events:
+            return
+
+        payload = build_payload(events, source_client_version=self._source_client_version, session_id=self._session_id)
+        try:
+            response = await client.post(NEMO_TELEMETRY_ENDPOINT, json=payload)
+            # 2xx, 400, 422 are all considered complete (no retry)
+            # 400/422 indicate bad payload which retrying won't fix
+            if response.status_code in (400, 422) or response.is_success:
+                return
+            # 413 (payload too large) - split and retry
+            if response.status_code == 413:
+                if len(events) == 1:
+                    # Can't split further, drop the event
+                    return
+                mid = len(events) // 2
+                await self._send_events_with_client(client, events[:mid])
+                await self._send_events_with_client(client, events[mid:])
+                return
+            if response.status_code == 408 or response.status_code >= 500:
+                self._add_to_dlq(events)
+        except httpx.HTTPError:
+            self._add_to_dlq(events)
+
+    def _add_to_dlq(self, events: list[QueuedEvent]) -> None:
+        for queued in events:
+            queued.retry_count += 1
+            if queued.retry_count > self._max_retries:
+                continue
+            self._dlq.append(queued)

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)]