fix: use run config for profiler token stats (#738)

Author: eric-tramel

packages/data-designer-engine/src/data_designer/engine/analysis/column_profilers/base.py

@@ -12,6 +12,7 @@
 
 from data_designer.config.base import ConfigBase, SingleColumnConfig
 from data_designer.config.column_types import DataDesignerColumnType
+from data_designer.config.run_config import JinjaRenderingEngine
 from data_designer.engine.configurable_task import ConfigurableTask, TaskConfigT
 
 logger = logging.getLogger(__name__)
@@ -20,6 +21,7 @@
 class ColumnConfigWithDataFrame(ConfigBase):
     column_config: SingleColumnConfig
     df: pd.DataFrame
+    jinja_rendering_engine: JinjaRenderingEngine = JinjaRenderingEngine.SECURE
 
     @model_validator(mode="after")
     def validate_column_exists(self) -> Self:

packages/data-designer-engine/src/data_designer/engine/analysis/column_statistics.py

@@ -73,7 +73,11 @@ def __repr__(self) -> str:
 
 class LLMTextColumnStatisticsCalculator(GeneralColumnStatisticsCalculator):
     def calculate_token_stats(self) -> dict[str, Any]:
-        return calculate_token_stats(self.column_config, self.df)
+        return calculate_token_stats(
+            self.column_config,
+            self.df,
+            jinja_rendering_engine=self.column_config_with_df.jinja_rendering_engine,
+        )
 
 
 class LLMCodeColumnStatisticsCalculator(LLMTextColumnStatisticsCalculator): ...

packages/data-designer-engine/src/data_designer/engine/analysis/dataset_profiler.py

@@ -76,7 +76,7 @@ def profile_dataset(
             logger.info(f"{LOG_INDENT}{c.get_column_emoji()} column: '{c.name}'")
             column_statistics.append(
                 get_column_statistics_calculator(c.column_type)(
-                    column_config_with_df=ColumnConfigWithDataFrame(column_config=c, df=dataset)
+                    column_config_with_df=self._create_column_config_with_df(c, dataset)
                 ).calculate()
             )
 
@@ -86,7 +86,7 @@ def profile_dataset(
             applicable_column_types = profiler.get_applicable_column_types()
             for c in self.config.column_configs:
                 if c.column_type in applicable_column_types:
-                    params = ColumnConfigWithDataFrame(column_config=c, df=dataset)
+                    params = self._create_column_config_with_df(c, dataset)
                     column_profiles.append(profiler.profile(params))
             if len(column_profiles) == 0:
                 logger.warning(
@@ -128,6 +128,17 @@ def _create_column_profiler(self, profiler_config: ColumnProfilerConfigT) -> Col
             config=profiler_config, resource_provider=self.resource_provider
         )
 
+    def _create_column_config_with_df(
+        self,
+        column_config: ColumnConfigT,
+        dataset: pd.DataFrame,
+    ) -> ColumnConfigWithDataFrame:
+        return ColumnConfigWithDataFrame(
+            column_config=column_config,
+            df=dataset,
+            jinja_rendering_engine=self.resource_provider.run_config.jinja_rendering_engine,
+        )
+
     def _validate_column_profiler_configs(self) -> None:
         if self.config.column_profiler_configs:
             if self.resource_provider.model_registry is None:

packages/data-designer-engine/src/data_designer/engine/analysis/utils/column_statistics_calculations.py

@@ -17,6 +17,7 @@
 from data_designer.config.column_configs import (
     LLMTextColumnConfig,
 )
+from data_designer.config.run_config import JinjaRenderingEngine
 from data_designer.engine.column_generators.utils.prompt_renderer import (
     PromptType,
     RecordBasedPromptRenderer,
@@ -95,12 +96,18 @@ def calculate_general_column_info(column_name: str, df: pd.DataFrame) -> dict[st
 
 
 def calculate_input_token_stats(
-    column_config: LLMTextColumnConfig, df: pd.DataFrame
+    column_config: LLMTextColumnConfig,
+    df: pd.DataFrame,
+    *,
+    jinja_rendering_engine: JinjaRenderingEngine = JinjaRenderingEngine.SECURE,
 ) -> dict[str, float | MissingValue]:
     try:
         num_tokens = []
         num_samples = min(MAX_PROMPT_SAMPLE_SIZE, len(df))
-        renderer = RecordBasedPromptRenderer(response_recipe=create_response_recipe(column_config))
+        renderer = RecordBasedPromptRenderer(
+            response_recipe=create_response_recipe(column_config),
+            jinja_rendering_engine=jinja_rendering_engine,
+        )
         for record in df.sample(num_samples, random_state=RANDOM_SEED).to_dict(orient="records"):
             system_prompt = renderer.render(
                 prompt_template=column_config.system_prompt, record=record, prompt_type=PromptType.SYSTEM_PROMPT
@@ -143,9 +150,14 @@ def calculate_output_token_stats(
         }
 
 
-def calculate_token_stats(column_config: LLMTextColumnConfig, df: pd.DataFrame) -> dict[str, float | MissingValue]:
+def calculate_token_stats(
+    column_config: LLMTextColumnConfig,
+    df: pd.DataFrame,
+    *,
+    jinja_rendering_engine: JinjaRenderingEngine = JinjaRenderingEngine.SECURE,
+) -> dict[str, float | MissingValue]:
     return {
-        **calculate_input_token_stats(column_config, df),
+        **calculate_input_token_stats(column_config, df, jinja_rendering_engine=jinja_rendering_engine),
         **calculate_output_token_stats(column_config, df),
     }
 

packages/data-designer-engine/tests/engine/analysis/test_dataset_profiler.py

@@ -1,18 +1,26 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
+from pathlib import Path
 from unittest.mock import patch
 
 import pytest
 
+import data_designer.lazy_heavy_imports as lazy
+from data_designer.config.analysis.column_statistics import MissingValue
 from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
-from data_designer.config.column_configs import SamplerColumnConfig
+from data_designer.config.column_configs import LLMTextColumnConfig, SamplerColumnConfig
+from data_designer.config.run_config import JinjaRenderingEngine, RunConfig
 from data_designer.config.sampler_params import CategorySamplerParams, SamplerType
 from data_designer.engine.analysis.column_profilers.judge_score_profiler import JudgeScoreProfilerConfig
 from data_designer.engine.analysis.dataset_profiler import DataDesignerDatasetProfiler, DatasetProfilerConfig
 from data_designer.engine.analysis.errors import DatasetProfilerConfigurationError
 from data_designer.engine.analysis.utils.judge_score_processing import JudgeScoreSample
 from data_designer.engine.dataset_builders.multi_column_configs import SamplerMultiColumnConfig
+from data_designer.engine.resources.resource_provider import ResourceProvider
+from data_designer.engine.storage.artifact_storage import ArtifactStorage
 
 
 def test_dataset_profiler_config_flattens_multi_column_configs():
@@ -88,6 +96,62 @@ def test_dataset_profiler_profile_dataset_with_column_profilers(
     stub_model_facade.generate.assert_called()
 
 
+@pytest.mark.parametrize(
+    (
+        "jinja_rendering_engine",
+        "expected_input_tokens_mean",
+        "expected_input_tokens_median",
+        "expected_input_tokens_stddev",
+    ),
+    [
+        (JinjaRenderingEngine.NATIVE, 10.0, 10.0, 0.0),
+        (
+            JinjaRenderingEngine.SECURE,
+            MissingValue.CALCULATION_FAILED,
+            MissingValue.CALCULATION_FAILED,
+            MissingValue.CALCULATION_FAILED,
+        ),
+    ],
+)
+def test_dataset_profiler_uses_run_config_jinja_engine_for_input_token_stats(
+    tmp_path: Path,
+    jinja_rendering_engine: JinjaRenderingEngine,
+    expected_input_tokens_mean: float | MissingValue,
+    expected_input_tokens_median: float | MissingValue,
+    expected_input_tokens_stddev: float | MissingValue,
+) -> None:
+    column_config = LLMTextColumnConfig(
+        name="summary",
+        prompt="Trajectory: {{ messages }}",
+        system_prompt="System prompt",
+        model_alias="nano",
+    )
+    dataset = lazy.pd.DataFrame(
+        {
+            "summary": ["response"],
+            "messages": ["x" * 512_001],
+        }
+    )
+    profiler = DataDesignerDatasetProfiler(
+        config=DatasetProfilerConfig(column_configs=[column_config]),
+        resource_provider=ResourceProvider(
+            artifact_storage=ArtifactStorage(artifact_path=tmp_path),
+            run_config=RunConfig(jinja_rendering_engine=jinja_rendering_engine),
+        ),
+    )
+
+    with patch(
+        "data_designer.engine.analysis.utils.column_statistics_calculations.count_text_tokens",
+        return_value=10,
+    ):
+        profile = profiler.profile_dataset(target_num_records=1, dataset=dataset)
+
+    stats = profile.column_statistics[0]
+    assert stats.input_tokens_mean == expected_input_tokens_mean
+    assert stats.input_tokens_median == expected_input_tokens_median
+    assert stats.input_tokens_stddev == expected_input_tokens_stddev
+
+
 @patch(
     "data_designer.engine.analysis.dataset_profiler.DataDesignerDatasetProfiler._validate_schema_consistency",
     autospec=True,

packages/data-designer-engine/tests/engine/analysis/utils/test_column_statistics_calculations.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 from itertools import cycle
+from unittest.mock import patch
 
 import pytest
 
@@ -16,6 +17,7 @@
     NumericalDistribution,
 )
 from data_designer.config.column_configs import LLMTextColumnConfig
+from data_designer.config.run_config import JinjaRenderingEngine
 from data_designer.config.utils.numerical_helpers import prepare_number_for_reporting
 from data_designer.engine.analysis.utils.column_statistics_calculations import (
     calculate_column_distribution,
@@ -188,6 +190,46 @@ def test_calculate_input_token_stats(mock_prompt_renderer_render, stub_column_co
     assert result["input_tokens_median"] == MissingValue.CALCULATION_FAILED
 
 
+@pytest.mark.parametrize(
+    ("prompt", "messages", "expected_token_count"),
+    [
+        ("Joined: {{ messages | join('-') }}", ["Hello", "World"], 4),
+        ("Trajectory: {{ messages }}", "x" * 512_001, 10),
+    ],
+)
+def test_calculate_input_token_stats_respects_native_jinja_engine(
+    prompt: str,
+    messages: list[str] | str,
+    expected_token_count: int,
+) -> None:
+    column_config = LLMTextColumnConfig(
+        name="test_column",
+        prompt=prompt,
+        system_prompt="System prompt",
+        model_alias="test_model_alias",
+    )
+    df = lazy.pd.DataFrame(
+        {
+            "test_column": ["response"],
+            "messages": [messages],
+        }
+    )
+
+    with patch(
+        "data_designer.engine.analysis.utils.column_statistics_calculations.count_text_tokens",
+        return_value=expected_token_count,
+    ):
+        result = calculate_input_token_stats(
+            column_config,
+            df,
+            jinja_rendering_engine=JinjaRenderingEngine.NATIVE,
+        )
+
+    assert result["input_tokens_mean"] == float(expected_token_count)
+    assert result["input_tokens_median"] == float(expected_token_count)
+    assert result["input_tokens_stddev"] == 0.0
+
+
 def test_calculate_output_token_stats(stub_column_config, stub_df_responses):
     result = calculate_output_token_stats(stub_column_config, stub_df_responses)
     assert "output_tokens_mean" in result