diff --git a/api/apps/dataset_api.py b/api/apps/dataset_api.py index 9eec34148..94b75fe6a 100644 --- a/api/apps/dataset_api.py +++ b/api/apps/dataset_api.py @@ -135,9 +135,52 @@ def list_datasets(): @manager.route('/', methods=['DELETE']) @login_required +@validate_request("dataset_id") def remove_dataset(dataset_id): - return construct_json_result(code=RetCode.DATA_ERROR, message=f"attempt to remove dataset: {dataset_id}") + req = request.json + try: + kbs = KnowledgebaseService.query( + created_by=current_user.id, id=req["dataset_id"]) + if not kbs: + return construct_json_result( + data=False, message=f'Only owner of knowledgebase authorized for this operation.', + code=RetCode.OPERATING_ERROR) + for doc in DocumentService.query(kb_id=req["dataset_id"]): + if not DocumentService.remove_document(doc, kbs[0].tenant_id): + return construct_json_result( + message="Database error (Document removal)!") + f2d = File2DocumentService.get_by_document_id(doc.id) + FileService.filter_delete([File.source_type == FileSource.KNOWLEDGEBASE, File.id == f2d[0].file_id]) + File2DocumentService.delete_by_document_id(doc.id) + + if not KnowledgebaseService.delete_by_id(req["dataset_id"]): + return construct_json_result( + message="Database error (Knowledgebase removal)!") + return construct_json_result(code=RetCode.DATA_ERROR, message=f"attempt to remove dataset: {dataset_id}") + except Exception as e: + return construct_error_response(e) + +# ------------------------------ get details of a dataset ---------------------------------------- +@manager.route('/', methods=['GET']) +@login_required +@validate_request("dataset_id") +def get_dataset(): + dataset_id = request.args["dataset_id"] + try: + dataset = KnowledgebaseService.get_detail(dataset_id) + if not dataset: + return construct_json_result( + message="Can't find this knowledgebase!") + return construct_json_result(code=RetCode.DATA_ERROR, message=f"attempt to get detail of dataset: {dataset_id}") + except Exception as e: + return construct_json_result(e) + +# ------------------------------ update a dataset -------------------------------------------- +@manager.route('/', methods=['GET']) +@login_required +def get_dataset(dataset_id): + return construct_json_result(code=RetCode.DATA_ERROR, message=f"attempt to get detail of dataset: {dataset_id}") @manager.route('/', methods=['PUT']) @login_required @@ -146,10 +189,7 @@ def update_dataset(dataset_id): return construct_json_result(code=RetCode.DATA_ERROR, message=f"attempt to update dataset: {dataset_id}") -@manager.route('/', methods=['GET']) -@login_required -def get_dataset(dataset_id): - return construct_json_result(code=RetCode.DATA_ERROR, message=f"attempt to get detail of dataset: {dataset_id}") + diff --git a/sdk/python/ragflow/ragflow.py b/sdk/python/ragflow/ragflow.py index 2deb69ab2..8f2871037 100644 --- a/sdk/python/ragflow/ragflow.py +++ b/sdk/python/ragflow/ragflow.py @@ -39,8 +39,23 @@ class RAGFlow: result_dict = json.loads(res.text) return result_dict - def delete_dataset(self, dataset_name=None, dataset_id=None): - return dataset_name + def delete_dataset(self, dataset_name): + dataset_id = self.find_dataset_id_by_name(dataset_name) + if not dataset_id: + return {"success": False, "message": "Dataset not found."} + + res = requests.delete(f"{self.dataset_url}/{dataset_id}", headers=self.authorization_header) + if res.status_code == 200: + return {"success": True, "message": "Dataset deleted successfully!"} + else: + return {"success": False, "message": f"Other status code: {res.status_code}"} + + def find_dataset_id_by_name(self, dataset_name): + res = requests.get(self.dataset_url, headers=self.authorization_header) + for dataset in res.json()['data']: + if dataset['name'] == dataset_name: + return dataset['id'] + return None def list_dataset(self, offset=0, count=-1, orderby="create_time", desc=True): params = { diff --git a/sdk/python/test/test_dataset.py b/sdk/python/test/test_dataset.py index 04e168a75..315bf7ac6 100644 --- a/sdk/python/test/test_dataset.py +++ b/sdk/python/test/test_dataset.py @@ -101,6 +101,21 @@ class TestDataset(TestSdk): _, res = response assert "IndexError" in res['message'] + def test_delete_one_dataset_with_success(self): + # get the real name of the created dataset + ragflow = RAGFlow(API_KEY, HOST_ADDRESS) + res = ragflow.create_dataset("kb0") + real_dataset_name = res['data']['dataset_name'] + print("name", real_dataset_name) + # delete this dataset + result = ragflow.delete_dataset(real_dataset_name) + print(result) + assert result["success"] is True + + def test_delete_dataset_with_not_existing_dataset(self): + ragflow = RAGFlow(API_KEY, HOST_ADDRESS) + res = ragflow.delete_dataset("weird_dataset") + assert res["success"] is False