Handle LLMMessage serialization at the storage layer (#5024)

Summary:

Move `LLMMessage` dict conversion from the `experiment.llm_messages`
getter/setter to the storage encoders/decoders, following Ax convention
that domain objects hold domain types and serialization happens at the
storage boundary.

**`experiment.py`**: The setter now stores `LLMMessage` objects directly
in `_properties`. The getter handles both `LLMMessage` objects (new path)
and plain dicts (backward compat with previously stored data).

**JSON store**: No explicit changes needed — the encoder's generic
dataclass fallback auto-serializes `LLMMessage` with a `__type` tag,
and `LLMMessage` is already registered in `CORE_DECODER_REGISTRY`.

**SQA store**: The encoder converts `LLMMessage` → dict via
`dataclasses.asdict()` in the properties copy before DB write (same
pattern as `pruning_target_parameterization`). The decoder converts
dicts → `LLMMessage` after loading properties, in both
`_init_experiment_from_sqa` and `_init_mt_experiment_from_sqa`.

Reviewed By: lena-kashtelyan

Differential Revision: D96434290

Author: ItsMrLin

ax/storage/sqa_store/decoder.py

@@ -31,6 +31,7 @@
 from ax.core.data import Data
 from ax.core.experiment import Experiment
 from ax.core.generator_run import GeneratorRun
+from ax.core.llm_provider import LLMMessage
 from ax.core.metric import Metric
 from ax.core.multi_type_experiment import MultiTypeExperiment
 from ax.core.objective import MultiObjective, Objective, ScalarizedObjective
@@ -220,6 +221,10 @@ def _init_experiment_from_sqa(
         # `experiment_sqa.properties` is `sqlalchemy.ext.mutable.MutableDict`
         # so need to convert it to regular dict.
         properties = dict(experiment_sqa.properties or {})
+        if Keys.LLM_MESSAGES in properties:
+            properties[Keys.LLM_MESSAGES] = [
+                LLMMessage(**m) for m in properties[Keys.LLM_MESSAGES]
+            ]
         opt_config, tracking_metrics = self.opt_config_and_tracking_metrics_from_sqa(
             metrics_sqa=experiment_sqa.metrics,
             pruning_target_parameterization=(
@@ -281,6 +286,10 @@ def _init_mt_experiment_from_sqa(
     ) -> MultiTypeExperiment:
         """First step of conversion within experiment_from_sqa."""
         properties = dict(experiment_sqa.properties or {})
+        if Keys.LLM_MESSAGES in properties:
+            properties[Keys.LLM_MESSAGES] = [
+                LLMMessage(**m) for m in properties[Keys.LLM_MESSAGES]
+            ]
         opt_config, tracking_metrics = self.opt_config_and_tracking_metrics_from_sqa(
             metrics_sqa=experiment_sqa.metrics,
             pruning_target_parameterization=(

ax/storage/sqa_store/encoder.py

@@ -6,6 +6,7 @@
 
 # pyre-strict
 
+import dataclasses
 from enum import Enum
 from logging import Logger
 from typing import Any, cast
@@ -28,6 +29,7 @@
 from ax.core.evaluations_to_data import DataType
 from ax.core.experiment import Experiment
 from ax.core.generator_run import GeneratorRun
+from ax.core.llm_provider import LLMMessage
 from ax.core.metric import Metric
 from ax.core.multi_type_experiment import MultiTypeExperiment
 from ax.core.objective import MultiObjective, Objective, ScalarizedObjective
@@ -234,6 +236,11 @@ def experiment_to_sqa(self, experiment: Experiment) -> SQAExperiment:
             properties["pruning_target_parameterization"] = arm_to_dict(
                 oc.pruning_target_parameterization
             )
+        if Keys.LLM_MESSAGES in properties:
+            properties[Keys.LLM_MESSAGES] = [
+                dataclasses.asdict(m) if isinstance(m, LLMMessage) else m
+                for m in properties[Keys.LLM_MESSAGES]
+            ]
 
         # pyre-ignore[9]: Expected `Base` for 1st...yping.Type[Experiment]`.
         experiment_class: type[SQAExperiment] = self.config.class_to_sqa_class[