From f71485ea2b918edab0c60ca25e71a3af331c864c Mon Sep 17 00:00:00 2001 From: vasilije Date: Tue, 17 Dec 2024 20:20:10 +0100 Subject: [PATCH] Add versioning to the data point model --- .github/workflows/profiling.yaml | 1 + cognee/infrastructure/engine/models/DataPoint.py | 9 +++++++++ 2 files changed, 10 insertions(+) diff --git a/.github/workflows/profiling.yaml b/.github/workflows/profiling.yaml index 93ed82f82..6633e376e 100644 --- a/.github/workflows/profiling.yaml +++ b/.github/workflows/profiling.yaml @@ -57,6 +57,7 @@ jobs: run: | poetry install --no-interaction --all-extras poetry run pip install pyinstrument + poetry run pip install parso # Set environment variables for SHAs diff --git a/cognee/infrastructure/engine/models/DataPoint.py b/cognee/infrastructure/engine/models/DataPoint.py index cf49cb4e1..9cc818228 100644 --- a/cognee/infrastructure/engine/models/DataPoint.py +++ b/cognee/infrastructure/engine/models/DataPoint.py @@ -81,3 +81,12 @@ class DataPoint(BaseModel): """Deserialize the instance from pickled bytes.""" data = pickle.loads(pickled_data) return self(**data) + + def to_dict(self, **kwargs) -> Dict[str, Any]: + """Serialize model to a dictionary.""" + return self.model_dump(**kwargs) + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "DataPoint": + """Deserialize model from a dictionary.""" + return cls.model_validate(data)