add GenericModelType to fix serialization (#1672)

* fix: add GenericModelType to fix serialization

also sets all GenericModelType fields to Optional with default=None

* chore: linting

* fix: whitespace

Author: tmchartrand

src/aind_data_schema/base.py

@@ -14,6 +14,7 @@
     Field,
     NaiveDatetime,
     PrivateAttr,
+    SerializeAsAny,
     ValidationError,
     ValidatorFunctionWrapHandler,
     create_model,
@@ -79,10 +80,7 @@ def has_corrupt_keys(input) -> bool:
 
 class GenericModel(BaseModel, extra="allow"):
     """Base class for generic types that can be used in AIND schema"""
-
     # extra="allow" is needed because BaseModel by default drops extra parameters.
-    # Alternatively, consider using 'SerializeAsAny' once this issue is resolved
-    # https://github.com/pydantic/pydantic/issues/6423
 
     @model_validator(mode="after")
     def validate_fieldnames(self):
@@ -95,6 +93,8 @@ def validate_fieldnames(self):
         return self
 
 
+GenericModelType = SerializeAsAny[GenericModel]
+
 T = TypeVar("T")
 Discriminated = Annotated[T, Field(discriminator="object_type")]
 DiscriminatedList = List[Discriminated[T]]

src/aind_data_schema/components/configs.py

@@ -22,7 +22,7 @@
 from pydantic import Field, field_validator, model_validator
 from pydantic_core.core_schema import ValidationInfo
 
-from aind_data_schema.base import AwareDatetimeWithDefault, DataModel, DiscriminatedList, GenericModel
+from aind_data_schema.base import AwareDatetimeWithDefault, DataModel, DiscriminatedList, GenericModelType
 from aind_data_schema.components.coordinates import (
     TRANSFORM_TYPES,
     AtlasCoordinate,
@@ -542,7 +542,7 @@ class MRIScan(DeviceConfig):
     # other fields
     resolution: Optional[Scale] = Field(default=None, title="Voxel resolution")
     resolution_unit: Optional[SizeUnit] = Field(default=None, title="Voxel resolution unit")
-    additional_scan_parameters: GenericModel = Field(..., title="Parameters")
+    additional_scan_parameters: Optional[GenericModelType] = Field(default=None, title="Parameters")
     notes: Optional[str] = Field(default=None, title="Notes", validate_default=True)
 
     @field_validator("notes", mode="after")

src/aind_data_schema/components/devices.py

@@ -41,7 +41,7 @@
 )
 from pydantic import Field, ValidationInfo, field_validator, model_validator
 
-from aind_data_schema.base import DataModel, Discriminated, GenericModel
+from aind_data_schema.base import DataModel, Discriminated, GenericModelType
 from aind_data_schema.components.coordinates import TRANSFORM_TYPES, AxisName, CoordinateSystem, Scale
 from aind_data_schema.components.identifiers import Software
 
@@ -78,7 +78,7 @@ class Device(DataModel):
     model: Optional[str] = Field(default=None, title="Model")
 
     # Additional fields
-    additional_settings: Optional[GenericModel] = Field(default=None, title="Additional parameters")
+    additional_settings: Optional[GenericModelType] = Field(default=None, title="Additional parameters")
     notes: Optional[str] = Field(default=None, title="Notes")
 
     @model_validator(mode="after")

src/aind_data_schema/components/identifiers.py

@@ -7,7 +7,7 @@
 from aind_data_schema_models.registries import Registry
 from pydantic import Field
 
-from aind_data_schema.base import DataModel, DiscriminatedList, GenericModel
+from aind_data_schema.base import DataModel, DiscriminatedList, GenericModelType
 
 
 class Database(str, Enum):
@@ -83,7 +83,7 @@ class Code(DataModel):
     input_data: Optional[DiscriminatedList[DataAsset | CombinedData]] = Field(
         default=None, title="Input data", description="Input data used in the code or script"
     )
-    parameters: Optional[GenericModel] = Field(
+    parameters: Optional[GenericModelType] = Field(
         default=None, title="Parameters", description="Parameters used in the code or script"
     )
 

src/aind_data_schema/components/measurements.py

@@ -6,7 +6,7 @@
 from aind_data_schema_models.units import UNITS, PowerUnit, TimeUnit, VolumeUnit, VoltageUnit
 from pydantic import model_validator
 
-from aind_data_schema.base import AwareDatetimeWithDefault, DataModel, Discriminated, Field, GenericModel
+from aind_data_schema.base import AwareDatetimeWithDefault, DataModel, Discriminated, Field, GenericModelType
 from aind_data_schema.components.configs import DeviceConfig
 from aind_data_schema.components.reagent import Reagent
 from aind_data_schema.utils.validators import TimeValidation
@@ -27,7 +27,7 @@ class CalibrationFit(DataModel):
         ...,
         title="Fit type",
     )
-    fit_parameters: Optional[GenericModel] = Field(
+    fit_parameters: Optional[GenericModelType] = Field(
         default=None,
         title="Fit parameters",
         description="Parameters of the fit equation, e.g. slope and intercept for linear fit",

src/aind_data_schema/components/stimulus.py

@@ -7,7 +7,7 @@
 from aind_data_schema_models.units import FrequencyUnit, PowerUnit, TimeUnit, ConcentrationUnit
 from pydantic import Field, model_validator
 
-from aind_data_schema.base import DataModel, GenericModel
+from aind_data_schema.base import DataModel, GenericModel, GenericModelType
 
 
 class PulseShape(str, Enum):
@@ -48,16 +48,16 @@ class OptoStimulation(GenericModel):
         description="Duration of baseline recording prior to first pulse train",
     )
     baseline_duration_unit: TimeUnit = Field(default=TimeUnit.S, title="Baseline duration unit")
-    other_parameters: GenericModel = Field(GenericModel(), title="Other parameters")
+    other_parameters: Optional[GenericModelType] = Field(default=None, title="Other parameters")
     notes: Optional[str] = Field(default=None, title="Notes")
 
 
 class VisualStimulation(GenericModel):
     """Description of visual stimulus parameters. Provides a high level description of stimulus."""
 
     stimulus_name: str = Field(..., title="Stimulus name")
-    stimulus_parameters: GenericModel = Field(
-        GenericModel(),
+    stimulus_parameters: Optional[GenericModelType] = Field(
+        default=None,
         title="Stimulus parameters",
         description="Define and list the parameter values used (e.g. all TF or orientation values)",
     )
@@ -82,7 +82,7 @@ class PhotoStimulationGroup(DataModel):
     spiral_duration_unit: TimeUnit = Field(default=TimeUnit.S, title="Spiral duration unit")
     inter_spiral_interval: Decimal = Field(..., title="Inter trial interval (s)")
     inter_spiral_interval_unit: TimeUnit = Field(default=TimeUnit.S, title="Inter trial interval unit")
-    other_parameters: GenericModel = Field(GenericModel(), title="Other parameters")
+    other_parameters: Optional[GenericModelType] = Field(default=None, title="Other parameters")
     notes: Optional[str] = Field(default=None, title="Notes")
 
 
@@ -94,7 +94,7 @@ class PhotoStimulation(GenericModel):
     groups: List[PhotoStimulationGroup] = Field(..., title="Groups")
     inter_trial_interval: Decimal = Field(..., title="Inter trial interval (s)")
     inter_trial_interval_unit: TimeUnit = Field(default=TimeUnit.S, title="Inter trial interval unit")
-    other_parameters: GenericModel = Field(GenericModel(), title="Other parameters")
+    other_parameters: Optional[GenericModelType] = Field(default=None, title="Other parameters")
     notes: Optional[str] = Field(default=None, title="Notes")
 
 

src/aind_data_schema/core/acquisition.py

@@ -10,7 +10,13 @@
 from aind_data_schema.utils.validators import TimeValidation
 from pydantic import Field, SkipValidation, model_validator
 
-from aind_data_schema.base import AwareDatetimeWithDefault, DataCoreModel, DataModel, DiscriminatedList, GenericModel
+from aind_data_schema.base import (
+    AwareDatetimeWithDefault,
+    DataCoreModel,
+    DataModel,
+    DiscriminatedList,
+    GenericModelType
+)
 from aind_data_schema.components.configs import (
     AirPuffConfig,
     CatheterConfig,
@@ -83,7 +89,7 @@ class AcquisitionSubjectDetails(DataModel):
 class PerformanceMetrics(DataModel):
     """Summary of a StimulusEpoch"""
 
-    output_parameters: GenericModel = Field(default=GenericModel(), title="Additional metrics")
+    output_parameters: Optional[GenericModelType] = Field(default=None, title="Additional metrics")
     reward_consumed_during_epoch: Optional[Decimal] = Field(default=None, title="Reward consumed during training (uL)")
     reward_consumed_unit: Optional[VolumeUnit] = Field(default=None, title="Reward consumed unit")
     trials_total: Optional[int] = Field(default=None, title="Total trials")

src/aind_data_schema/core/model.py

@@ -5,7 +5,7 @@
 from aind_data_schema_models.system_architecture import ModelArchitecture
 from pydantic import Field
 
-from aind_data_schema.base import DataCoreModel, DataModel, DiscriminatedList, GenericModel
+from aind_data_schema.base import DataCoreModel, DataModel, DiscriminatedList, GenericModelType
 from aind_data_schema.components.identifiers import Code, Software
 from aind_data_schema.core.processing import DataProcess, ProcessName
 
@@ -59,8 +59,8 @@ class Model(DataCoreModel):
     )
     architecture: ModelArchitecture = Field(..., title="architecture", description="Model architecture / type of model")
     software_framework: Optional[Software] = Field(default=None, title="Software framework")
-    architecture_parameters: GenericModel = Field(
-        default=GenericModel(),
+    architecture_parameters: Optional[GenericModelType] = Field(
+        default=None,
         title="Architecture parameters",
         description="Parameters of model architecture, such as input signature or number of layers.",
     )

src/aind_data_schema/core/processing.py

@@ -9,7 +9,7 @@
 from aind_data_schema_models.units import MemoryUnit, UnitlessUnit
 from pydantic import Field, SkipValidation, ValidationInfo, field_validator, model_validator
 
-from aind_data_schema.base import AwareDatetimeWithDefault, DataCoreModel, DataModel, GenericModel
+from aind_data_schema.base import AwareDatetimeWithDefault, DataCoreModel, DataModel, GenericModelType
 from aind_data_schema.components.identifiers import Code
 from aind_data_schema.components.wrappers import AssetPath
 from aind_data_schema.utils.merge import merge_notes, merge_optional_list, merge_process_graph
@@ -71,7 +71,11 @@ class DataProcess(DataModel):
     output_path: Optional[AssetPath] = Field(
         default=None, title="Output path", description="Path to processing outputs, if stored."
     )
-    output_parameters: GenericModel = Field(default=GenericModel(), description="Output parameters", title="Outputs")
+    output_parameters: Optional[GenericModelType] = Field(
+        default=None,
+        description="Output parameters",
+        title="Outputs"
+    )
     notes: Optional[str] = Field(default=None, title="Notes", validate_default=True)
     resources: Optional[ResourceUsage] = Field(default=None, title="Process resource usage")
 

tests/test_aind_generic.py

@@ -2,19 +2,19 @@
 
 import unittest
 
-from pydantic import BaseModel, Field
+from pydantic import Field
 
-from aind_data_schema.base import DataModel, GenericModel
+from aind_data_schema.base import DataModel, GenericModel, GenericModelType
 
 
 class GenericContainer(DataModel):
     """Represents a generic container"""
 
-    contains_model: GenericModel
-    contains_dict: GenericModel
+    contains_model: GenericModelType
+    contains_dict: GenericModelType
 
 
-class Bar(BaseModel):
+class Bar(GenericModel):
     """Represents a mock model"""
 
     bar: str = Field(default="bar")
@@ -58,7 +58,7 @@ def test_sub_container_from_container(self):
             contains_dict={"foodict": 1, "bardict": "bar"},
         )
         parent_container = GenericContainer(
-            contains_model=Bar(bar="baz", foo=2).model_dump(),
+            contains_model=Bar(bar="baz", foo=2),
             contains_dict={"foodict": 1, "bardict": "bar"},
         )
         deserialized = SubGenericContainer.model_validate_json(parent_container.model_dump_json())