data-designer-engine 0.4.0rc1__tar.gz → 0.4.0rc2__tar.gz

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-designer-engine
-Version: 0.4.0rc1
+Version: 0.4.0rc2
 Summary: Generation engine for DataDesigner synthetic data generation
 License-Expression: Apache-2.0
 Classifier: Development Status :: 4 - Beta

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc2}/src/data_designer/engine/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.4.0rc1'
-__version_tuple__ = version_tuple = (0, 4, 0, 'rc1')
+__version__ = version = '0.4.0rc2'
+__version_tuple__ = version_tuple = (0, 4, 0, 'rc2')
 
 __commit_id__ = commit_id = None

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc2}/src/data_designer/engine/column_generators/generators/llm_completion.py

@@ -62,9 +62,9 @@ class ColumnGeneratorWithModelChatCompletion(ColumnGeneratorWithModel[TaskConfig
 
         multi_modal_context = None
         if self.config.multi_modal_context is not None and len(self.config.multi_modal_context) > 0:
-            multi_modal_context = [
-                context.get_context(deserialized_record) for context in self.config.multi_modal_context
-            ]
+            multi_modal_context = []
+            for context in self.config.multi_modal_context:
+                multi_modal_context.extend(context.get_contexts(deserialized_record))
 
         response, reasoning_trace = self.model.generate(
             prompt=self.prompt_renderer.render(

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc2}/src/data_designer/engine/models/registry.py

@@ -107,6 +107,11 @@ class ModelRegistry:
     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}..."

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc2}/src/data_designer/engine/models/telemetry.py

@@ -8,6 +8,7 @@ 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
@@ -18,15 +19,12 @@ import platform
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from enum import Enum
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import Any, ClassVar
 
 from pydantic import BaseModel, Field
 
 from data_designer.lazy_heavy_imports import httpx
 
-if TYPE_CHECKING:
-    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"
@@ -35,6 +33,7 @@ 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):
@@ -231,7 +230,11 @@ class TelemetryHandler:
         self._timer_task: asyncio.Task | None = None
         self._running = False
         self._source_client_version = source_client_version
-        self._session_id = session_id
+        # 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:

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc2}/src/data_designer/engine/processing/processors/schema_transform.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 
 import json
 import logging
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from data_designer.config.processors import SchemaTransformProcessorConfig
 from data_designer.engine.dataset_builders.artifact_storage import BatchStage
@@ -20,6 +20,26 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def _json_escape_record(record: dict[str, Any]) -> dict[str, Any]:
+    """Escape record values for safe insertion into a JSON template."""
+
+    def escape_for_json_string(s: str) -> str:
+        """Use json.dumps to escape, then strip the surrounding quotes."""
+        return json.dumps(s)[1:-1]
+
+    escaped = {}
+    for key, value in record.items():
+        if isinstance(value, str):
+            escaped[key] = escape_for_json_string(value)
+        elif isinstance(value, (dict, list)):
+            escaped[key] = escape_for_json_string(json.dumps(value))
+        elif value is None:
+            escaped[key] = "null"
+        else:
+            escaped[key] = str(value)
+    return escaped
+
+
 class SchemaTransformProcessor(WithJinja2UserTemplateRendering, Processor[SchemaTransformProcessorConfig]):
     @property
     def template_as_str(self) -> str:
@@ -27,10 +47,12 @@ class SchemaTransformProcessor(WithJinja2UserTemplateRendering, Processor[Schema
 
     def process(self, data: pd.DataFrame, *, current_batch_number: int | None = None) -> pd.DataFrame:
         self.prepare_jinja2_template_renderer(self.template_as_str, data.columns.to_list())
-        formatted_records = [
-            json.loads(self.render_template(deserialize_json_values(record)).replace("\n", "\\n"))
-            for record in data.to_dict(orient="records")
-        ]
+        formatted_records = []
+        for record in data.to_dict(orient="records"):
+            deserialized = deserialize_json_values(record)
+            escaped = _json_escape_record(deserialized)
+            rendered = self.render_template(escaped)
+            formatted_records.append(json.loads(rendered))
         formatted_data = pd.DataFrame(formatted_records)
         if current_batch_number is not None:
             self.artifact_storage.write_batch_to_parquet_file(

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc2}/tests/engine/models/test_model_registry.py

@@ -313,6 +313,54 @@ 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)
+def test_run_health_check_skip_health_check_flag(
+    mock_completion,
+    stub_secrets_resolver,
+    stub_model_provider_registry,
+):
+    # Create model configs: one with skip_health_check=True, others with default (False)
+    model_configs = [
+        ModelConfig(
+            alias="skip-model",
+            model="skip-model-id",
+            provider="stub-model-provider",
+            inference_parameters=ChatCompletionInferenceParams(),
+            skip_health_check=True,
+        ),
+        ModelConfig(
+            alias="check-model",
+            model="check-model-id",
+            provider="stub-model-provider",
+            inference_parameters=ChatCompletionInferenceParams(),
+            skip_health_check=False,
+        ),
+        ModelConfig(
+            alias="default-model",
+            model="default-model-id",
+            provider="stub-model-provider",
+            inference_parameters=ChatCompletionInferenceParams(),
+        ),
+    ]
+
+    # Create a fresh model registry with the test configs
+    model_registry = create_model_registry(
+        model_configs=model_configs,
+        secret_resolver=stub_secrets_resolver,
+        model_provider_registry=stub_model_provider_registry,
+    )
+
+    model_aliases = ["skip-model", "check-model", "default-model"]
+    model_registry.run_health_check(model_aliases)
+
+    # Only check-model and default-model should be checked (skip-model is skipped)
+    assert mock_completion.call_count == 2  # check-model and default-model
+
+    # Verify the correct models were called
+    called_model_aliases = {call[0][0].model_alias for call in mock_completion.call_args_list}
+    assert called_model_aliases == {"check-model", "default-model"}
+
+
 @pytest.mark.parametrize(
     "alias,expected_result,expected_error",
     [

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc2}/tests/engine/models/test_model_utils.py

@@ -26,11 +26,11 @@ def test_prompt_to_messages():
         {"content": "hello", "role": "user"},
     ]
     assert prompt_to_messages(user_prompt="hello", multi_modal_context=[mult_modal_context]) == [
-        {"content": [{"type": "text", "text": "hello"}, mult_modal_context], "role": "user"}
+        {"content": [mult_modal_context, {"type": "text", "text": "hello"}], "role": "user"}
     ]
     assert prompt_to_messages(
         user_prompt="hello", system_prompt=stub_system_prompt, multi_modal_context=[mult_modal_context]
     ) == [
         {"content": stub_system_prompt, "role": "system"},
-        {"content": [{"type": "text", "text": "hello"}, mult_modal_context], "role": "user"},
+        {"content": [mult_modal_context, {"type": "text", "text": "hello"}], "role": "user"},
     ]

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc2}/tests/engine/processing/processors/test_schema_transform.py

@@ -129,7 +129,65 @@ def test_process_with_json_serialized_values(stub_processor: SchemaTransformProc
     assert written_dataframe is not None
     assert len(written_dataframe) == 2
 
-    # Verify that nested JSON values are properly deserialized in template rendering
+    # Verify that nested JSON values are properly serialized as JSON strings in template rendering
     first_output = written_dataframe.iloc[0].to_dict()
     assert first_output["text"] == "hello"
-    assert first_output["value"] == "{'nested': 'value1'}"
+    # Nested JSON should be properly serialized as JSON string (not Python repr)
+    assert first_output["value"] == '{"nested": "value1"}'
+
+
+def test_process_with_special_characters_in_llm_output(stub_processor: SchemaTransformProcessor) -> None:
+    """Test that LLM outputs with special characters are properly escaped for JSON.
+
+    This addresses GitHub issue #227 where SchemaTransformProcessor fails with JSONDecodeError
+    when LLM-generated content contains quotes, backslashes, or newlines.
+    """
+    df_with_special_chars = pd.DataFrame(
+        {
+            "col1": [
+                'He said "Hello"',
+                "Line1\nLine2",
+                "Path: C:\\Users\\test",
+                "Tab\there",
+            ],
+            "col2": [1, 2, 3, 4],
+        }
+    )
+
+    # Process should not raise JSONDecodeError
+    stub_processor.process(df_with_special_chars, current_batch_number=0)
+    written_dataframe: pd.DataFrame = stub_processor.artifact_storage.write_batch_to_parquet_file.call_args.kwargs[
+        "dataframe"
+    ]
+
+    # Verify all rows were processed successfully
+    assert written_dataframe is not None
+    assert len(written_dataframe) == 4
+
+    # Verify the special characters are preserved in the output
+    outputs = written_dataframe.to_dict(orient="records")
+    assert outputs[0]["text"] == 'He said "Hello"'
+    assert outputs[1]["text"] == "Line1\nLine2"
+    assert outputs[2]["text"] == "Path: C:\\Users\\test"
+    assert outputs[3]["text"] == "Tab\there"
+
+
+def test_process_with_mixed_special_characters(stub_processor: SchemaTransformProcessor) -> None:
+    """Test complex LLM output with multiple types of special characters."""
+    df_complex = pd.DataFrame(
+        {
+            "col1": [
+                'She replied: "I\'m not sure about that\\nLet me think..."',
+            ],
+            "col2": [42],
+        }
+    )
+
+    stub_processor.process(df_complex, current_batch_number=0)
+    written_dataframe: pd.DataFrame = stub_processor.artifact_storage.write_batch_to_parquet_file.call_args.kwargs[
+        "dataframe"
+    ]
+
+    assert len(written_dataframe) == 1
+    output = written_dataframe.iloc[0].to_dict()
+    assert output["text"] == 'She replied: "I\'m not sure about that\\nLet me think..."'

{data_designer_engine-0.4.0rc1 → data_designer_engine-0.4.0rc2}/src/data_designer/engine/models/utils.py

@@ -21,9 +21,9 @@ def prompt_to_messages(
     user_content = user_prompt
     if multi_modal_context and len(multi_modal_context) > 0:
         user_content = []
-        user_content.append({"type": "text", "text": user_prompt})
         for context in multi_modal_context:
             user_content.append(context)
+        user_content.append({"type": "text", "text": user_prompt})
     return (
         [
             str_to_message(content=system_prompt, role="system"),