Replace parse_obj with model_validate

topoteretes · Nov 4, 2024 · cb20506 · cb20506
1 parent 115e93b
commit cb20506
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py b/cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py
@@ -59,7 +59,7 @@ class Keyword(BaseModel):
             data_points.append(
                 DataPoint[Keyword](
                     id=str(classification_type_id),
-                    payload=Keyword.parse_obj({
+                    payload=Keyword.model_validate({
                         "uuid": str(classification_type_id),
                         "text": classification_type_label,
                         "chunk_id": str(data_chunk.chunk_id),
@@ -98,7 +98,7 @@ class Keyword(BaseModel):
                 data_points.append(
                     DataPoint[Keyword](
                         id=str(classification_subtype_id),
-                        payload=Keyword.parse_obj({
+                        payload=Keyword.model_validate({
                             "uuid": str(classification_subtype_id),
                             "text": classification_subtype_label,
                             "chunk_id": str(data_chunk.chunk_id),

diff --git a/cognee/tasks/infer_data_ontology/infer_data_ontology.py b/cognee/tasks/infer_data_ontology/infer_data_ontology.py
@@ -56,7 +56,7 @@ async def recursive_flatten(self, items: Union[List[Dict[str, Any]], Dict[str, A
             for item in items:
                 flat_list.extend(await self.recursive_flatten(item, parent_id))
         elif isinstance(items, dict):
-            model = NodeModel.parse_obj(items)
+            model = NodeModel.model_validate(items)
             flat_list.append(await self.flatten_model(model, parent_id))
             for child in model.children:
                 flat_list.extend(await self.recursive_flatten(child, model.node_id))

diff --git a/cognee/tasks/save_chunks_to_store/save_chunks_to_store.py b/cognee/tasks/save_chunks_to_store/save_chunks_to_store.py
@@ -11,7 +11,7 @@ async def save_chunks_to_store(data_chunks: list[DocumentChunk], collection_name
 
     # Remove and unlink existing chunks
     if await vector_engine.has_collection(collection_name):
-        existing_chunks = [DocumentChunk.parse_obj(chunk.payload) for chunk in (await vector_engine.retrieve(
+        existing_chunks = [DocumentChunk.model_validate(chunk.payload) for chunk in (await vector_engine.retrieve(
             collection_name,
             [str(chunk.chunk_id) for chunk in data_chunks],
         ))]