Serialise BaseModels

src/great_ai/great_ai/helper/__init__.py

@@ -1,4 +1,4 @@
-from .freeze_arguments import freeze_arguments
+from .freeze_arguments import freeze, freeze_arguments
 from .get_arguments import get_arguments
 from .get_function_metadata_store import get_function_metadata_store
 from .hashable_base_model import HashableBaseModel

src/great_ai/great_ai/helper/freeze_arguments.py

@@ -1,20 +1,22 @@
 from functools import wraps
 from typing import Any, Callable, Dict, List, Set, Union
 
+from pydantic import BaseModel
+
 
 class FrozenDict(dict):
     def __hash__(self) -> int:
-        return hash(frozenset(self.items()))
+        return hash(frozenset((k, freeze(v)) for k, v in self.items()))
 
 
 class FrozenList(list):
     def __hash__(self) -> int:
-        return hash(tuple(self))
+        return hash(tuple(freeze(i) for i in self))
 
 
 class FrozenSet(set):
     def __hash__(self) -> int:
-        return hash(frozenset(self))
+        return hash(frozenset(freeze(i) for i in self))
 
 
 def freeze_arguments(func: Callable[..., Any]) -> Callable[..., Any]:
@@ -26,14 +28,14 @@ def freeze_arguments(func: Callable[..., Any]) -> Callable[..., Any]:
 
     @wraps(func)
     def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
-        args = tuple(_freeze(arg) for arg in args)
-        kwargs = {k: _freeze(v) for k, v in kwargs.items()}
+        args = tuple(freeze(arg) for arg in args)
+        kwargs = {k: freeze(v) for k, v in kwargs.items()}
         return func(*args, **kwargs)
 
     return wrapper
 
 
-def _freeze(value: Union[List[Any], Dict[str, Any], Set[Any]]) -> Any:
+def freeze(value: Union[List[Any], Dict[str, Any], Set[Any]]) -> Any:
     if isinstance(value, dict):
         return FrozenDict(value)
 
@@ -43,4 +45,12 @@ def _freeze(value: Union[List[Any], Dict[str, Any], Set[Any]]) -> Any:
     if isinstance(value, set):
         return FrozenSet(value)
 
+    if isinstance(value, BaseModel):
+
+        class HashableValue(type(value)):
+            def __hash__(self) -> int:
+                return hash(frozenset((k, freeze(v)) for k, v in self.dict().items()))
+
+        return HashableValue(**value.dict())
+
     return value

src/great_ai/great_ai/persistence/mongodb_driver.py

@@ -121,7 +121,7 @@ class MongodbDriver(TracingDatabaseDriver):
         return documents, len(documents)
 
     def update(self, id: str, new_version: Trace) -> None:
-        serialized = new_version.dict()
+        serialized = new_version.to_flat_dict()
         serialized["_id"] = new_version.trace_id
         with MongoClient(self.mongo_connection_string) as client:
             client[self.mongo_database].traces.update_one(id, new_version)