fix(generation): keep timeseries prompt context in the training serialization

src/nemo_safe_synthesizer/data_processing/assembler.py

@@ -986,8 +986,11 @@ def _get_initial_prefill(self) -> dict[str, str]:
             if len(seen_groups[group_value]) < 3:
                 seen_groups[group_value].append(record["text"])
 
-        # Convert lists to joined strings
-        return {group: " " + "\n".join(samples) for group, samples in seen_groups.items()}
+        # Each sample line is already newline-terminated (see
+        # _convert_records_to_jsonl), so concatenate directly: joining with
+        # "\n" would insert blank lines between records, a shape that never
+        # occurs in training examples.
+        return {group: " " + "".join(samples) for group, samples in seen_groups.items()}
 
     def _apply_train_test_split(self, dataset: Dataset) -> None:
         """Override split logic to preserve record order and split along group boundaries."""

src/nemo_safe_synthesizer/generation/timeseries_backend.py

@@ -67,8 +67,14 @@ class GroupState:
     current_prefill: str
     """Current prefill string, updated as generation progresses to include recently generated records."""
 
-    recent_records: list[dict] = field(default_factory=list)
-    """Sliding window of recently generated records used to build the next prompt context."""
+    recent_records: list[ParsedRecord] = field(default_factory=list)
+    """Sliding window of recently generated records used to build the next prompt context.
+
+    Holds ``ParsedRecord`` so the next prompt context can reuse the exact
+    bytes the model emitted (``.text``): re-serializing the parsed dicts
+    would drift the prompt into a JSON dialect the model was not trained
+    on (spacing, slash escaping, float formatting).
+    """
 
     expected_records: int = 0
     """Target record count, calculated from ``(stop_timestamp - start_timestamp) / interval_seconds``."""
@@ -513,12 +519,18 @@ def _check_chronological_for_group(self, batch: Batch, group_state: GroupState)
                 if record.is_valid:
                     record.invalidate(error)
 
-    def _update_group_state(self, group_state: GroupState, records: list[dict]) -> None:
-        """Update a group's state with new records.
+    def _update_group_state(self, group_state: GroupState, records: list[ParsedRecord]) -> None:
+        """Update a group's state with new valid records.
+
+        The rebuilt prefill reuses each record's original ``text`` (the
+        bytes the model emitted, which match the training serialization)
+        and mirrors the initial prefill's shape from
+        ``SequentialExampleAssembler._get_initial_prefill``: a leading
+        space, then newline-terminated records.
 
         Args:
             group_state: The group state to update.
-            records: The new valid records.
+            records: The new valid records (``parsed`` is set on each).
         """
         if not records:
             return
@@ -529,12 +541,12 @@ def _update_group_state(self, group_state: GroupState, records: list[dict]) -> N
 
         # Update prefill
         tail = group_state.recent_records[-self._prefill_context_size :]
-        lines = [json.dumps(record, ensure_ascii=False) for record in tail]
-        group_state.current_prefill = "\n".join(lines) + "\n"
+        group_state.current_prefill = " " + "".join(f"{record.text}\n" for record in tail)
 
         # Update last timestamp
-        last_record = records[-1]
-        timestamp_seconds = self._parse_timestamp_seconds(last_record.get(self._time_column))
+        last_parsed = records[-1].parsed or {}
+        timestamp_value = last_parsed.get(self._time_column) if self._time_column is not None else None
+        timestamp_seconds = self._parse_timestamp_seconds(timestamp_value)
         if timestamp_seconds is not None:
             group_state.last_timestamp_seconds = timestamp_seconds
 
@@ -621,9 +633,9 @@ def _process_group_result(
         if self.config.time_series.timestamp_interval_seconds is not None:
             self._check_chronological_for_group(batch, state)
 
-        batch_records = self._retain_single_valid_response(batch)
-        reached_stop = self._has_reached_stop_time(batch_records)
-        self._update_group_state(state, batch_records)
+        retained_records = self._retain_single_valid_response(batch)
+        reached_stop = self._has_reached_stop_time([r.parsed for r in retained_records if r.parsed is not None])
+        self._update_group_state(state, retained_records)
 
         # Check if batch has high invalid fraction
         invalid_fraction = 1.0 - batch.valid_record_fraction
@@ -852,7 +864,7 @@ def _generate_parallel_groups(
 
         return all_groups_succeeded
 
-    def _retain_single_valid_response(self, batch: Batch) -> list[dict]:
+    def _retain_single_valid_response(self, batch: Batch) -> list[ParsedRecord]:
         """Retain the response with the most valid records, discarding all others.
 
         For time-series sliding window generation, only one response can be used
@@ -865,9 +877,11 @@ def _retain_single_valid_response(self, batch: Batch) -> list[dict]:
             batch: The batch to retain the response from.
 
         Returns:
-            List of valid records from the retained response.
+            The retained response's valid ``ParsedRecord`` objects, keeping
+            both the parsed dicts and the original emitted text (the latter
+            feeds the next prompt context).
         """
-        final_records: list[dict] = []
+        final_records: list[ParsedRecord] = []
 
         # Find the index of the response with the most valid records.
         max_valid_idx = None
@@ -886,7 +900,7 @@ def _retain_single_valid_response(self, batch: Batch) -> list[dict]:
                 # error statistics without carrying stale text/token counts.
                 response.records = [ParsedRecord(text="", error=trim_error)]
             else:
-                final_records.extend(response.valid_records)
+                final_records.extend(r for r in response.records if r.is_valid and r.parsed is not None)
 
         return final_records
 

tests/data_processing/test_assembler.py

@@ -809,21 +809,15 @@ def test_sequential_assembler_initial_prefill(
     assert "A" in prefill
     assert "B" in prefill
 
-    # Each prefill should contain up to 3 records as JSONL (newline-separated)
-    # Group A has 3 records, Group B has 2 records
-    # Filter out empty lines that may appear between records
-    prefill_a_lines = [line for line in prefill["A"].strip().split("\n") if line]
-    prefill_b_lines = [line for line in prefill["B"].strip().split("\n") if line]
-
-    assert len(prefill_a_lines) == 3  # All 3 records from group A
-    assert len(prefill_b_lines) == 2  # Both records from group B
-
-    # Verify the records contain expected values (order should be by time)
-    assert '"value": 10' in prefill_a_lines[0] or '"value":10' in prefill_a_lines[0]
-    assert '"value": 20' in prefill_a_lines[1] or '"value":20' in prefill_a_lines[1]
-    assert '"value": 30' in prefill_a_lines[2] or '"value":30' in prefill_a_lines[2]
-    assert '"value": 100' in prefill_b_lines[0] or '"value":100' in prefill_b_lines[0]
-    assert '"value": 200' in prefill_b_lines[1] or '"value":200' in prefill_b_lines[1]
+    # Pin the exact byte shape: a leading space, then newline-terminated
+    # training-dialect (pandas to_json) records with single newlines between
+    # them -- the same shape training examples use. Blank lines or Python
+    # json.dumps spacing here would put the generation prompt in a dialect
+    # the model never saw in training.
+    assert prefill["A"] == (
+        ' {"group":"A","time":1,"value":10}\n{"group":"A","time":2,"value":20}\n{"group":"A","time":3,"value":30}\n'
+    )
+    assert prefill["B"] == ' {"group":"B","time":1,"value":100}\n{"group":"B","time":2,"value":200}\n'
 
 
 def test_should_flush_example_boundary_conditions():

tests/generation/test_timeseries_backend.py

@@ -22,6 +22,7 @@
     TimeSeriesParameters,
     TrainingHyperparams,
 )
+from nemo_safe_synthesizer.data_processing.record_utils import ParsedRecord
 from nemo_safe_synthesizer.defaults import DEFAULT_MAX_SEQ_LENGTH, PSEUDO_GROUP_COLUMN
 from nemo_safe_synthesizer.generation.processors import TimeSeriesDataProcessor
 from nemo_safe_synthesizer.generation.results import GenerationBatches
@@ -387,16 +388,47 @@ def test_appends_records_and_updates_prefill(self, timeseries_base_params, times
         )
 
         records = [
-            {"timestamp": "2024-01-01 00:00:00", "value": 1},
-            {"timestamp": "2024-01-01 01:00:00", "value": 2},
+            ParsedRecord(
+                text='{"timestamp":"2024-01-01 00:00:00","value":1}',
+                parsed={"timestamp": "2024-01-01 00:00:00", "value": 1},
+            ),
+            ParsedRecord(
+                text='{"timestamp":"2024-01-01 01:00:00","value":2}',
+                parsed={"timestamp": "2024-01-01 01:00:00", "value": 2},
+            ),
         ]
 
         backend._update_group_state(state, records)
 
         assert len(state.recent_records) == 2
-        assert '"timestamp"' in state.current_prefill
+        # The rebuilt prefill must reuse the records' emitted text verbatim
+        # (training dialect) and mirror the initial prefill's shape: leading
+        # space, newline-terminated records, single newlines between them.
+        assert state.current_prefill == (
+            ' {"timestamp":"2024-01-01 00:00:00","value":1}\n{"timestamp":"2024-01-01 01:00:00","value":2}\n'
+        )
         assert state.last_timestamp_seconds is not None
 
+    def test_updates_timestamp_for_empty_string_column_name(
+        self, timeseries_base_params, timeseries_model_metadata, mock_workdir
+    ):
+        """Test that an empty-string timestamp column name remains valid."""
+        backend = create_timeseries_backend(timeseries_base_params, timeseries_model_metadata, mock_workdir)
+        backend._time_column = ""
+        state = GroupState(
+            group_id="test",
+            initial_prefill="",
+            current_prefill="",
+            expected_records=10,
+            last_timestamp_seconds=0,
+        )
+        timestamp = "2024-01-01 01:00:00"
+        records = [ParsedRecord(text='{"":"2024-01-01 01:00:00"}', parsed={"": timestamp})]
+
+        backend._update_group_state(state, records)
+
+        assert state.last_timestamp_seconds == backend._parse_timestamp_seconds(timestamp)
+
 
 class TestGetTimestampFromPrefill:
     """Tests for the _get_timestamp_from_prefill method."""