data-designer-engine 0.4.0__py3-none-any.whl

data_designer/engine/models/recipes/base.py

@@ -0,0 +1,81 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import abc
+from collections.abc import Callable
+from typing import Generic, TypeVar
+
+T = TypeVar("T")
+
+
+class ResponseRecipe(abc.ABC, Generic[T]):
+    """Base class for defining response recipes.
+
+    Response recipes contain all necessary information for
+    getting an LLM to perform a particular common task,
+    like outputting code in a desired format or following
+    structured output.
+    """
+
+    @abc.abstractmethod
+    def _build_parser_fn(self) -> Callable[[str], T]:
+        """Build the recipe's output parser function."""
+        ...
+
+    @property
+    @abc.abstractmethod
+    def example_template(self) -> str: ...
+
+    @abc.abstractmethod
+    def serialize_output(self, output: T) -> str:
+        """Serialize an instance of the parser output."""
+        ...
+
+    @abc.abstractmethod
+    def deserialize_output(self, serialized_output: str) -> T:
+        """Deserialize a serialized instance of the parser output."""
+        ...
+
+    def __init__(self):
+        self._parse_fn = self._build_parser_fn()
+
+    @property
+    def task_instructions(self) -> str | None:
+        """Specifies task instructions.
+
+        These instructions lay out the particular task information the
+        LLM requires in order to carry out the function of the recipe.
+        """
+        return None
+
+    def parse(self, response: str) -> T:
+        """Apply the recipe's parser to a raw model output."""
+        return self._parse_fn(response)
+
+    def generate_response_example(self, example: T) -> str:
+        """Create a serialized response example that the parser would admit."""
+        return self.example_template.format(example=example)
+
+    def apply_recipe_to_user_prompt(self, user_prompt: str) -> str:
+        """Appends recipe specific task instructions if applicable.
+
+        Args:
+            user_prompt (str): User prompt to be appended with recipe specific task instructions if applicable.
+
+        Returns:
+            str: Final user prompt
+        """
+        return f"{user_prompt}\n\n{self.task_instructions}" if self.task_instructions else user_prompt
+
+    def apply_recipe_to_system_prompt(self, system_prompt: str | None) -> str:
+        """Appends recipe specific task instructions if applicable.
+
+        Args:
+            system_prompt (str): System prompt to be appended with recipe specific task instructions if applicable.
+
+        Returns:
+            str: Final system prompt
+        """
+        return system_prompt

data_designer/engine/models/recipes/response_recipes.py

@@ -0,0 +1,293 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import json
+from collections.abc import Callable
+
+from pydantic import BaseModel
+
+from data_designer.config.utils.code_lang import CodeLang
+from data_designer.engine.models.parsers.errors import ParserException
+from data_designer.engine.models.parsers.parser import LLMResponseParser
+from data_designer.engine.models.parsers.postprocessors import (
+    StructuredDataBlock,
+    deserialize_json_code,
+    merge_text_blocks,
+)
+from data_designer.engine.models.parsers.types import CodeBlock
+from data_designer.engine.models.recipes.base import (
+    ResponseRecipe,
+)
+from data_designer.engine.processing.gsonschema.validators import JSONSchemaValidationError, validate
+
+
+class TextResponseRecipe(ResponseRecipe[str]):
+    """Default text-parser.
+
+    This parser is meant to cover the "pass-through" case of natural language LLM responses.
+    """
+
+    @property
+    def example_template(self) -> str:
+        return "{example}"
+
+    def serialize_output(self, output: str) -> str:
+        return output
+
+    def deserialize_output(self, serialized_output: str) -> str:
+        return serialized_output
+
+    def _build_parser_fn(self) -> Callable[[str], str]:
+        parser = LLMResponseParser(
+            postprocessors=[
+                merge_text_blocks,
+            ]
+        )
+
+        return lambda x: parser.parse(x).response
+
+
+class StructuredResponseRecipe(ResponseRecipe[dict]):
+    """Recipe for structured responses.
+
+    This recipe is intended to cover the generic case of
+    prompting-based requests for structured data outputs,
+    and the structure in question is determined by a
+    provided JSON Schema.
+
+    The LLM's response us validated against the provided
+    JSON Schema, however the object returned is python
+    dictionary obtained from deserializing the LLM's
+    JSON response.
+    """
+
+    json_schema: dict
+    pruning: bool
+    no_extra_properties: bool
+
+    def __init__(
+        self,
+        json_schema: dict,
+        pruning: bool = True,
+        no_extra_properties: bool = True,
+        **kwargs,
+    ):
+        """Initialize StructuredResponseRecipe.
+
+        Args:
+            json_schema (dict): A target JSON schema that the LLM
+                should adhere to when making its response.
+            pruning (bool): If `True`, then any extra fields in the returned
+                JSON object will be removed. Otherwise, they are retained,
+                which could raise validation errors. Default=True
+            no_extra_properties (bool) If `True`, then validation will fail
+                if extra properties are encountered in the returned JSON response.
+                Default=True.
+        """
+        super().__init__(**kwargs)
+        self.json_schema = json_schema
+        self.pruning = pruning
+        self.no_extra_properties = no_extra_properties
+
+    @property
+    def task_instructions(self) -> str:
+        return (
+            "* Your response must be in JSON format.\n"
+            "* Your JSON response must be returned within a Markdown, ```json code fence.\n"
+            "* The JSON format is given as a JSON Schema description within <response_schema> markup tags.\n\n"
+            f"<response_schema>\n{self.schema}\n</response_schema>"
+        )
+
+    @property
+    def example_template(self) -> str:
+        return "```json\n{example}\n```"
+
+    def generate_response_example(self, example: dict) -> str:
+        return self.example_template.format(example=json.dumps(example))
+
+    @property
+    def schema(self) -> str:
+        return json.dumps(self.json_schema)
+
+    def serialize_output(self, output: dict) -> str:
+        return json.dumps(output, ensure_ascii=False)
+
+    def deserialize_output(self, serialized_output: str) -> dict:
+        return json.loads(serialized_output)
+
+    @property
+    def _validate_args(self):
+        return {
+            "schema": self.json_schema,
+            "pruning": self.pruning,
+            "no_extra_properties": self.no_extra_properties,
+        }
+
+    def _build_parser_fn(self) -> Callable[[str], dict]:
+        parser = LLMResponseParser(
+            postprocessors=[
+                merge_text_blocks,
+                deserialize_json_code,
+            ]
+        )
+
+        def parse_fn(response: str) -> dict:
+            try:
+                obj = parser.parse(response).filter([StructuredDataBlock]).parsed.pop().obj
+                return validate(obj, **self._validate_args)
+            except IndexError:
+                raise ParserException(
+                    "No parsable JSON structure within ```json markdown fence.",
+                    source=response,
+                ) from None
+            except JSONSchemaValidationError as exc:
+                raise ParserException(
+                    "Response doesn't match requested <response_schema>\n" + str(exc),
+                    source=response,
+                ) from None
+
+        return parse_fn
+
+
+class PydanticResponseRecipe(ResponseRecipe[BaseModel]):
+    """Recipe for Pydantic responses.
+
+    This recipe covers the case that we have a Pydantic
+    data type (BaseModel) already specified in the runtime
+    making LLM calls, and we want to obtain an object of
+    that same data type as the output from the parser.
+
+    This recipe operates in a very similar fashion to
+    `StructuredResponseRecipe` except that it is initialized
+    from a Pydantic `BaseModel` and does the extra step of
+    validating against that `BaseModel` using
+    `BaseModel.model_validate` for its return.
+    """
+
+    data_type: type[BaseModel]
+
+    def __init__(self, data_type: type[BaseModel], **kwargs):
+        """Initialize a PydanticResponseRecipe.
+
+        Args:
+            data_type (type(BaseModel)): The target Pydantic BaseModel
+                subclass that the LLM should adhere to in its response,
+                and defines the output type of the parser.
+        """
+        super().__init__(**kwargs)
+        self.data_type = data_type
+
+    @property
+    def schema(self) -> str:
+        return json.dumps(self.data_type.model_json_schema())
+
+    @property
+    def task_instructions(self) -> str:
+        return (
+            "* Your response must be in JSON format.\n"
+            "* Your JSON response must be returned within a Markdown, ```json code fence.\n"
+            "* The JSON format is given as a JSON Schema description within <response_schema> markup tags.\n\n"
+            f"<response_schema>\n{self.schema}\n</response_schema>"
+        )
+
+    @property
+    def example_template(self) -> str:
+        return "```json\n{example}\n```"
+
+    def generate_response_example(self, example: BaseModel) -> str:
+        return self.example_template.format(example=example.model_dump_json())
+
+    def serialize_output(self, output: BaseModel) -> str:
+        return output.model_dump_json()
+
+    def deserialize_output(self, serialized_output: str) -> BaseModel:
+        return self.data_type.model_validate_json(serialized_output)
+
+    def _build_parser_fn(self) -> Callable[[str], BaseModel]:
+        parser = LLMResponseParser(
+            postprocessors=[
+                merge_text_blocks,
+                deserialize_json_code,
+            ]
+        )
+
+        def parse_fn(response: str) -> BaseModel:
+            try:
+                obj = parser.parse(response).filter([StructuredDataBlock]).parsed.pop().obj
+                return self.data_type.model_validate(obj)
+            except IndexError:
+                raise ParserException(
+                    "No parsable JSON structure within ```json markdown fence.",
+                    source=response,
+                ) from None
+            except Exception as exc:
+                raise ParserException(
+                    "Response doesn't match requested <response_schema>\n" + str(exc),
+                    source=response,
+                ) from None
+
+        return parse_fn
+
+
+class CodeResponseRecipe(ResponseRecipe[str]):
+    """Obtain a code snippet from an LLM."""
+
+    def __init__(self, syntax: str | CodeLang, **kwargs):
+        """Initialize a CodeResponseRecipe.
+
+        Args:
+            syntax (str | CodeLang): The code syntax that the
+                LLM should adhere to, e.g. `"python"`, `"sql"`, etc.
+        """
+        super().__init__(**kwargs)
+        self.syntax = CodeLang.parse_lang(syntax)
+
+    @property
+    def task_instructions(self) -> str:
+        return (
+            f"* Your response must be code written in {self.syntax}.\n"
+            "* You will follow accepted and common syntax and best-practices.\n"
+            f"* Your response will be given in markdown code fences specifying the correct language.\n"
+            "* Only respond with a SINGLE code block."
+        )
+
+    @property
+    def example_template(self) -> str:
+        return f"```{self.syntax}\n{{example}}\n```\n"
+
+    def serialize_output(self, output: str) -> str:
+        return output
+
+    def deserialize_output(self, serialized_output: str) -> str:
+        return serialized_output
+
+    def _build_parser_fn(self) -> Callable[[str], str]:
+        parser = LLMResponseParser(
+            postprocessors=[
+                merge_text_blocks,
+            ]
+        )
+
+        def parse_fn(response: str) -> str:
+            try:
+                code_block = parser.parse(response).filter([CodeBlock]).parsed.pop()
+                # For the type checker -- should always pass
+                assert isinstance(code_block, CodeBlock)
+            except IndexError:
+                raise ParserException(
+                    "No parsable code response.",
+                    source=response,
+                ) from None
+
+            # Only report this as a parser error if there was a mismatch.
+            if code_block.code_lang and code_block.code_lang != self.syntax:
+                raise ParserException(
+                    f"Responded with code not matching the requested syntax ({self.syntax}).",
+                    source=response,
+                )
+
+            return code_block.code.strip()
+
+        return parse_fn

data_designer/engine/models/registry.py

@@ -0,0 +1,151 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+from data_designer.config.models import GenerationType, ModelConfig
+from data_designer.engine.model_provider import ModelProvider, ModelProviderRegistry
+from data_designer.engine.models.usage import ModelUsageStats, RequestUsageStats, TokenUsageStats
+from data_designer.engine.secret_resolver import SecretResolver
+
+if TYPE_CHECKING:
+    from data_designer.engine.models.facade import ModelFacade
+
+logger = logging.getLogger(__name__)
+
+
+class ModelRegistry:
+    def __init__(
+        self,
+        *,
+        secret_resolver: SecretResolver,
+        model_provider_registry: ModelProviderRegistry,
+        model_configs: list[ModelConfig] | None = None,
+        model_facade_factory: Callable[[ModelConfig, SecretResolver, ModelProviderRegistry], ModelFacade] | None = None,
+    ):
+        self._secret_resolver = secret_resolver
+        self._model_provider_registry = model_provider_registry
+        self._model_facade_factory = model_facade_factory
+        self._model_configs: dict[str, ModelConfig] = {}
+        self._models: dict[str, ModelFacade] = {}
+        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_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: list[str]) -> None:
+        logger.info("🩺 Running health checks for models...")
+        for model_alias in model_aliases:
+            model_config = self.get_model_config(model_alias=model_alias)
+            if model_config.skip_health_check:
+                logger.info(f"  |-- ⏭️  Skipping health check for model alias {model_alias!r} (skip_health_check=True)")
+                continue
+
+            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:
+                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!")
+                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:
+        if self._model_facade_factory is None:
+            raise RuntimeError("ModelRegistry was not initialized with a model_facade_factory")
+        return self._model_facade_factory(model_config, self._secret_resolver, self._model_provider_registry)