data-designer 0.1.0__py3-none-any.whl

data_designer/engine/models/registry.py

@@ -0,0 +1,118 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+
+from data_designer.config.models import 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.secret_resolver import SecretResolver
+
+logger = logging.getLogger(__name__)
+
+
+class ModelRegistry:
+    def __init__(
+        self,
+        *,
+        secret_resolver: SecretResolver,
+        model_provider_registry: ModelProviderRegistry,
+        model_configs: list[ModelConfig] | None = None,
+    ):
+        self._secret_resolver = secret_resolver
+        self._model_provider_registry = model_provider_registry
+        self._model_configs = {}
+        self._models = {}
+        self._set_model_configs(model_configs)
+
+    @property
+    def model_configs(self) -> dict[str, ModelConfig]:
+        return self._model_configs
+
+    @property
+    def models(self) -> dict[str, ModelFacade]:
+        return self._models
+
+    def register_model_configs(self, model_configs: list[ModelConfig]) -> None:
+        """Register a new Model configuration at runtime.
+
+        Args:
+            model_config: A new Model configuration to register. If an
+                Model configuration already exists in the registry
+                with the same name, then it will be overwritten.
+        """
+        self._set_model_configs(list(self._model_configs.values()) + model_configs)
+
+    def get_model(self, *, model_alias: str) -> ModelFacade:
+        # Check if model config exists first
+        if model_alias not in self._model_configs:
+            raise ValueError(f"No model config with alias {model_alias!r} found!")
+
+        # Lazy initialization: only create model facade when first requested
+        if model_alias not in self._models:
+            self._models[model_alias] = self._get_model(self._model_configs[model_alias])
+
+        return self._models[model_alias]
+
+    def get_model_config(self, *, model_alias: str) -> ModelConfig:
+        if model_alias not in self._model_configs:
+            raise ValueError(f"No model config with alias {model_alias!r} found!")
+        return self._model_configs[model_alias]
+
+    def get_model_usage_stats(self, total_time_elapsed: float) -> dict[str, dict]:
+        return {
+            model.model_name: model.usage_stats.get_usage_stats(total_time_elapsed=total_time_elapsed)
+            for model in self._models.values()
+            if model.usage_stats.has_usage
+        }
+
+    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:
+        logger.info("🩺 Running health checks for models...")
+        for model_alias in model_aliases:
+            model = self.get_model(model_alias=model_alias)
+            logger.info(
+                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",
+                )
+                logger.info("  |-- ✅ Passed!")
+            except Exception as e:
+                logger.error("  |-- ❌ Failed!")
+                raise e
+
+    def _set_model_configs(self, model_configs: list[ModelConfig]) -> None:
+        model_configs = model_configs or []
+        self._model_configs = {mc.alias: mc for mc in model_configs}
+        # 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,
+    )

data_designer/engine/models/usage.py

@@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+
+from pydantic import BaseModel, computed_field
+
+logger = logging.getLogger(__name__)
+
+
+class TokenUsageStats(BaseModel):
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+
+    @computed_field
+    def total_tokens(self) -> int:
+        return self.prompt_tokens + self.completion_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
+
+
+class RequestUsageStats(BaseModel):
+    successful_requests: int = 0
+    failed_requests: int = 0
+
+    @computed_field
+    def total_requests(self) -> int:
+        return self.successful_requests + self.failed_requests
+
+    @property
+    def has_usage(self) -> bool:
+        return self.total_requests > 0
+
+    def extend(self, *, successful_requests: int, failed_requests: int) -> None:
+        self.successful_requests += successful_requests
+        self.failed_requests += failed_requests
+
+
+class ModelUsageStats(BaseModel):
+    token_usage: TokenUsageStats = TokenUsageStats()
+    request_usage: RequestUsageStats = RequestUsageStats()
+
+    @property
+    def has_usage(self) -> bool:
+        return self.token_usage.has_usage and self.request_usage.has_usage
+
+    def extend(
+        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
+            )
+        if request_usage is not None:
+            self.request_usage.extend(
+                successful_requests=request_usage.successful_requests, failed_requests=request_usage.failed_requests
+            )
+
+    def get_usage_stats(self, *, total_time_elapsed: float) -> dict:
+        return self.model_dump() | {
+            "tokens_per_second": int(self.token_usage.total_tokens / total_time_elapsed)
+            if total_time_elapsed > 0
+            else 0,
+            "requests_per_minute": int(self.request_usage.total_requests / total_time_elapsed * 60)
+            if total_time_elapsed > 0
+            else 0,
+        }

data_designer/engine/models/utils.py

@@ -0,0 +1,38 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def prompt_to_messages(
+    *,
+    user_prompt: str,
+    system_prompt: str | None = None,
+    multi_modal_context: list[dict[str, Any]] | None = None,
+) -> list[dict[str, str | list[dict]]]:
+    """Convert a user and system prompt into Messages format.
+
+    Args:
+        user_prompt (str): A user prompt.
+        system_prompt (str, optional): An optional system prompt.
+    """
+    user_content = user_prompt
+    if multi_modal_context and len(multi_modal_context) > 0:
+        user_content = []
+        user_content.append({"type": "text", "text": user_prompt})
+        for context in multi_modal_context:
+            user_content.append(context)
+    return (
+        [
+            str_to_message(content=system_prompt, role="system"),
+            str_to_message(content=user_content, role="user"),
+        ]
+        if system_prompt
+        else [str_to_message(content=user_content, role="user")]
+    )
+
+
+def str_to_message(content: str | list[dict], role: str = "user") -> dict[str, str | list[dict]]:
+    return {"content": content, "role": role}

data_designer/engine/processing/ginja/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

data_designer/engine/processing/ginja/ast.py

@@ -0,0 +1,64 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from collections import deque
+from typing import Optional, Type
+
+from jinja2 import nodes as j_nodes
+
+
+def ast_max_depth(node: j_nodes.Node) -> int:
+    """Calculate the depth of a Jinja AST from a given node.
+
+    Args:
+        node (jinja2.nodes.Node): The starting Jinja2 AST node
+
+    Returns:
+        int: The maximum depth of the tree
+    """
+    # Each entry is (node, depth)
+    queue = deque([(node, 1)])
+    max_depth = 0
+
+    while queue:
+        current_node, current_depth = queue.popleft()
+
+        # Update maximum depth seen so far
+        max_depth = max(max_depth, current_depth)
+
+        # Add all children with incremented depth
+        for child in current_node.iter_child_nodes():
+            queue.append((child, current_depth + 1))
+
+    return max_depth
+
+
+def ast_descendant_count(ast: j_nodes.Node, only_type: Optional[Type[j_nodes.Node]] = None) -> int:
+    """Count the number of nodes which descend from the given node.
+
+    Args:
+        ast (jinja2.nodes.Node): The starting Jinja2 AST node
+        only_type (Type[jinja2.nodes.Node]): If specified, then only
+            nodes of this type will be counted.
+
+    Returns:
+        int: The number of nodes descended from the given node.
+    """
+    if only_type is None:
+        only_type = j_nodes.Node
+
+    return len(list(ast.find_all(only_type)))
+
+
+def ast_count_name_references(ast: j_nodes.Node, name: str) -> int:
+    """Count the number of nodes descended from the current that refer to name.
+
+    Args:
+        ast (jinja2.nodes.Node): The starting Jinja2 AST node
+
+    Returns:
+        int: The number of nodes descended from the provided node whose
+            name field matches the given name.
+    """
+    referenced_names = [node.name for node in ast.find_all(j_nodes.Name) if node.name == name]
+    return len(referenced_names)