diff --git a/memos/crud.py b/memos/crud.py index 7b3154f..67f0fb5 100644 --- a/memos/crud.py +++ b/memos/crud.py @@ -67,6 +67,15 @@ def get_entity_by_filepath(filepath: str, db: Session) -> Entity | None: return db.query(EntityModel).filter(EntityModel.filepath == filepath).first() +def remove_entity(entity_id: int, db: Session): + entity = db.query(EntityModel).filter(EntityModel.id == entity_id).first() + if entity: + db.delete(entity) + db.commit() + else: + raise ValueError(f"Entity with id {entity_id} not found") + + def create_plugin(newPlugin: NewPluginParam, db: Session) -> Plugin: db_plugin = PluginModel(**newPlugin.model_dump(mode='json')) db.add(db_plugin) diff --git a/memos/server.py b/memos/server.py index 4457af2..b51f753 100644 --- a/memos/server.py +++ b/memos/server.py @@ -141,6 +141,15 @@ def update_entity( return entity +@app.delete("/libraries/{library_id}/entities/{entity_id}", status_code=status.HTTP_204_NO_CONTENT) +def remove_entity(library_id: int, entity_id: int, db: Session = Depends(get_db)): + entity = crud.get_entity_by_id(entity_id, db) + if entity is None or entity.library_id != library_id: + raise HTTPException(status_code=404, detail="Entity not found in the specified library") + + crud.remove_entity(entity_id, db) + + @app.post("/plugins", response_model=Plugin) def new_plugin(new_plugin: NewPluginParam, db: Session = Depends(get_db)): plugin = crud.create_plugin(new_plugin, db) diff --git a/memos/test_server.py b/memos/test_server.py index f30d242..b87a25c 100644 --- a/memos/test_server.py +++ b/memos/test_server.py @@ -314,3 +314,44 @@ def test_list_entities_in_folder(client): invalid_list_response = client.get(f"/libraries/9999/folders/{folder_id}/entities") assert invalid_list_response.status_code == 404 assert invalid_list_response.json() == {"detail": "Library not found"} + + +def test_remove_entity(client): + # Create a new library + new_library = NewLibraryParam(name="Test Library") + library_response = client.post("/libraries", json=new_library.model_dump(mode="json")) + library_id = library_response.json()["id"] + + # Create a new folder in the library + new_folder = NewFolderParam(path="/tmp") + folder_response = client.post(f"/libraries/{library_id}/folders", json=new_folder.model_dump(mode="json")) + folder_id = folder_response.json()["id"] + + # Create a new entity to be deleted + new_entity = NewEntityParam( + filename="test_delete.txt", + filepath="test_delete.txt", + size=100, + file_created_at="2023-01-01T00:00:00", + file_last_modified_at="2023-01-01T00:00:00", + file_type="text/plain", + folder_id=folder_id, + ) + entity_response = client.post( + f"/libraries/{library_id}/entities", json=new_entity.model_dump(mode="json") + ) + entity_id = entity_response.json()["id"] + + # Delete the entity + delete_response = client.delete(f"/libraries/{library_id}/entities/{entity_id}") + assert delete_response.status_code == 204 + + # Verify the entity is deleted + get_response = client.get(f"/libraries/{library_id}/entities/{entity_id}") + assert get_response.status_code == 404 + assert get_response.json() == {"detail": "Entity not found"} + + # Test for entity not found in the specified library + invalid_delete_response = client.delete(f"/libraries/{library_id}/entities/9999") + assert invalid_delete_response.status_code == 404 + assert invalid_delete_response.json() == {"detail": "Entity not found in the specified library"}