From cb205069bc6266f26d86be40baaf7be5ef3a2066 Mon Sep 17 00:00:00 2001 From: Leon Luithlen Date: Mon, 4 Nov 2024 17:20:01 +0100 Subject: [PATCH] Replace parse_obj with model_validate --- .../chunk_naive_llm_classifier/chunk_naive_llm_classifier.py | 4 ++-- cognee/tasks/infer_data_ontology/infer_data_ontology.py | 2 +- cognee/tasks/save_chunks_to_store/save_chunks_to_store.py | 2 +- 3 files changed, 4 insertions(+), 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 index b4bcb9467..4a0c8af96 100644 --- 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 @@ async def chunk_naive_llm_classifier(data_chunks: list[DocumentChunk], classific 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 @@ async def chunk_naive_llm_classifier(data_chunks: list[DocumentChunk], classific 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 index 6415eb005..b9e2acd12 100644 --- a/cognee/tasks/infer_data_ontology/infer_data_ontology.py +++ b/cognee/tasks/infer_data_ontology/infer_data_ontology.py @@ -56,7 +56,7 @@ class OntologyEngine: 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 index 435fd0208..90286cf94 100644 --- 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], ))]