.
+ drop_block = isinstance(parsed_block, tp.TextBlock) and not parsed_block.text.strip()
+
+ if not drop_block:
+ output.parsed.append(parsed_block)
+
+ # Check tails -- inelegant, but they're always text.
+ # Don't add the tail if it is just blank space.
+ if element.tail and element.tail.strip():
+ output.parsed.append(tp.TextBlock(text=element.tail))
+
+ return self.postprocess(output)
diff --git a/packages/data-designer-engine/src/data_designer/engine/models_v2/parsers/postprocessors.py b/packages/data-designer-engine/src/data_designer/engine/models_v2/parsers/postprocessors.py
new file mode 100644
index 00000000..4635562c
--- /dev/null
+++ b/packages/data-designer-engine/src/data_designer/engine/models_v2/parsers/postprocessors.py
@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import json_repair
+from pydantic import BaseModel, ValidationError
+
+from data_designer.engine.models.parsers.types import (
+ CodeBlock,
+ LLMStructuredResponse,
+ PydanticTypeBlock,
+ StructuredDataBlock,
+ TextBlock,
+)
+
+
+def merge_text_blocks(
+ structured_response: LLMStructuredResponse,
+) -> LLMStructuredResponse:
+ processed_response = structured_response.model_copy()
+ processed_response.parsed = []
+ accumulator = None
+ for block in structured_response.parsed:
+ if isinstance(block, TextBlock):
+ if accumulator is not None:
+ accumulator = TextBlock(text=accumulator.text + block.text)
+ else:
+ accumulator = block
+ else:
+ if accumulator is not None:
+ processed_response.parsed.append(accumulator)
+ accumulator = None
+
+ processed_response.parsed.append(block)
+
+ if accumulator:
+ processed_response.parsed.append(accumulator)
+
+ return processed_response
+
+
+def deserialize_json_code(
+ structured_response: LLMStructuredResponse,
+) -> LLMStructuredResponse:
+ processed_response = structured_response.model_copy()
+ processed_response.parsed = []
+
+ for block in structured_response.parsed:
+ if isinstance(block, CodeBlock) and block.code_lang == "json":
+ deserialized = json_repair.loads(block.code)
+
+ block = StructuredDataBlock(serialized=block.code, obj=deserialized)
+
+ processed_response.parsed.append(block)
+ else:
+ processed_response.parsed.append(block)
+
+ return processed_response
+
+
+class RealizePydanticTypes:
+ types: list[type[BaseModel]]
+
+ def __init__(self, types: list[type[BaseModel]]):
+ self.types = types
+
+ def _fit_types(self, obj: dict) -> BaseModel | None:
+ final_obj = None
+
+ for t in self.types:
+ try:
+ final_obj = t.model_validate(obj)
+ except ValidationError:
+ pass
+
+ return final_obj
+
+ def __call__(self, structured_response: LLMStructuredResponse) -> LLMStructuredResponse:
+ processed_response = structured_response.model_copy()
+ processed_response.parsed = []
+
+ for block in structured_response.parsed:
+ if isinstance(block, StructuredDataBlock):
+ new_block = block
+ pydantic_obj = self._fit_types(block.obj)
+ if pydantic_obj:
+ new_block = PydanticTypeBlock(serialized=block.serialized, obj=pydantic_obj)
+ processed_response.parsed.append(new_block)
+ else:
+ processed_response.parsed.append(block)
+
+ return processed_response
diff --git a/packages/data-designer-engine/src/data_designer/engine/models_v2/parsers/tag_parsers.py b/packages/data-designer-engine/src/data_designer/engine/models_v2/parsers/tag_parsers.py
new file mode 100644
index 00000000..ba09eed0
--- /dev/null
+++ b/packages/data-designer-engine/src/data_designer/engine/models_v2/parsers/tag_parsers.py
@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from lxml.etree import _Element
+
+from data_designer.engine.models.parsers.types import CodeBlock, TextBlock
+
+
+def text_parser(element: _Element) -> TextBlock:
+ return TextBlock(text=element.text if element.text else "")
+
+
+def text_parser_keep_markup(element: _Element) -> TextBlock:
+ body = element.text if element.text else ""
+ return TextBlock(text=f"<{element.tag}>{body}")
+
+
+def inline_code_parser(element: _Element) -> TextBlock:
+ return TextBlock(text=f"`{element.text if element.text else ''}`")
+
+
+def code_block_parser(element: _Element) -> CodeBlock:
+ """Parse a element node.
+
+ This parser handles the special case of Markdown->HTML conversion
+ for fenced code blocks. These take on the form:
+
+ ```xx
+ ...
+ ```
+
+ ...
+
+ This parser is intended to be attached to the special case of "pre.code"
+ tag hierarchies.
+
+ Syntax Handling
+
+ If the syntax is not specified, e.g. ``...`` or
+ ``...``, then the syntax field is returned
+ as None. However, the parser does not _enforce_ the prefix
+ `language-` on the value of the class attribute.
+ If it is not present, then the entire value
+
+ Args:
+ element (lxml.etree._Element): An element of the lxml-parsed
+ element tree.
+
+ Returns:
+ CodeBlock: Datat structured containing both the body of the code
+ as well as the specified synax of the code block.
+
+ """
+ prefix = "language-"
+ language_identifier = element.attrib.get("class", "")
+ language_identifier = language_identifier.removeprefix(prefix)
+ return CodeBlock(
+ code=element.text.strip() if element.text else "",
+ code_lang=language_identifier if language_identifier else None,
+ )
diff --git a/packages/data-designer-engine/src/data_designer/engine/models_v2/parsers/types.py b/packages/data-designer-engine/src/data_designer/engine/models_v2/parsers/types.py
new file mode 100644
index 00000000..81575ef9
--- /dev/null
+++ b/packages/data-designer-engine/src/data_designer/engine/models_v2/parsers/types.py
@@ -0,0 +1,84 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Any, Protocol, runtime_checkable
+
+from lxml.etree import _Element
+from pydantic import BaseModel, Field
+from typing_extensions import Self
+
+
+class LLMStructuredResponse(BaseModel):
+ """Output format for the LLM Response Parser."""
+
+ response: str = Field(description="Raw Markdown/Markup response received from the LLM and input to the parser.")
+ markup: str = Field(description="Markup/HTML resulting from running Markdown parsing on response.")
+ parsed: list[BaseModel] = Field(
+ default_factory=list,
+ description="Structured content parsed from markup. Elements of this list are in document-order.",
+ )
+
+ def head(self, n: int) -> Self:
+ """Retain only the first n elements of the parsed response."""
+ out = self.model_copy()
+ out.parsed = out.parsed[:n]
+ return out
+
+ def tail(self, n: int) -> Self:
+ """Retain only the last n elements of the parsed response."""
+ out = self.model_copy()
+ out.parsed = out.parsed[-n:]
+ return out
+
+ def filter(self, block_types: list[type[BaseModel]]) -> Self:
+ out = self.model_copy()
+ out.parsed = [b for b in out.parsed if isinstance(b, tuple(block_types))]
+ return out
+
+
+@runtime_checkable
+class TagParser(Protocol):
+ """Protocol for tag parsing implementations.
+
+ All TagParsers are objects which can take as input an `lxml`
+ element, do some computation, and return some kind of structured
+ output, represented as a subclass of Pydantic `BaseModel`.
+ This protocol implementation can cover both classes as well
+ as curried functions as parsers (e.g. `partial`).
+ """
+
+ def __call__(self, element: _Element) -> BaseModel: ...
+
+
+@runtime_checkable
+class PostProcessor(Protocol):
+ """Protocol for parsed output postprocessing implementations.
+
+ Implementations of this protocol are used to transform the results of
+ the LLM response parser while retaining the same output structure.
+ This is done so that PostProcessor implementations can be chained
+ together.
+ """
+
+ def __call__(self, structured_response: LLMStructuredResponse) -> LLMStructuredResponse: ...
+
+
+class TextBlock(BaseModel):
+ text: str
+
+
+class CodeBlock(BaseModel):
+ code: str
+ code_lang: str | None = None
+
+
+class StructuredDataBlock(BaseModel):
+ serialized: str
+ obj: Any
+
+
+class PydanticTypeBlock(BaseModel):
+ serialized: str
+ obj: BaseModel
diff --git a/packages/data-designer-engine/src/data_designer/engine/models_v2/recipes/base.py b/packages/data-designer-engine/src/data_designer/engine/models_v2/recipes/base.py
new file mode 100644
index 00000000..ab3e313a
--- /dev/null
+++ b/packages/data-designer-engine/src/data_designer/engine/models_v2/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
diff --git a/packages/data-designer-engine/src/data_designer/engine/models_v2/recipes/response_recipes.py b/packages/data-designer-engine/src/data_designer/engine/models_v2/recipes/response_recipes.py
new file mode 100644
index 00000000..deba050e
--- /dev/null
+++ b/packages/data-designer-engine/src/data_designer/engine/models_v2/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 markup tags.\n\n"
+ f"\n{self.schema}\n"
+ )
+
+ @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 \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 markup tags.\n\n"
+ f"\n{self.schema}\n"
+ )
+
+ @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 \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
diff --git a/packages/data-designer-engine/src/data_designer/engine/models_v2/registry.py b/packages/data-designer-engine/src/data_designer/engine/models_v2/registry.py
new file mode 100644
index 00000000..8f6a0e9d
--- /dev/null
+++ b/packages/data-designer-engine/src/data_designer/engine/models_v2/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)
diff --git a/packages/data-designer-engine/src/data_designer/engine/models_v2/telemetry.py b/packages/data-designer-engine/src/data_designer/engine/models_v2/telemetry.py
new file mode 100644
index 00000000..f90e39c0
--- /dev/null
+++ b/packages/data-designer-engine/src/data_designer/engine/models_v2/telemetry.py
@@ -0,0 +1,362 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 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.
+- NEMO_SESSION_PREFIX: Optional prefix to add to session IDs.
+"""
+
+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
+
+from pydantic import BaseModel, Field
+
+from data_designer.lazy_heavy_imports import httpx
+
+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
+SESSION_PREFIX = os.getenv("NEMO_SESSION_PREFIX")
+
+
+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
+ # Apply session prefix if environment variable is set
+ if SESSION_PREFIX:
+ self._session_id = f"{SESSION_PREFIX}{session_id}"
+ else:
+ 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)
diff --git a/packages/data-designer-engine/src/data_designer/engine/models_v2/usage.py b/packages/data-designer-engine/src/data_designer/engine/models_v2/usage.py
new file mode 100644
index 00000000..b239a9f3
--- /dev/null
+++ b/packages/data-designer-engine/src/data_designer/engine/models_v2/usage.py
@@ -0,0 +1,73 @@
+# 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 pydantic import BaseModel, computed_field
+
+logger = logging.getLogger(__name__)
+
+
+class TokenUsageStats(BaseModel):
+ input_tokens: int = 0
+ output_tokens: int = 0
+
+ @computed_field
+ def total_tokens(self) -> int:
+ return self.input_tokens + self.output_tokens
+
+ @property
+ def has_usage(self) -> bool:
+ return self.total_tokens > 0
+
+ def extend(self, *, input_tokens: int, output_tokens: int) -> None:
+ self.input_tokens += input_tokens
+ self.output_tokens += output_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(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
+ )
+
+ 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,
+ }
diff --git a/packages/data-designer-engine/src/data_designer/engine/models_v2/utils.py b/packages/data-designer-engine/src/data_designer/engine/models_v2/utils.py
new file mode 100644
index 00000000..6c0418c4
--- /dev/null
+++ b/packages/data-designer-engine/src/data_designer/engine/models_v2/utils.py
@@ -0,0 +1,101 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+
+@dataclass
+class ChatMessage:
+ """A chat message in an LLM conversation.
+
+ This dataclass represents messages exchanged in a conversation with an LLM,
+ supporting various message types including user prompts, assistant responses,
+ system instructions, and tool interactions.
+
+ Attributes:
+ role: The role of the message sender. One of 'user', 'assistant', 'system', or 'tool'.
+ content: The message content. Can be a string or a list of content blocks
+ for multimodal messages (e.g., text + images).
+ reasoning_content: Optional reasoning/thinking content from the assistant,
+ typically from extended thinking or chain-of-thought models.
+ tool_calls: Optional list of tool calls requested by the assistant.
+ Each tool call contains 'id', 'type', and 'function' keys.
+ tool_call_id: Optional ID linking a tool response to its corresponding
+ tool call. Required for messages with role='tool'.
+ """
+
+ role: Literal["user", "assistant", "system", "tool"]
+ content: str | list[dict[str, Any]] = ""
+ reasoning_content: str | None = None
+ tool_calls: list[dict[str, Any]] = field(default_factory=list)
+ tool_call_id: str | None = None
+
+ def to_dict(self) -> dict[str, Any]:
+ """Convert the message to a dictionary format for API calls.
+
+ Returns:
+ A dictionary containing the message fields. Only includes non-empty
+ optional fields to keep the output clean.
+ """
+ result: dict[str, Any] = {"role": self.role, "content": self.content}
+ if self.reasoning_content:
+ result["reasoning_content"] = self.reasoning_content
+ if self.tool_calls:
+ result["tool_calls"] = self.tool_calls
+ if self.tool_call_id:
+ result["tool_call_id"] = self.tool_call_id
+ return result
+
+ @classmethod
+ def as_user(cls, content: str | list[dict[str, Any]]) -> ChatMessage:
+ """Create a user message."""
+ return cls(role="user", content=content)
+
+ @classmethod
+ def as_assistant(
+ cls,
+ content: str = "",
+ reasoning_content: str | None = None,
+ tool_calls: list[dict[str, Any]] | None = None,
+ ) -> ChatMessage:
+ """Create an assistant message."""
+ return cls(
+ role="assistant",
+ content=content,
+ reasoning_content=reasoning_content,
+ tool_calls=tool_calls or [],
+ )
+
+ @classmethod
+ def as_system(cls, content: str) -> ChatMessage:
+ """Create a system message."""
+ return cls(role="system", content=content)
+
+ @classmethod
+ def as_tool(cls, content: str, tool_call_id: str) -> ChatMessage:
+ """Create a tool response message."""
+ return cls(role="tool", content=content, tool_call_id=tool_call_id)
+
+
+def prompt_to_messages(
+ *,
+ user_prompt: str,
+ system_prompt: str | None = None,
+ multi_modal_context: list[dict[str, Any]] | None = None,
+) -> list[ChatMessage]:
+ """Convert a user and system prompt into ChatMessage list.
+
+ Args:
+ user_prompt (str): A user prompt.
+ system_prompt (str, optional): An optional system prompt.
+ """
+ user_content: str | list[dict[str, Any]] = user_prompt
+ if multi_modal_context:
+ user_content = [*multi_modal_context, {"type": "text", "text": user_prompt}]
+
+ if system_prompt:
+ return [ChatMessage.as_system(system_prompt), ChatMessage.as_user(user_content)]
+ return [ChatMessage.as_user(user_content)]
diff --git a/packages/data-designer-engine/tests/engine/models/test_async_engine_switch.py b/packages/data-designer-engine/tests/engine/models/test_async_engine_switch.py
new file mode 100644
index 00000000..951efa07
--- /dev/null
+++ b/packages/data-designer-engine/tests/engine/models/test_async_engine_switch.py
@@ -0,0 +1,61 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import importlib
+import sys
+from pathlib import Path
+from types import ModuleType
+
+import pytest
+
+_MODEL_PREFIXES = (
+ "data_designer.engine.models",
+ "data_designer.engine.models_v2",
+)
+
+
+def _matches_models_namespace(module_name: str) -> bool:
+ return any(module_name == prefix or module_name.startswith(f"{prefix}.") for prefix in _MODEL_PREFIXES)
+
+
+def _purge_models_modules() -> dict[str, ModuleType]:
+ saved: dict[str, ModuleType] = {}
+ for module_name in list(sys.modules):
+ if _matches_models_namespace(module_name):
+ module = sys.modules.pop(module_name, None)
+ if isinstance(module, ModuleType):
+ saved[module_name] = module
+ return saved
+
+
+def _restore_models_modules(saved: dict[str, ModuleType]) -> None:
+ for module_name in list(sys.modules):
+ if _matches_models_namespace(module_name):
+ sys.modules.pop(module_name, None)
+ sys.modules.update(saved)
+
+
+def _module_path_parts(module: ModuleType) -> tuple[str, ...]:
+ module_file = module.__file__
+ assert module_file is not None
+ return Path(module_file).parts
+
+
+def test_async_engine_env_switch(monkeypatch: pytest.MonkeyPatch) -> None:
+ saved_modules = _purge_models_modules()
+ try:
+ monkeypatch.delenv("DATA_DESIGNER_ASYNC_ENGINE", raising=False)
+ default_facade = importlib.import_module("data_designer.engine.models.facade")
+ default_parts = _module_path_parts(default_facade)
+ assert "models_v2" not in default_parts
+ assert "models" in default_parts
+
+ _purge_models_modules()
+ monkeypatch.setenv("DATA_DESIGNER_ASYNC_ENGINE", "1")
+ async_facade = importlib.import_module("data_designer.engine.models.facade")
+ async_parts = _module_path_parts(async_facade)
+ assert "models_v2" in async_parts
+ finally:
+ _restore_models_modules(saved_modules)
diff --git a/packages/data-designer-engine/tests/engine/models/test_facade.py b/packages/data-designer-engine/tests/engine/models/test_facade.py
index 89eb9b10..81810392 100644
--- a/packages/data-designer-engine/tests/engine/models/test_facade.py
+++ b/packages/data-designer-engine/tests/engine/models/test_facade.py
@@ -9,7 +9,7 @@
from data_designer.engine.mcp.errors import MCPConfigurationError
from data_designer.engine.models.errors import ModelGenerationValidationFailureError
-from data_designer.engine.models.facade import ModelFacade
+from data_designer.engine.models.facade import CustomRouter, ModelFacade
from data_designer.engine.models.parsers.errors import ParserException
from data_designer.engine.models.utils import ChatMessage
@@ -77,18 +77,18 @@ def stub_expected_embedding_response():
(3, 3, 16),
],
)
-@patch("data_designer.engine.models.facade.ModelFacade.completion", autospec=True)
+@patch.object(ModelFacade, "completion", autospec=True)
def test_generate(
- mock_completion,
- stub_model_facade,
- max_correction_steps,
- max_conversation_restarts,
- total_calls,
-):
+ mock_completion: Any,
+ stub_model_facade: ModelFacade,
+ max_correction_steps: int,
+ max_conversation_restarts: int,
+ total_calls: int,
+) -> None:
bad_response = mock_oai_response_object("bad response")
mock_completion.side_effect = lambda *args, **kwargs: bad_response
- def _failing_parser(response: str):
+ def _failing_parser(response: str) -> str:
raise ParserException("parser exception")
with pytest.raises(ModelGenerationValidationFailureError):
@@ -119,7 +119,7 @@ def _failing_parser(response: str):
("hello!", [ChatMessage.as_system("hello!"), ChatMessage.as_user("does not matter")]),
],
)
-@patch("data_designer.engine.models.facade.ModelFacade.completion", autospec=True)
+@patch.object(ModelFacade, "completion", autospec=True)
def test_generate_with_system_prompt(
mock_completion: Any,
stub_model_facade: ModelFacade,
@@ -183,7 +183,7 @@ def test_consolidate_kwargs(stub_model_configs, stub_model_facade):
True,
],
)
-@patch("data_designer.engine.models.facade.CustomRouter.completion", autospec=True)
+@patch.object(CustomRouter, "completion", autospec=True)
def test_completion_success(
mock_router_completion: Any,
stub_completion_messages: list[ChatMessage],
@@ -204,7 +204,7 @@ def test_completion_success(
}
-@patch("data_designer.engine.models.facade.CustomRouter.completion", autospec=True)
+@patch.object(CustomRouter, "completion", autospec=True)
def test_completion_with_exception(
mock_router_completion: Any,
stub_completion_messages: list[ChatMessage],
@@ -216,7 +216,7 @@ def test_completion_with_exception(
stub_model_facade.completion(stub_completion_messages)
-@patch("data_designer.engine.models.facade.CustomRouter.completion", autospec=True)
+@patch.object(CustomRouter, "completion", autospec=True)
def test_completion_with_kwargs(
mock_router_completion: Any,
stub_completion_messages: list[ChatMessage],
@@ -240,29 +240,36 @@ def mock_completion(self: Any, model: str, messages: list[dict[str, Any]], **kwa
assert captured_kwargs == {**stub_model_configs[0].inference_parameters.generate_kwargs, **kwargs}
-@patch("data_designer.engine.models.facade.CustomRouter.embedding", autospec=True)
-def test_generate_text_embeddings_success(mock_router_embedding, stub_model_facade, stub_expected_embedding_response):
+@patch.object(CustomRouter, "embedding", autospec=True)
+def test_generate_text_embeddings_success(
+ mock_router_embedding: Any,
+ stub_model_facade: ModelFacade,
+ stub_expected_embedding_response: EmbeddingResponse,
+) -> None:
mock_router_embedding.side_effect = lambda self, model, input, **kwargs: stub_expected_embedding_response
input_texts = ["test1", "test2"]
result = stub_model_facade.generate_text_embeddings(input_texts)
assert result == [data["embedding"] for data in stub_expected_embedding_response.data]
-@patch("data_designer.engine.models.facade.CustomRouter.embedding", autospec=True)
-def test_generate_text_embeddings_with_exception(mock_router_embedding, stub_model_facade):
+@patch.object(CustomRouter, "embedding", autospec=True)
+def test_generate_text_embeddings_with_exception(mock_router_embedding: Any, stub_model_facade: ModelFacade) -> None:
mock_router_embedding.side_effect = Exception("Router error")
with pytest.raises(Exception, match="Router error"):
stub_model_facade.generate_text_embeddings(["test1", "test2"])
-@patch("data_designer.engine.models.facade.CustomRouter.embedding", autospec=True)
+@patch.object(CustomRouter, "embedding", autospec=True)
def test_generate_text_embeddings_with_kwargs(
- mock_router_embedding, stub_model_configs, stub_model_facade, stub_expected_embedding_response
-):
+ mock_router_embedding: Any,
+ stub_model_configs: Any,
+ stub_model_facade: ModelFacade,
+ stub_expected_embedding_response: EmbeddingResponse,
+) -> None:
captured_kwargs = {}
- def mock_embedding(self, model, input, **kwargs):
+ def mock_embedding(self: Any, model: str, input: list[str], **kwargs: Any) -> EmbeddingResponse:
captured_kwargs.update(kwargs)
return stub_expected_embedding_response
diff --git a/packages/data-designer-engine/tests/engine/models/test_litellm_overrides.py b/packages/data-designer-engine/tests/engine/models/test_litellm_overrides.py
index 05f3ed79..6ccbe63a 100644
--- a/packages/data-designer-engine/tests/engine/models/test_litellm_overrides.py
+++ b/packages/data-designer-engine/tests/engine/models/test_litellm_overrides.py
@@ -8,6 +8,7 @@
import litellm
import pytest
+from data_designer.engine.models import litellm_overrides
from data_designer.engine.models.litellm_overrides import (
DEFAULT_MAX_CALLBACKS,
CustomRouter,
@@ -56,9 +57,9 @@ def test_apply_litellm_patches_no_exceptions():
pytest.fail(f"apply_litellm_patches() raised an unexpected exception: {e}")
-@patch("data_designer.engine.models.litellm_overrides.quiet_noisy_logger", autospec=True)
-def test_apply_litellm_patches(mock_quiet_noisy_logger):
- apply_litellm_patches()
+@patch.object(litellm_overrides, "quiet_noisy_logger", autospec=True)
+def test_apply_litellm_patches(mock_quiet_noisy_logger: object) -> None:
+ litellm_overrides.apply_litellm_patches()
assert isinstance(litellm.in_memory_llm_clients_cache, ThreadSafeCache)
assert (
litellm.litellm_core_utils.logging_callback_manager.LoggingCallbackManager.MAX_CALLBACKS
diff --git a/packages/data-designer-engine/tests/engine/models/test_model_registry.py b/packages/data-designer-engine/tests/engine/models/test_model_registry.py
index be109a01..a4bfc8e6 100644
--- a/packages/data-designer-engine/tests/engine/models/test_model_registry.py
+++ b/packages/data-designer-engine/tests/engine/models/test_model_registry.py
@@ -6,6 +6,7 @@
import pytest
from data_designer.config.models import ChatCompletionInferenceParams, ModelConfig
+from data_designer.engine.models import litellm_overrides
from data_designer.engine.models.errors import ModelAuthenticationError
from data_designer.engine.models.facade import ModelFacade
from data_designer.engine.models.factory import create_model_registry
@@ -40,10 +41,13 @@ def stub_no_usage_config():
)
-@patch("data_designer.engine.models.litellm_overrides.apply_litellm_patches", autospec=True)
+@patch.object(litellm_overrides, "apply_litellm_patches", autospec=True)
def test_create_model_registry(
- mock_apply_litellm_patches, stub_model_configs, stub_secrets_resolver, stub_model_provider_registry
-):
+ mock_apply_litellm_patches: object,
+ stub_model_configs: list[ModelConfig],
+ stub_secrets_resolver: object,
+ stub_model_provider_registry: object,
+) -> None:
model_registry = create_model_registry(
model_configs=stub_model_configs,
secret_resolver=stub_secrets_resolver,
@@ -272,20 +276,26 @@ def test_get_usage_deltas(
assert deltas == {}
-@patch("data_designer.engine.models.facade.ModelFacade.generate_text_embeddings", autospec=True)
-@patch("data_designer.engine.models.facade.ModelFacade.completion", autospec=True)
-def test_run_health_check_success(mock_completion, mock_generate_text_embeddings, stub_model_registry):
+@patch.object(ModelFacade, "generate_text_embeddings", autospec=True)
+@patch.object(ModelFacade, "completion", autospec=True)
+def test_run_health_check_success(
+ mock_completion: object,
+ mock_generate_text_embeddings: object,
+ stub_model_registry: ModelRegistry,
+) -> None:
model_aliases = {"stub-text", "stub-reasoning", "stub-embedding"}
stub_model_registry.run_health_check(model_aliases)
assert mock_completion.call_count == 2
assert mock_generate_text_embeddings.call_count == 1
-@patch("data_designer.engine.models.facade.ModelFacade.generate_text_embeddings", autospec=True)
-@patch("data_designer.engine.models.facade.ModelFacade.completion", autospec=True)
+@patch.object(ModelFacade, "generate_text_embeddings", autospec=True)
+@patch.object(ModelFacade, "completion", autospec=True)
def test_run_health_check_completion_authentication_error(
- mock_completion, mock_generate_text_embeddings, stub_model_registry
-):
+ mock_completion: object,
+ mock_generate_text_embeddings: object,
+ stub_model_registry: ModelRegistry,
+) -> None:
auth_error = ModelAuthenticationError("Invalid API key for completion model")
mock_completion.side_effect = auth_error
model_aliases = ["stub-text", "stub-reasoning", "stub-embedding"]
@@ -297,11 +307,13 @@ def test_run_health_check_completion_authentication_error(
mock_generate_text_embeddings.assert_not_called()
-@patch("data_designer.engine.models.facade.ModelFacade.generate_text_embeddings", autospec=True)
-@patch("data_designer.engine.models.facade.ModelFacade.completion", autospec=True)
+@patch.object(ModelFacade, "generate_text_embeddings", autospec=True)
+@patch.object(ModelFacade, "completion", autospec=True)
def test_run_health_check_embedding_authentication_error(
- mock_completion, mock_generate_text_embeddings, stub_model_registry
-):
+ mock_completion: object,
+ mock_generate_text_embeddings: object,
+ stub_model_registry: ModelRegistry,
+) -> None:
auth_error = ModelAuthenticationError("Invalid API key for embedding model")
mock_generate_text_embeddings.side_effect = auth_error
model_aliases = ["stub-text", "stub-reasoning", "stub-embedding"]
@@ -313,12 +325,12 @@ def test_run_health_check_embedding_authentication_error(
mock_generate_text_embeddings.assert_called_once()
-@patch("data_designer.engine.models.facade.ModelFacade.completion", autospec=True)
+@patch.object(ModelFacade, "completion", autospec=True)
def test_run_health_check_skip_health_check_flag(
- mock_completion,
- stub_secrets_resolver,
- stub_model_provider_registry,
-):
+ mock_completion: object,
+ stub_secrets_resolver: object,
+ stub_model_provider_registry: object,
+) -> None:
# Create model configs: one with skip_health_check=True, others with default (False)
model_configs = [
ModelConfig(
diff --git a/packages/data-designer-engine/tests/engine/validators/test_sql.py b/packages/data-designer-engine/tests/engine/validators/test_sql.py
index 756d0856..3ab599e1 100644
--- a/packages/data-designer-engine/tests/engine/validators/test_sql.py
+++ b/packages/data-designer-engine/tests/engine/validators/test_sql.py
@@ -1,12 +1,17 @@
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
+from unittest.mock import patch
+
+import pytest
+
from data_designer.config.utils.code_lang import CodeLang
from data_designer.config.validator_params import CodeValidatorParams
+from data_designer.engine.validators import sql as sql_validator_module
from data_designer.engine.validators.sql import SQLValidator
-def test_valid_ansi_sql_code():
+def test_valid_ansi_sql_code() -> None:
sql_validator = SQLValidator(CodeValidatorParams(code_lang=CodeLang.SQL_ANSI))
code = "SELECT category, COUNT(*) as total_incidents FROM security_incidents_2 GROUP BY category;"
result = sql_validator.run_validation([{"sql": code}])
@@ -14,9 +19,31 @@ def test_valid_ansi_sql_code():
assert result.data[0].error_messages == ""
-def test_invalid_ansi_sql_code():
+def test_invalid_ansi_sql_code() -> None:
sql_validator = SQLValidator(CodeValidatorParams(code_lang=CodeLang.SQL_ANSI))
code = "NOT SQL"
result = sql_validator.run_validation([{"sql": code}])
assert not result.data[0].is_valid
assert result.data[0].error_messages == "PRS: Line 1, Position 1: Found unparsable section: 'NOT SQL'"
+
+
+def test_sql_validator_multi_column_input_raises() -> None:
+ sql_validator = SQLValidator(CodeValidatorParams(code_lang=CodeLang.SQL_ANSI))
+ with pytest.raises(ValueError, match="single column input"):
+ sql_validator.run_validation([{"sql": "SELECT 1", "extra": "ignored"}])
+
+
+def test_sql_validator_decimal_without_scale_fails() -> None:
+ sql_validator = SQLValidator(CodeValidatorParams(code_lang=CodeLang.SQL_ANSI))
+ code = "CREATE TABLE example (amount DECIMAL(10));"
+ result = sql_validator.run_validation([{"sql": code}])
+ assert not result.data[0].is_valid
+ assert "DECIMAL definitions without a scale" in result.data[0].error_messages
+
+
+def test_sql_validator_handles_lint_exception() -> None:
+ sql_validator = SQLValidator(CodeValidatorParams(code_lang=CodeLang.SQL_ANSI))
+ with patch.object(sql_validator_module.sqlfluff, "lint", side_effect=RuntimeError("boom")):
+ result = sql_validator.run_validation([{"sql": "SELECT 1"}])
+ assert not result.data[0].is_valid
+ assert "Exception during SQL parsing" in result.data[0].error_messages
diff --git a/scripts/benchmarks/benchmark_engine_v2.py b/scripts/benchmarks/benchmark_engine_v2.py
new file mode 100644
index 00000000..f6249021
--- /dev/null
+++ b/scripts/benchmarks/benchmark_engine_v2.py
@@ -0,0 +1,822 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Benchmark DataDesigner engine performance with mock LLMs."""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import hashlib
+import json
+import math
+import os
+import random
+import statistics
+import subprocess
+import sys
+import tempfile
+import time
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from data_designer.config.column_configs import LLMTextColumnConfig, SamplerColumnConfig, ValidationColumnConfig
+from data_designer.config.config_builder import DataDesignerConfigBuilder
+from data_designer.config.mcp import MCPProvider, ToolConfig
+from data_designer.config.models import ChatCompletionInferenceParams, ModelConfig, ModelProvider
+from data_designer.config.run_config import RunConfig
+from data_designer.config.sampler_params import SamplerType, UniformSamplerParams
+from data_designer.config.validator_params import LocalCallableValidatorParams, ValidatorType
+from data_designer.engine.mcp.registry import MCPToolDefinition, MCPToolResult
+from data_designer.lazy_heavy_imports import np, pd
+
+if TYPE_CHECKING:
+ import numpy as np
+ import pandas as pd
+
+
+RESULT_PREFIX = "BENCHMARK_RESULT="
+DEFAULT_NUM_RECORDS = 1024
+DEFAULT_BUFFER_SIZE = 1024
+DEFAULT_SEED = 11
+DEFAULT_MAX_PARALLEL_REQUESTS = 16
+DEFAULT_VALIDATOR_BATCH_SIZE = 256
+DEFAULT_ITERATIONS = 5
+
+MOCK_MCP_PROVIDER_NAME = "mock-mcp"
+MOCK_TOOL_ALIAS = "mock-tools"
+MOCK_TOOL_NAME = "mock_lookup"
+MOCK_TOOL_DESCRIPTION = "Mock lookup tool for benchmark runs."
+MOCK_TOOL_SCHEMA = {
+ "type": "object",
+ "properties": {
+ "query": {"type": "string"},
+ "limit": {"type": "integer"},
+ },
+ "required": ["query"],
+}
+
+
+@dataclass(frozen=True)
+class BenchmarkSettings:
+ num_records: int
+ buffer_size: int
+ seed: int
+ max_parallel_requests: int
+ validator_batch_size: int
+
+ def to_cli_args(self) -> list[str]:
+ return [
+ "--num-records",
+ str(self.num_records),
+ "--buffer-size",
+ str(self.buffer_size),
+ "--seed",
+ str(self.seed),
+ "--max-parallel-requests",
+ str(self.max_parallel_requests),
+ "--validator-batch-size",
+ str(self.validator_batch_size),
+ ]
+
+
+@dataclass(frozen=True)
+class BenchmarkResult:
+ engine_mode: str
+ num_records: int
+ buffer_size: int
+ build_time_sec: float
+ total_time_sec: float
+ dataset_hash: str
+ row_count: int
+ column_count: int
+
+ def to_dict(self) -> dict[str, Any]:
+ return {
+ "engine_mode": self.engine_mode,
+ "num_records": self.num_records,
+ "buffer_size": self.buffer_size,
+ "build_time_sec": self.build_time_sec,
+ "total_time_sec": self.total_time_sec,
+ "dataset_hash": self.dataset_hash,
+ "row_count": self.row_count,
+ "column_count": self.column_count,
+ }
+
+ @classmethod
+ def from_dict(cls, payload: dict[str, Any]) -> BenchmarkResult:
+ return cls(
+ engine_mode=str(payload["engine_mode"]),
+ num_records=int(payload["num_records"]),
+ buffer_size=int(payload["buffer_size"]),
+ build_time_sec=float(payload["build_time_sec"]),
+ total_time_sec=float(payload["total_time_sec"]),
+ dataset_hash=str(payload["dataset_hash"]),
+ row_count=int(payload["row_count"]),
+ column_count=int(payload["column_count"]),
+ )
+
+
+@dataclass(frozen=True)
+class MetricStats:
+ mean: float
+ stdev: float
+ ci_half_width: float
+ n: int
+
+ @property
+ def ci_low(self) -> float:
+ return self.mean - self.ci_half_width
+
+ @property
+ def ci_high(self) -> float:
+ return self.mean + self.ci_half_width
+
+
+@dataclass(frozen=True)
+class ResponseProfile:
+ label: str
+ score_mu: float
+ score_sigma: float
+ latency_alpha: float
+ latency_beta: float
+ volatility_sigma: float
+ categories: tuple[str, ...]
+ category_weights: tuple[float, ...]
+
+
+MODEL_PROFILES: dict[str, ResponseProfile] = {
+ "mock-alpha": ResponseProfile(
+ label="alpha",
+ score_mu=0.1,
+ score_sigma=0.35,
+ latency_alpha=2.2,
+ latency_beta=6.0,
+ volatility_sigma=0.25,
+ categories=("low", "mid", "high"),
+ category_weights=(0.25, 0.55, 0.2),
+ ),
+ "mock-beta": ResponseProfile(
+ label="beta",
+ score_mu=0.3,
+ score_sigma=0.45,
+ latency_alpha=2.6,
+ latency_beta=4.8,
+ volatility_sigma=0.3,
+ categories=("low", "mid", "high"),
+ category_weights=(0.2, 0.5, 0.3),
+ ),
+ "mock-gamma": ResponseProfile(
+ label="gamma",
+ score_mu=0.5,
+ score_sigma=0.5,
+ latency_alpha=3.0,
+ latency_beta=3.6,
+ volatility_sigma=0.35,
+ categories=("low", "mid", "high"),
+ category_weights=(0.15, 0.45, 0.4),
+ ),
+}
+
+DEFAULT_PROFILE = ResponseProfile(
+ label="default",
+ score_mu=0.2,
+ score_sigma=0.4,
+ latency_alpha=2.4,
+ latency_beta=5.0,
+ volatility_sigma=0.3,
+ categories=("low", "mid", "high"),
+ category_weights=(0.3, 0.5, 0.2),
+)
+
+
+@dataclass(frozen=True)
+class FakeMessage:
+ content: str
+ tool_calls: list[dict[str, Any]] | None = None
+ reasoning_content: str | None = None
+
+
+@dataclass(frozen=True)
+class FakeChoice:
+ message: FakeMessage
+
+
+@dataclass(frozen=True)
+class FakeResponse:
+ choices: list[FakeChoice]
+ usage: Any | None = None
+ model: str | None = None
+
+
+def _distinct_parallel_requests(base: int) -> tuple[int, int, int]:
+ if base < 3:
+ raise ValueError("max_parallel_requests must be >= 3 to create distinct per-model limits.")
+ high = base
+ mid = max(1, int(round(high / 2)))
+ low = max(1, int(round(high / 5)))
+
+ if mid >= high:
+ mid = high - 1
+ if low >= mid:
+ low = max(1, mid - 1)
+
+ return high, mid, low
+
+
+def _t_critical_95(df: int) -> float:
+ table = {
+ 1: 12.706,
+ 2: 4.303,
+ 3: 3.182,
+ 4: 2.776,
+ 5: 2.571,
+ 6: 2.447,
+ 7: 2.365,
+ 8: 2.306,
+ 9: 2.262,
+ 10: 2.228,
+ 11: 2.201,
+ 12: 2.179,
+ 13: 2.160,
+ 14: 2.145,
+ 15: 2.131,
+ 16: 2.120,
+ 17: 2.110,
+ 18: 2.101,
+ 19: 2.093,
+ 20: 2.086,
+ 21: 2.080,
+ 22: 2.074,
+ 23: 2.069,
+ 24: 2.064,
+ 25: 2.060,
+ 26: 2.056,
+ 27: 2.052,
+ 28: 2.048,
+ 29: 2.045,
+ 30: 2.042,
+ }
+ return table.get(df, 1.96)
+
+
+def _compute_stats(values: list[float]) -> MetricStats:
+ if not values:
+ return MetricStats(mean=0.0, stdev=0.0, ci_half_width=0.0, n=0)
+ if len(values) == 1:
+ return MetricStats(mean=values[0], stdev=0.0, ci_half_width=0.0, n=1)
+ stdev = statistics.stdev(values)
+ mean = statistics.mean(values)
+ t_value = _t_critical_95(len(values) - 1)
+ ci_half_width = t_value * stdev / math.sqrt(len(values))
+ return MetricStats(mean=mean, stdev=stdev, ci_half_width=ci_half_width, n=len(values))
+
+
+def _format_stats(stats: MetricStats, *, unit: str, precision: int = 3) -> str:
+ fmt = f"{{:.{precision}f}}"
+ mean = fmt.format(stats.mean)
+ ci = fmt.format(stats.ci_half_width)
+ stdev = fmt.format(stats.stdev)
+ return f"{mean}{unit} ± {ci}{unit} (stdev {stdev}{unit}, n={stats.n})"
+
+
+def _format_speed_stats(stats: MetricStats, *, precision: int = 2) -> str:
+ fmt = f"{{:.{precision}f}}"
+ mean = fmt.format(stats.mean)
+ ci = fmt.format(stats.ci_half_width)
+ stdev = fmt.format(stats.stdev)
+ return f"{mean}x ± {ci}x (stdev {stdev}x, n={stats.n})"
+
+
+def _significant_diff(stats: MetricStats) -> bool:
+ return stats.n > 1 and abs(stats.mean) > stats.ci_half_width
+
+
+def _json_default(value: Any) -> Any:
+ if isinstance(value, np.generic):
+ return value.item()
+ if isinstance(value, np.ndarray):
+ return value.tolist()
+ if isinstance(value, (pd.Timestamp, pd.Timedelta)):
+ return value.isoformat()
+ if isinstance(value, set):
+ return sorted(value)
+ if isinstance(value, bytes):
+ return value.decode("utf-8", errors="replace")
+ return str(value)
+
+
+def _stable_seed(model: str, messages: list[dict[str, Any]]) -> int:
+ payload = json.dumps(
+ {"model": model, "messages": messages},
+ sort_keys=True,
+ separators=(",", ":"),
+ ensure_ascii=True,
+ default=_json_default,
+ )
+ digest = hashlib.sha256(payload.encode("utf-8")).digest()
+ return int.from_bytes(digest[:8], "big")
+
+
+def _profile_for_model(model: str) -> ResponseProfile:
+ for key, profile in MODEL_PROFILES.items():
+ if key in model:
+ return profile
+ return DEFAULT_PROFILE
+
+
+def _mock_response_text(model: str, messages: list[dict[str, Any]]) -> str:
+ profile = _profile_for_model(model)
+ rng = random.Random(_stable_seed(model, messages))
+ category = rng.choices(profile.categories, weights=profile.category_weights, k=1)[0]
+ score = rng.lognormvariate(profile.score_mu, profile.score_sigma)
+ latency_ms = int(rng.betavariate(profile.latency_alpha, profile.latency_beta) * 900.0)
+ volatility = rng.gauss(0.0, profile.volatility_sigma)
+ return f"{profile.label}:{category}|score={score:.3f}|latency_ms={latency_ms}|vol={volatility:.3f}"
+
+
+def _tool_call_id(model: str, messages: list[dict[str, Any]]) -> str:
+ call_seed = _stable_seed(model, messages)
+ return f"tool-{call_seed:016x}"
+
+
+def _tool_call_arguments(model: str, messages: list[dict[str, Any]]) -> dict[str, Any]:
+ rng = random.Random(_stable_seed(model, messages))
+ return {
+ "query": f"{model}-lookup-{rng.randint(1000, 9999)}",
+ "limit": rng.randint(1, 3),
+ }
+
+
+def _build_tool_call(model: str, messages: list[dict[str, Any]]) -> dict[str, Any]:
+ arguments = _tool_call_arguments(model, messages)
+ return {
+ "id": _tool_call_id(model, messages),
+ "type": "function",
+ "function": {"name": MOCK_TOOL_NAME, "arguments": json.dumps(arguments)},
+ }
+
+
+def _should_request_tool(messages: list[dict[str, Any]]) -> bool:
+ return not any(message.get("role") == "tool" for message in messages)
+
+
+def _mock_tool_definition() -> MCPToolDefinition:
+ return MCPToolDefinition(
+ name=MOCK_TOOL_NAME,
+ description=MOCK_TOOL_DESCRIPTION,
+ input_schema=MOCK_TOOL_SCHEMA,
+ )
+
+
+def _mock_tool_result(tool_name: str, arguments: dict[str, Any], provider_name: str) -> MCPToolResult:
+ payload = {
+ "tool": tool_name,
+ "provider": provider_name,
+ "query": arguments.get("query", ""),
+ "limit": arguments.get("limit", 0),
+ "status": "ok",
+ }
+ return MCPToolResult(content=json.dumps(payload))
+
+
+def _fake_response(model: str, messages: list[dict[str, Any]], **kwargs: Any) -> FakeResponse:
+ if kwargs.get("tools") and _should_request_tool(messages):
+ tool_call = _build_tool_call(model, messages)
+ return FakeResponse(
+ choices=[FakeChoice(message=FakeMessage(content="Using tool.", tool_calls=[tool_call]))],
+ model=model,
+ )
+ response_text = _mock_response_text(model, messages)
+ return FakeResponse(choices=[FakeChoice(message=FakeMessage(content=response_text))], model=model)
+
+
+@contextlib.contextmanager
+def _patch_llm_responses() -> Iterator[None]:
+ # Imports are deferred so engine selection respects DATA_DESIGNER_ASYNC_ENGINE.
+ from data_designer.engine.models.litellm_overrides import CustomRouter
+
+ original_completion = CustomRouter.completion
+ original_acompletion = getattr(CustomRouter, "acompletion", None)
+
+ def fake_completion(self: Any, model: str, messages: list[dict[str, Any]], **kwargs: Any) -> FakeResponse:
+ _ = self
+ return _fake_response(model, messages, **kwargs)
+
+ async def fake_acompletion(self: Any, model: str, messages: list[dict[str, Any]], **kwargs: Any) -> FakeResponse:
+ _ = self
+ return _fake_response(model, messages, **kwargs)
+
+ CustomRouter.completion = fake_completion
+ CustomRouter.acompletion = fake_acompletion
+ try:
+ yield
+ finally:
+ CustomRouter.completion = original_completion
+ if original_acompletion is not None:
+ CustomRouter.acompletion = original_acompletion
+ else:
+ try:
+ delattr(CustomRouter, "acompletion")
+ except AttributeError:
+ pass
+
+
+@contextlib.contextmanager
+def _patch_mcp_io() -> Iterator[None]:
+ import data_designer.engine.mcp.io as mcp_io
+
+ original_list_tools = mcp_io.list_tools
+ original_call_tools = mcp_io.call_tools
+
+ def fake_list_tools(provider: Any, timeout_sec: float | None = None) -> tuple[MCPToolDefinition, ...]:
+ if getattr(provider, "name", None) != MOCK_MCP_PROVIDER_NAME:
+ return original_list_tools(provider, timeout_sec=timeout_sec)
+ return (_mock_tool_definition(),)
+
+ def fake_call_tools(
+ calls: list[tuple[Any, str, dict[str, Any]]],
+ *,
+ timeout_sec: float | None = None,
+ ) -> list[MCPToolResult]:
+ if any(getattr(call[0], "name", None) != MOCK_MCP_PROVIDER_NAME for call in calls):
+ return original_call_tools(calls, timeout_sec=timeout_sec)
+ return [_mock_tool_result(tool_name, arguments, provider.name) for provider, tool_name, arguments in calls]
+
+ mcp_io.list_tools = fake_list_tools
+ mcp_io.call_tools = fake_call_tools
+ try:
+ yield
+ finally:
+ mcp_io.list_tools = original_list_tools
+ mcp_io.call_tools = original_call_tools
+
+
+def _extract_metric(text: str, key: str) -> float | None:
+ marker = f"{key}="
+ start = text.find(marker)
+ if start == -1:
+ return None
+ start += len(marker)
+ end = start
+ while end < len(text) and (text[end].isdigit() or text[end] in {".", "-"}):
+ end += 1
+ try:
+ return float(text[start:end])
+ except ValueError:
+ return None
+
+
+def _validate_recommendation(df: pd.DataFrame) -> pd.DataFrame:
+ series = df["llm_stage3"].astype(str)
+ scores = series.map(lambda text: _extract_metric(text, "score"))
+ latencies = series.map(lambda text: _extract_metric(text, "latency_ms"))
+ scores_numeric = pd.to_numeric(scores, errors="coerce")
+ latency_numeric = pd.to_numeric(latencies, errors="coerce")
+ is_valid = scores_numeric.between(0.0, 10.0) & latency_numeric.between(0.0, 900.0)
+ return pd.DataFrame(
+ {
+ "is_valid": is_valid.fillna(False).astype(bool),
+ "score": scores_numeric,
+ "latency_ms": latency_numeric,
+ }
+ )
+
+
+def _build_config(settings: BenchmarkSettings) -> DataDesignerConfigBuilder:
+ high_parallel, mid_parallel, low_parallel = _distinct_parallel_requests(settings.max_parallel_requests)
+ model_configs = [
+ ModelConfig(
+ alias="mock-alpha",
+ model="mock-alpha",
+ provider="mock-provider",
+ inference_parameters=ChatCompletionInferenceParams(max_parallel_requests=high_parallel),
+ skip_health_check=True,
+ ),
+ ModelConfig(
+ alias="mock-beta",
+ model="mock-beta",
+ provider="mock-provider",
+ inference_parameters=ChatCompletionInferenceParams(max_parallel_requests=low_parallel),
+ skip_health_check=True,
+ ),
+ ModelConfig(
+ alias="mock-gamma",
+ model="mock-gamma",
+ provider="mock-provider",
+ inference_parameters=ChatCompletionInferenceParams(max_parallel_requests=mid_parallel),
+ skip_health_check=True,
+ ),
+ ]
+
+ builder = DataDesignerConfigBuilder(model_configs=model_configs)
+ builder.add_tool_config(
+ ToolConfig(
+ tool_alias=MOCK_TOOL_ALIAS,
+ providers=[MOCK_MCP_PROVIDER_NAME],
+ allow_tools=[MOCK_TOOL_NAME],
+ max_tool_call_turns=1,
+ timeout_sec=1.0,
+ )
+ )
+ builder.add_column(
+ SamplerColumnConfig(
+ name="seed_value",
+ sampler_type=SamplerType.UNIFORM,
+ params=UniformSamplerParams(low=0.0, high=100.0, decimal_places=3),
+ )
+ )
+ builder.add_column(
+ LLMTextColumnConfig(
+ name="llm_stage1",
+ model_alias="mock-alpha",
+ prompt="Summarize the signal for seed {{ seed_value }}.",
+ )
+ )
+ builder.add_column(
+ LLMTextColumnConfig(
+ name="llm_stage2",
+ model_alias="mock-beta",
+ tool_alias=MOCK_TOOL_ALIAS,
+ prompt="Analyze {{ llm_stage1 }} and produce a risk assessment.",
+ )
+ )
+ builder.add_column(
+ LLMTextColumnConfig(
+ name="llm_stage3",
+ model_alias="mock-gamma",
+ prompt="Generate a recommendation from {{ llm_stage2 }} with seed {{ seed_value }}.",
+ )
+ )
+ builder.add_column(
+ ValidationColumnConfig(
+ name="llm_stage3_validation",
+ target_columns=["llm_stage3"],
+ validator_type=ValidatorType.LOCAL_CALLABLE,
+ validator_params=LocalCallableValidatorParams(validation_function=_validate_recommendation),
+ batch_size=settings.validator_batch_size,
+ )
+ )
+ return builder
+
+
+def _dataset_fingerprint(df: pd.DataFrame) -> str:
+ normalized = df.reset_index(drop=True)
+ normalized = normalized.reindex(sorted(normalized.columns), axis=1)
+ records = normalized.to_dict(orient="records")
+ payload = json.dumps(
+ records,
+ sort_keys=True,
+ separators=(",", ":"),
+ ensure_ascii=True,
+ default=_json_default,
+ )
+ return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+
+
+def _run_single_benchmark(settings: BenchmarkSettings, engine_mode: str) -> BenchmarkResult:
+ # Imports are deferred so engine selection respects DATA_DESIGNER_ASYNC_ENGINE.
+ from data_designer.engine.dataset_builders.artifact_storage import ArtifactStorage
+ from data_designer.engine.dataset_builders.column_wise_builder import ColumnWiseDatasetBuilder
+ from data_designer.engine.model_provider import resolve_model_provider_registry
+ from data_designer.engine.resources.resource_provider import create_resource_provider
+ from data_designer.engine.resources.seed_reader import SeedReaderRegistry
+ from data_designer.engine.secret_resolver import CompositeResolver, EnvironmentResolver, PlaintextResolver
+
+ random.seed(settings.seed)
+ np.random.seed(settings.seed)
+
+ run_config = RunConfig(
+ buffer_size=settings.buffer_size,
+ disable_early_shutdown=True,
+ max_conversation_restarts=0,
+ max_conversation_correction_steps=0,
+ )
+ builder = _build_config(settings)
+
+ provider = ModelProvider(
+ name="mock-provider",
+ endpoint="https://mock.local",
+ provider_type="openai",
+ api_key="mock-key",
+ )
+ mcp_provider = MCPProvider(
+ name=MOCK_MCP_PROVIDER_NAME,
+ endpoint="https://mock.local/mcp",
+ api_key="mock-mcp-key",
+ )
+ model_provider_registry = resolve_model_provider_registry([provider], default_provider_name=provider.name)
+ secret_resolver = CompositeResolver([EnvironmentResolver(), PlaintextResolver()])
+
+ with tempfile.TemporaryDirectory() as temp_dir:
+ artifact_storage = ArtifactStorage(artifact_path=temp_dir, dataset_name=f"benchmark-{engine_mode}")
+ resource_provider = create_resource_provider(
+ artifact_storage=artifact_storage,
+ model_configs=builder.model_configs,
+ secret_resolver=secret_resolver,
+ model_provider_registry=model_provider_registry,
+ seed_reader_registry=SeedReaderRegistry(readers=[]),
+ blob_storage=None,
+ seed_dataset_source=None,
+ run_config=run_config,
+ mcp_providers=[mcp_provider],
+ tool_configs=builder.tool_configs,
+ )
+ dataset_builder = ColumnWiseDatasetBuilder(
+ data_designer_config=builder.build(),
+ resource_provider=resource_provider,
+ )
+
+ total_start = time.perf_counter()
+ with _patch_llm_responses(), _patch_mcp_io():
+ build_start = time.perf_counter()
+ dataset_builder.build(num_records=settings.num_records)
+ build_time = time.perf_counter() - build_start
+ dataset = dataset_builder.artifact_storage.load_dataset_with_dropped_columns()
+
+ dataset_hash = _dataset_fingerprint(dataset)
+ total_time = time.perf_counter() - total_start
+
+ return BenchmarkResult(
+ engine_mode=engine_mode,
+ num_records=settings.num_records,
+ buffer_size=settings.buffer_size,
+ build_time_sec=build_time,
+ total_time_sec=total_time,
+ dataset_hash=dataset_hash,
+ row_count=int(dataset.shape[0]),
+ column_count=int(dataset.shape[1]),
+ )
+
+
+def _run_subprocess(settings: BenchmarkSettings, engine_mode: str) -> BenchmarkResult:
+ env = os.environ.copy()
+ if engine_mode == "async":
+ env["DATA_DESIGNER_ASYNC_ENGINE"] = "1"
+ else:
+ env.pop("DATA_DESIGNER_ASYNC_ENGINE", None)
+
+ script_path = Path(__file__).resolve()
+ cmd = [sys.executable, str(script_path), "--mode", "run", "--engine", engine_mode, *settings.to_cli_args()]
+ completed = subprocess.run(cmd, capture_output=True, text=True, env=env, check=False)
+
+ if completed.returncode != 0:
+ raise RuntimeError(f"Benchmark subprocess failed.\nstdout:\n{completed.stdout}\nstderr:\n{completed.stderr}")
+
+ for line in reversed(completed.stdout.splitlines()):
+ if line.startswith(RESULT_PREFIX):
+ payload = json.loads(line.removeprefix(RESULT_PREFIX))
+ return BenchmarkResult.from_dict(payload)
+
+ raise RuntimeError(
+ f"Benchmark subprocess did not emit a result payload.\nstdout:\n{completed.stdout}\nstderr:\n{completed.stderr}"
+ )
+
+
+def _format_speedup(sync_time: float, async_time: float) -> str:
+ if async_time <= 0:
+ return "n/a"
+ return f"{(sync_time / async_time):.2f}x"
+
+
+def _run_with_progress(settings: BenchmarkSettings, engine_mode: str, iteration: int, total: int) -> BenchmarkResult:
+ print(f"[{iteration}/{total}] Running {engine_mode} benchmark...", end="", flush=True)
+ result = _run_subprocess(settings, engine_mode)
+ print(f" done ({result.total_time_sec:.3f}s)")
+ return result
+
+
+def _compare_runs(settings: BenchmarkSettings, iterations: int) -> int:
+ sync_results: list[BenchmarkResult] = []
+ async_results: list[BenchmarkResult] = []
+ expected_hash: str | None = None
+
+ for iteration in range(1, iterations + 1):
+ sync_result = _run_with_progress(settings, "sync", iteration, iterations)
+ async_result = _run_with_progress(settings, "async", iteration, iterations)
+
+ if sync_result.dataset_hash != async_result.dataset_hash:
+ print(
+ "Content mismatch detected: "
+ f"sync hash {sync_result.dataset_hash} vs async hash {async_result.dataset_hash}"
+ )
+ return 1
+
+ if expected_hash is None:
+ expected_hash = sync_result.dataset_hash
+ elif expected_hash != sync_result.dataset_hash or expected_hash != async_result.dataset_hash:
+ print("Content mismatch detected across iterations.")
+ return 1
+
+ sync_results.append(sync_result)
+ async_results.append(async_result)
+
+ build_sync = [result.build_time_sec for result in sync_results]
+ build_async = [result.build_time_sec for result in async_results]
+ total_sync = [result.total_time_sec for result in sync_results]
+ total_async = [result.total_time_sec for result in async_results]
+
+ build_speedups = [sync / async_ for sync, async_ in zip(build_sync, build_async)]
+ total_speedups = [sync / async_ for sync, async_ in zip(total_sync, total_async)]
+ build_diffs = [sync - async_ for sync, async_ in zip(build_sync, build_async)]
+ total_diffs = [sync - async_ for sync, async_ in zip(total_sync, total_async)]
+
+ build_sync_stats = _compute_stats(build_sync)
+ build_async_stats = _compute_stats(build_async)
+ total_sync_stats = _compute_stats(total_sync)
+ total_async_stats = _compute_stats(total_async)
+
+ build_speed_stats = _compute_stats(build_speedups)
+ total_speed_stats = _compute_stats(total_speedups)
+ build_diff_stats = _compute_stats(build_diffs)
+ total_diff_stats = _compute_stats(total_diffs)
+
+ print("\nEngine benchmark summary (95% CI)")
+ print(f"- runs: {iterations} | content match: yes | hash {expected_hash}")
+ print(f"- build time sync: {_format_stats(build_sync_stats, unit='s')}")
+ print(f"- build time async: {_format_stats(build_async_stats, unit='s')}")
+ print(
+ f"- build speedup: {_format_speed_stats(build_speed_stats)} | "
+ f"paired diff {_format_stats(build_diff_stats, unit='s')} | "
+ f"significant: {'yes' if _significant_diff(build_diff_stats) else 'no'}"
+ )
+ print(f"- total time sync: {_format_stats(total_sync_stats, unit='s')}")
+ print(f"- total time async: {_format_stats(total_async_stats, unit='s')}")
+ print(
+ f"- total speedup: {_format_speed_stats(total_speed_stats)} | "
+ f"paired diff {_format_stats(total_diff_stats, unit='s')} | "
+ f"significant: {'yes' if _significant_diff(total_diff_stats) else 'no'}"
+ )
+
+ return 0
+
+
+def _parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(
+ description="Benchmark DataDesigner engine with mock LLMs and compare async execution."
+ )
+ parser.add_argument(
+ "--mode",
+ type=str,
+ choices=("compare", "run"),
+ default="compare",
+ help="Run both engines in subprocesses, or run once in the current process.",
+ )
+ parser.add_argument(
+ "--engine",
+ type=str,
+ choices=("sync", "async"),
+ default="sync",
+ help="Engine mode for --mode run.",
+ )
+ parser.add_argument("--num-records", type=int, default=DEFAULT_NUM_RECORDS, help="Records to generate.")
+ parser.add_argument("--buffer-size", type=int, default=DEFAULT_BUFFER_SIZE, help="Batch buffer size.")
+ parser.add_argument("--seed", type=int, default=DEFAULT_SEED, help="Random seed for determinism.")
+ parser.add_argument(
+ "--iterations",
+ type=int,
+ default=DEFAULT_ITERATIONS,
+ help="Number of sync/async runs to include in the compare mode.",
+ )
+ parser.add_argument(
+ "--max-parallel-requests",
+ type=int,
+ default=DEFAULT_MAX_PARALLEL_REQUESTS,
+ help="Max parallel LLM requests per model.",
+ )
+ parser.add_argument(
+ "--validator-batch-size",
+ type=int,
+ default=DEFAULT_VALIDATOR_BATCH_SIZE,
+ help="Batch size for the local validator.",
+ )
+ return parser.parse_args()
+
+
+def main() -> None:
+ args = _parse_args()
+ settings = BenchmarkSettings(
+ num_records=args.num_records,
+ buffer_size=args.buffer_size,
+ seed=args.seed,
+ max_parallel_requests=args.max_parallel_requests,
+ validator_batch_size=args.validator_batch_size,
+ )
+
+ if args.mode == "compare":
+ sys.exit(_compare_runs(settings, args.iterations))
+
+ if args.engine == "async":
+ os.environ["DATA_DESIGNER_ASYNC_ENGINE"] = "1"
+ else:
+ os.environ.pop("DATA_DESIGNER_ASYNC_ENGINE", None)
+
+ print(f"Running {args.engine} benchmark...")
+ result = _run_single_benchmark(settings, args.engine)
+ print(f"{RESULT_PREFIX}{json.dumps(result.to_dict(), sort_keys=True)}")
+
+
+if __name__ == "__main__":
+ main()