feat: Add update endpoint to Cognee

cognee/api/client.py

@@ -28,6 +28,7 @@ from cognee.api.v1.add.routers import get_add_router
 from cognee.api.v1.delete.routers import get_delete_router
 from cognee.api.v1.responses.routers import get_responses_router
 from cognee.api.v1.sync.routers import get_sync_router
+from cognee.api.v1.update.routers import get_update_router
 from cognee.api.v1.users.routers import (
     get_auth_router,
     get_register_router,
@@ -263,6 +264,8 @@ app.include_router(get_visualize_router(), prefix="/api/v1/visualize", tags=["vi
 
 app.include_router(get_delete_router(), prefix="/api/v1/delete", tags=["delete"])
 
+app.include_router(get_update_router(), prefix="/api/v1/update", tags=["update"])
+
 app.include_router(get_responses_router(), prefix="/api/v1/responses", tags=["responses"])
 
 app.include_router(get_sync_router(), prefix="/api/v1/sync", tags=["sync"])

cognee/api/v1/update/__init__.py

@@ -0,0 +1 @@
+from .update import update

cognee/api/v1/update/routers/__init__.py

@@ -0,0 +1 @@
+from .get_update_router import get_update_router

cognee/api/v1/update/routers/get_update_router.py

@@ -0,0 +1,61 @@
+from fastapi.responses import JSONResponse
+from fastapi import File, UploadFile, Depends, Form
+from typing import Optional
+from fastapi import APIRouter
+from fastapi.encoders import jsonable_encoder
+from typing import List
+from uuid import UUID
+from cognee.shared.logging_utils import get_logger
+from cognee.modules.users.models import User
+from cognee.modules.users.methods import get_authenticated_user
+from cognee.shared.utils import send_telemetry
+from cognee.modules.pipelines.models.PipelineRunInfo import (
+    PipelineRunErrored,
+)
+
+logger = get_logger()
+
+
+def get_update_router() -> APIRouter:
+    router = APIRouter()
+
+    @router.patch("", response_model=None)
+    async def update(
+        data_id: UUID,
+        dataset_id: UUID,
+        data: List[UploadFile] = File(default=None),
+        node_set: Optional[List[str]] = Form(default=[""], example=[""]),
+        user: User = Depends(get_authenticated_user),
+    ):
+        send_telemetry(
+            "Update API Endpoint Invoked",
+            user.id,
+            additional_properties={
+                "endpoint": "PATCH /v1/update",
+                "dataset_id": str(dataset_id),
+                "data_id": str(data_id),
+                "node_set": str(node_set),
+            },
+        )
+
+        from cognee.api.v1.update import update as cognee_update
+
+        try:
+            update_run = await cognee_update(
+                data_id=data_id,
+                data=data,
+                dataset_id=dataset_id,
+                user=user,
+                node_set=node_set,
+            )
+
+            # If any cognify run errored return JSONResponse with proper error status code
+            if any(isinstance(v, PipelineRunErrored) for v in update_run.values()):
+                return JSONResponse(status_code=420, content=jsonable_encoder(update_run))
+            return update_run
+
+        except Exception as error:
+            logger.error(f"Error during deletion by data_id: {str(error)}")
+            return JSONResponse(status_code=409, content={"error": str(error)})
+
+    return router

cognee/api/v1/update/update.py

@@ -0,0 +1,46 @@
+from uuid import UUID
+from typing import Union, BinaryIO, List, Optional
+
+from cognee.modules.users.models import User
+from cognee.api.v1.delete import delete
+from cognee.api.v1.add import add
+from cognee.api.v1.cognify import cognify
+
+
+async def update(
+    data_id: UUID,
+    data: Union[BinaryIO, list[BinaryIO], str, list[str]],
+    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,
+    incremental_loading: bool = True,
+):
+    await delete(
+        data_id=data_id,
+        dataset_id=dataset_id,
+        user=user,
+    )
+
+    await add(
+        data=data,
+        dataset_id=dataset_id,
+        user=user,
+        node_set=node_set,
+        vector_db_config=vector_db_config,
+        graph_db_config=graph_db_config,
+        preferred_loaders=preferred_loaders,
+        incremental_loading=incremental_loading,
+    )
+
+    cognify_run = await cognify(
+        datasets=[dataset_id],
+        user=user,
+        vector_db_config=vector_db_config,
+        graph_db_config=graph_db_config,
+        incremental_loading=incremental_loading,
+    )
+
+    return cognify_run