From 74bc7c942054c3203ad20639ef3c08c7d59aea5f Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Mon, 29 Sep 2025 21:22:21 +0200 Subject: [PATCH 1/2] refactor: set node_set to None for endpoint --- cognee/api/v1/add/routers/get_add_router.py | 6 +++++- cognee/api/v1/update/routers/get_update_router.py | 2 +- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/cognee/api/v1/add/routers/get_add_router.py b/cognee/api/v1/add/routers/get_add_router.py index dfa7d275b..4d0063cc9 100644 --- a/cognee/api/v1/add/routers/get_add_router.py +++ b/cognee/api/v1/add/routers/get_add_router.py @@ -73,7 +73,11 @@ def get_add_router() -> APIRouter: try: add_run = await cognee_add( - data, datasetName, user=user, dataset_id=datasetId, node_set=node_set + data, + datasetName, + user=user, + dataset_id=datasetId, + node_set=node_set if node_set else None, ) if isinstance(add_run, PipelineRunErrored): diff --git a/cognee/api/v1/update/routers/get_update_router.py b/cognee/api/v1/update/routers/get_update_router.py index 144bdd9e6..4101e1e31 100644 --- a/cognee/api/v1/update/routers/get_update_router.py +++ b/cognee/api/v1/update/routers/get_update_router.py @@ -75,7 +75,7 @@ def get_update_router() -> APIRouter: data=data, dataset_id=dataset_id, user=user, - node_set=node_set, + node_set=node_set if node_set else None, ) # If any cognify run errored return JSONResponse with proper error status code From 7ab000d891fa6ed6fbe6804050456ea55fad28ce Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Tue, 30 Sep 2025 18:12:22 +0200 Subject: [PATCH 2/2] refactor: Add test for updating of docs and visualization --- cognee/__init__.py | 1 + cognee/api/v1/update/update.py | 2 +- cognee/tests/test_library.py | 32 +++++++++++++++++++++++++++++++- 3 files changed, 33 insertions(+), 2 deletions(-) diff --git a/cognee/__init__.py b/cognee/__init__.py index 27dea1ad5..6e4d2a903 100644 --- a/cognee/__init__.py +++ b/cognee/__init__.py @@ -19,6 +19,7 @@ from .api.v1.add import add from .api.v1.delete import delete from .api.v1.cognify import cognify from .modules.memify import memify +from .api.v1.update import update from .api.v1.config.config import config from .api.v1.datasets.datasets import datasets from .api.v1.prune import prune diff --git a/cognee/api/v1/update/update.py b/cognee/api/v1/update/update.py index b4b1f5e5a..a421b3dc0 100644 --- a/cognee/api/v1/update/update.py +++ b/cognee/api/v1/update/update.py @@ -10,9 +10,9 @@ from cognee.api.v1.cognify import cognify async def update( data_id: UUID, data: Union[BinaryIO, list[BinaryIO], str, list[str]], + dataset_id: UUID, user: User = None, node_set: Optional[List[str]] = None, - dataset_id: Optional[UUID] = None, vector_db_config: dict = None, graph_db_config: dict = None, preferred_loaders: List[str] = None, diff --git a/cognee/tests/test_library.py b/cognee/tests/test_library.py index c5e6cc64b..33eb64f29 100755 --- a/cognee/tests/test_library.py +++ b/cognee/tests/test_library.py @@ -6,6 +6,7 @@ from cognee.modules.search.operations import get_history from cognee.modules.users.methods import get_default_user from cognee.shared.logging_utils import get_logger from cognee.modules.search.types import SearchType +from cognee import update logger = get_logger() @@ -42,7 +43,7 @@ async def main(): await cognee.add([text], dataset_name) - await cognee.cognify([dataset_name]) + cognify_run_info = await cognee.cognify([dataset_name]) from cognee.infrastructure.databases.vector import get_vector_engine @@ -77,6 +78,35 @@ async def main(): assert len(history) == 6, "Search history is not correct." + # Test updating of documents + # Get Pipeline Run object + pipeline_run_obj = list(cognify_run_info.values())[0] + for data_item in pipeline_run_obj.data_ingestion_info: + # Update all documents in dataset to only contain Mark and Cindy information + await update( + dataset_id=pipeline_run_obj.dataset_id, + data_id=data_item["data_id"], + data="Mark met with Cindy at a cafe.", + ) + + search_results = await cognee.search( + query_type=SearchType.GRAPH_COMPLETION, query_text="What information do you contain?" + ) + assert "Mark" in search_results[0], ( + "Failed to update document, no mention of Mark in search results" + ) + assert "Cindy" in search_results[0], ( + "Failed to update document, no mention of Cindy in search results" + ) + assert "Artificial intelligence" not in search_results[0], ( + "Failed to update document, Artificial intelligence still mentioned in search results" + ) + + # Test visualization + from cognee import visualize_graph + + await visualize_graph() + # Assert local data files are cleaned properly await cognee.prune.prune_data() data_root_directory = get_storage_config()["data_root_directory"]