diff --git a/cognee/api/client.py b/cognee/api/client.py index 04c81841f..6766c12de 100644 --- a/cognee/api/client.py +++ b/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"]) diff --git a/cognee/api/v1/update/__init__.py b/cognee/api/v1/update/__init__.py new file mode 100644 index 000000000..3840db468 --- /dev/null +++ b/cognee/api/v1/update/__init__.py @@ -0,0 +1 @@ +from .update import update diff --git a/cognee/api/v1/update/routers/__init__.py b/cognee/api/v1/update/routers/__init__.py new file mode 100644 index 000000000..dbfa55603 --- /dev/null +++ b/cognee/api/v1/update/routers/__init__.py @@ -0,0 +1 @@ +from .get_update_router import get_update_router diff --git a/cognee/api/v1/update/routers/get_update_router.py b/cognee/api/v1/update/routers/get_update_router.py new file mode 100644 index 000000000..16acdf29c --- /dev/null +++ b/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 diff --git a/cognee/api/v1/update/update.py b/cognee/api/v1/update/update.py new file mode 100644 index 000000000..4de4333af --- /dev/null +++ b/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