fix(tracing): preserve string payloads in HDF5 round trips

openevolve/utils/trace_export_utils.py

@@ -80,7 +80,7 @@ def export_traces_json(
     logger.info(f"Exported {len(traces)} traces to {output_path}")
 
 
-# Names the metadata attributes that were JSON-encoded, so a list or dict can be told apart
+# Names the attributes that were JSON-encoded, so a list or dict can be told apart
 # from a string that merely looks like JSON
 _HDF5_JSON_KEYS_ATTR = "__json_encoded_keys__"
 
@@ -129,6 +129,7 @@ def export_traces_hdf5(
         for i, trace in enumerate(traces):
             trace_dict = trace.to_dict() if hasattr(trace, "to_dict") else trace
             trace_group = traces_group.create_group(f"trace_{i:06d}")
+            json_encoded_keys = []
 
             for key, value in trace_dict.items():
                 if value is None:
@@ -137,6 +138,7 @@ def export_traces_hdf5(
                 if isinstance(value, dict):
                     # Store dictionaries as JSON strings in attributes
                     trace_group.attrs[key] = json.dumps(value)
+                    json_encoded_keys.append(key)
                 elif isinstance(value, list):
                     # Convert lists to numpy arrays if possible
                     try:
@@ -145,6 +147,7 @@ def export_traces_hdf5(
                     except (ValueError, TypeError):
                         # Fall back to JSON for complex lists
                         trace_group.attrs[key] = json.dumps(value)
+                        json_encoded_keys.append(key)
                 elif isinstance(value, str):
                     # Store strings as attributes
                     trace_group.attrs[key] = value
@@ -154,6 +157,9 @@ def export_traces_hdf5(
                 else:
                     # Store other types as JSON
                     trace_group.attrs[key] = json.dumps(value)
+                    json_encoded_keys.append(key)
+
+            trace_group.attrs[_HDF5_JSON_KEYS_ATTR] = json.dumps(json_encoded_keys)
 
     logger.info(f"Exported {len(traces)} traces to {output_path}")
 
@@ -284,12 +290,21 @@ def load_traces_hdf5(input_path: Union[str, Path]) -> tuple[List[Dict[str, Any]]
             for trace_name in sorted(traces_group.keys()):
                 trace_group = traces_group[trace_name]
                 trace_dict = {}
+                json_encoded_keys = trace_group.attrs.get(_HDF5_JSON_KEYS_ATTR)
+                if json_encoded_keys is not None:
+                    json_encoded_keys = set(json.loads(json_encoded_keys))
 
                 # Load attributes
                 for key in trace_group.attrs:
+                    if key == _HDF5_JSON_KEYS_ATTR:
+                        continue
                     value = trace_group.attrs[key]
-                    # Try to parse JSON strings
-                    if isinstance(value, str) and (value.startswith("{") or value.startswith("[")):
+                    if json_encoded_keys is not None:
+                        trace_dict[key] = json.loads(value) if key in json_encoded_keys else value
+                    # Preserve the heuristic for files written without an encoding key list.
+                    elif isinstance(value, str) and (
+                        value.startswith("{") or value.startswith("[")
+                    ):
                         try:
                             trace_dict[key] = json.loads(value)
                         except json.JSONDecodeError:

tests/test_trace_export_hdf5.py

@@ -4,6 +4,7 @@
 import unittest
 from pathlib import Path
 
+from openevolve.evolution_trace import EvolutionTrace
 from openevolve.utils.trace_export_utils import export_traces, load_traces
 
 try:
@@ -65,5 +66,49 @@ def test_files_without_a_key_list_load_as_before(self):
         self.assertEqual(loaded["broken"], "{not json")
 
 
+@unittest.skipIf(h5py is None, "h5py is not installed")
+class TestHdf5TraceRoundTrip(unittest.TestCase):
+    def test_json_looking_code_and_response_stay_strings(self):
+        trace = EvolutionTrace(
+            iteration=1,
+            timestamp=123.0,
+            parent_id="parent",
+            child_id="child",
+            parent_metrics={"score": 1.0},
+            child_metrics={"score": 2.0},
+            parent_code="[1, 2]",
+            child_code='{"answer": 42}',
+            llm_response='{"code": "result = 42"}',
+            child_changes_description='["improved score"]',
+            artifacts={"nested": [{"value": 1}]},
+        )
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "trace.h5"
+            export_traces([trace], path, format="hdf5")
+            traces, _ = load_traces(path)
+        self.assertEqual(traces, [trace.to_dict()])
+
+    def test_numeric_and_json_encoded_lists_round_trip(self):
+        trace = {"scores": [1.0, 2.0], "labels": ["alpha"], "nested": [{"a": [1]}]}
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "trace.h5"
+            export_traces([trace], path, format="hdf5")
+            traces, _ = load_traces(path)
+        self.assertEqual(traces, [trace])
+
+    def test_trace_without_a_key_list_loads_as_before(self):
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "legacy.h5"
+            with h5py.File(path, "w") as f:
+                attrs = f.create_group("traces/trace_000000").attrs
+                attrs["metrics"] = '{"score": 1}'
+                attrs["labels"] = '["alpha"]'
+                attrs["broken"] = "{not json"
+            traces, _ = load_traces(path)
+        self.assertEqual(
+            traces, [{"metrics": {"score": 1}, "labels": ["alpha"], "broken": "{not json"}]
+        )
+
+
 if __name__ == "__main__":
     unittest.main()