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

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/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/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/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/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"),

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.dist-info → data_designer_engine-0.4.0rc2.dist-info}/METADATA

@@ -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.dist-info → data_designer_engine-0.4.0rc2.dist-info}/RECORD

@@ -1,5 +1,5 @@
 data_designer/engine/__init__.py,sha256=ObZ6NUPeEvvpGTJ5WIGKUyIrIjaI747OM6ErweRtHxQ,137
-data_designer/engine/_version.py,sha256=yib4WPM_pEWXdpIHBdFnf29aurTH5f4xrnwVlv7cijo,714
+data_designer/engine/_version.py,sha256=FvItxCBzPigrdVpFPfL1gQeV1-km5r7nCNGUzrYebTU,714
 data_designer/engine/compiler.py,sha256=4QAeCJjINtH0afSXygdhiKMyq2KIfaDthK3ApZLgrQ0,4152
 data_designer/engine/configurable_task.py,sha256=6R4FPXPzIeK0lqNVSEXzRDtK14B3dFz38lplr-nkvRE,2539
 data_designer/engine/errors.py,sha256=YXI7ny83BQ16sOK43CpTm384hJTKuZkPTEAjlHlDIfA,1303
@@ -20,7 +20,7 @@ data_designer/engine/column_generators/generators/__init__.py,sha256=ObZ6NUPeEvv
 data_designer/engine/column_generators/generators/base.py,sha256=QElk5KsaUQ3EYwlv40NcZgQsw3HIkX3YQV_0S3erl7Q,4209
 data_designer/engine/column_generators/generators/embedding.py,sha256=uB0jgHlCgctgIUf9ZfMqG1YThbJ0g-GCX3VdNbdDSko,1407
 data_designer/engine/column_generators/generators/expression.py,sha256=BiQcfVTinvQl3OI9nkdhB9B7FGBueWiHJwxTA8uNVuY,2330
-data_designer/engine/column_generators/generators/llm_completion.py,sha256=3S3ikNLLLGnutUdcuswL5dUfcLgT_-he8DiRZ9K706U,4721
+data_designer/engine/column_generators/generators/llm_completion.py,sha256=udYWE3lwaQhZqxRTHQc6w1kWGEvLAfIh2OUjX6vxMB0,4750
 data_designer/engine/column_generators/generators/samplers.py,sha256=gNzURmu9K8Zb5MHamKvZPIxmWlFgl2W4FIVgaFcy4f0,3371
 data_designer/engine/column_generators/generators/seed_dataset.py,sha256=CoQPbz4Ww7pBLaGw8-CYqIk1sjfkBaoRMKZQexdfgKY,6824
 data_designer/engine/column_generators/generators/validation.py,sha256=YfYbk-8_ZUye0No6_Q7hIqpZv_tunnEZ6HkLSMFXlDE,6659
@@ -43,10 +43,10 @@ data_designer/engine/models/errors.py,sha256=k9oZnmk8DRD8U2SVKJJRLwrcdsCcVoJiOb_
 data_designer/engine/models/facade.py,sha256=UBMpw_o2JcsWpJsPdpTPKfFZCh_i0eeG_oaWi1XeKds,12582
 data_designer/engine/models/factory.py,sha256=2NjI0iiGv8ayQ1c249lsJtha4pDmvmtSjdwvlvitRds,1581
 data_designer/engine/models/litellm_overrides.py,sha256=e9IZCFQ6BhNWlOTncm8ErL8w4rtE1_4USh2mtUYxCZI,6207
-data_designer/engine/models/registry.py,sha256=7hZ6TQwwZf259yRZmc3ZI20a4wAo3PCOozPi9Mc5KLo,6827
-data_designer/engine/models/telemetry.py,sha256=wmuekvPRZjNz7p7ImKx5H_hqDRhTv_dSB-u2S6Ze3uo,12502
+data_designer/engine/models/registry.py,sha256=Bid7Mv_ebzbTrlfzN-1wbcFxp_qQwilL0h2iwN5UPJ0,7099
+data_designer/engine/models/telemetry.py,sha256=_VZR6Iatr6-5Hypw3bes5Jr4y7Y3VagxFEVAv36eHcE,12733
 data_designer/engine/models/usage.py,sha256=A0LV9Ycuj_7snOsaqnirs4mlkAjozv2mzj2om2FpDoU,2410
-data_designer/engine/models/utils.py,sha256=HS5pXAAz7IcOcijeClC-xxq6R6DUmC2ykZu8Vr33Ivk,1259
+data_designer/engine/models/utils.py,sha256=sLBs-STJSe7BGzDAngRGGxo6GwAvFmtimqUs54zZ6DU,1259
 data_designer/engine/models/parsers/__init__.py,sha256=ObZ6NUPeEvvpGTJ5WIGKUyIrIjaI747OM6ErweRtHxQ,137
 data_designer/engine/models/parsers/errors.py,sha256=ODcZ4TOsmZyH4-MoNkKXhjiMm_4gLWPsz90qKtNF9_Q,1053
 data_designer/engine/models/parsers/parser.py,sha256=XkdDt2WEnolvsv2bArq4hhujfJ3kLmG6G2jkRXMYA8c,9489
@@ -69,7 +69,7 @@ data_designer/engine/processing/gsonschema/validators.py,sha256=ui3PzGjIclI6Hlw4
 data_designer/engine/processing/processors/base.py,sha256=bkAQO0yK6ATJ3zTwS7F9FXobenJqydCyfijSP2MM-70,472
 data_designer/engine/processing/processors/drop_columns.py,sha256=xT7ym2pQc-R0-YHIuYDQGFn2uAf74309-pV4H878Wlk,1866
 data_designer/engine/processing/processors/registry.py,sha256=ewuFY8QeXpql5CNTZZa_87aYPGPNv1H0hpJR7CBVuzI,1097
-data_designer/engine/processing/processors/schema_transform.py,sha256=RhLXXKoj9MFpOqsXZ2hfSaTr7_yUUNI3gmFBS4XtEy4,2006
+data_designer/engine/processing/processors/schema_transform.py,sha256=cpN5XAg_YNKpne_Ed3Vhk8_yuoTUiUy_pINgPZF2ASk,2822
 data_designer/engine/registry/base.py,sha256=eACpE7o_c2btiiXrOFJw7o0VvACo7DSqhj8AntkNkCQ,3579
 data_designer/engine/registry/data_designer_registry.py,sha256=mz8ksE49pS1JRVDNubYSxTs0j-8Q6sd08F_dYyTCWSE,1528
 data_designer/engine/registry/errors.py,sha256=k1EaV7egNQwNmRsI8EfymTfeNprcDutPf2M6Vc1nbn8,350
@@ -108,6 +108,6 @@ data_designer/engine/validators/local_callable.py,sha256=JaL-yOXrTFpubiO2QlSt4Qb
 data_designer/engine/validators/python.py,sha256=omXjwMaomQYiyq4g6XqKt2wexVuI_rWue9Dk-CYc-do,8039
 data_designer/engine/validators/remote.py,sha256=rythhIrH2GvqncMQeF3FiJa9Om0KZWeK3cWjW-ZubaM,3077
 data_designer/engine/validators/sql.py,sha256=AMaEdA-gj9j0zwVp809x3ycKltd51wVEhI8mMYGyxd4,2408
-data_designer_engine-0.4.0rc1.dist-info/METADATA,sha256=FybLz1fOjJ2bK0zQ93Ti17o7WZTxDFtrBeGx7Oa6jCo,1876
-data_designer_engine-0.4.0rc1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-data_designer_engine-0.4.0rc1.dist-info/RECORD,,
+data_designer_engine-0.4.0rc2.dist-info/METADATA,sha256=ZChyQl5ksGCWVi_XE6wD-GXG9-wWHko1vBDnd9ecLqw,1876
+data_designer_engine-0.4.0rc2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+data_designer_engine-0.4.0rc2.dist-info/RECORD,,