diff --git a/great_ai/helper/freeze_arguments.py b/great_ai/helper/freeze_arguments.py index e6ff56f..524906f 100644 --- a/great_ai/helper/freeze_arguments.py +++ b/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 diff --git a/tests/test_prediction_regressions.py b/tests/test_prediction_regressions.py index a556a9e..3971b11 100644 --- a/tests/test_prediction_regressions.py +++ b/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",