data-designer 0.3.8rc2__py3-none-any.whl → 0.4.0rc1__py3-none-any.whl

data_designer/engine/models/recipes/base.py

@@ -1,81 +0,0 @@
-# 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

@@ -1,293 +0,0 @@
-# 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

@@ -1,146 +0,0 @@
-# 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 = 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)