Preserve validated Pydantic values when freezing arguments

great_ai/helper/freeze_arguments.py

@@ -13,8 +13,10 @@ def _hashable_model_type(model_type: Type[BaseModel]) -> Type[BaseModel]:
     # operands to share a type, so creating a fresh subclass on every freeze()
     # call would make even identical models compare unequal.
     class HashableValue(model_type):  # type: ignore
+        __slots__ = ()
+
         def __hash__(self) -> int:
-            return hash(frozenset((k, freeze(v)) for k, v in self.model_dump().items()))
+            return hash(freeze(self.model_dump()))
 
     return HashableValue
 
@@ -67,6 +69,10 @@ def freeze(value: Union[Sequence[Any], Mapping[str, Any], Set[Any], BaseModel])
         return FrozenSet(value)
 
     if isinstance(value, BaseModel):
-        return _hashable_model_type(type(value))(**value.model_dump())
+        # Keep the already-validated values, aliases, nested types and private
+        # attributes. Reconstructing from model_dump() runs validators again.
+        frozen = value.model_copy()
+        object.__setattr__(frozen, "__class__", _hashable_model_type(type(value)))
+        return frozen
 
     return value

tests/test_prediction_regressions.py

@@ -1,13 +1,41 @@
 from datetime import datetime
-from typing import Any, Dict
+from typing import Any, Dict, List
 from unittest.mock import patch
 
 from great_ai import Trace
+from great_ai.helper.freeze_arguments import freeze
 from great_ai.persistence.mongodb_driver import MongoDbDriver
+from pydantic import BaseModel, Field, field_validator
 from pymongo import MongoClient
 from pymongo.synchronous.collection import Collection
 
 
+class Item(BaseModel):
+    value: int = Field(alias="number")
+
+    @field_validator("value")
+    @classmethod
+    def increment(cls, value: int) -> int:
+        return value + 1
+
+
+class Payload(BaseModel):
+    item: Item
+    others: List[Item]
+
+
+def test_freezing_preserves_validated_models() -> None:
+    original = Payload(item=Item(number=1), others=[Item(number=2)])
+    frozen = freeze(original)
+    assert isinstance(frozen, Payload)
+    assert isinstance(frozen.item, Item)
+    assert frozen.item.value == 2
+    assert frozen.others[0].value == 3
+    assert frozen.model_dump() == original.model_dump()
+    assert hash(frozen) == hash(freeze(original))
+    assert type(original) is Payload
+
+
 def test_flattening_preserves_reserved_trace_fields() -> None:
     trace = Trace(
         trace_id="actual-id",