pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

created get_dataset, update_dataset API and fixed: delete (#1201)

Browse Source 
### What problem does this PR solve?

Added get_dataset and update_dataset API.
Fixed delete_dataset.

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] New Feature (non-breaking change which adds functionality)
- [x] Documentation Update
This commit is contained in:
cecilia-uu
2024-06-19 18:01:38 +08:00
committed by GitHub
parent 0b90aab22c
commit 5ec19b5f53
5 changed files with 280 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
sdk/python/ragflow/ragflow.py
View File

@ -17,11 +17,9 @@ import os
import requests
import json
from httpx import HTTPError
class RAGFlow:
def __init__(self, user_key, base_url, version = 'v1'):
def __init__(self, user_key, base_url, version='v1'):
'''
api_url: http://<host_address>/api/v1
dataset_url: http://<host_address>/api/v1/dataset
@ -41,14 +39,10 @@ class RAGFlow:
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}"}
endpoint = f"{self.dataset_url}/{dataset_id}"
res = requests.delete(endpoint, headers=self.authorization_header)
return res.json()
def find_dataset_id_by_name(self, dataset_name):
res = requests.get(self.dataset_url, headers=self.authorization_header)
@ -64,42 +58,18 @@ class RAGFlow:
"orderby": orderby,
"desc": desc
}
try:
response = requests.get(url=self.dataset_url, params=params, headers=self.authorization_header)
response.raise_for_status() # if it is not 200
original_data = response.json()
# TODO: format the data
# print(original_data)
# # Process the original data into the desired format
# formatted_data = {
# "datasets": [
# {
# "id": dataset["id"],
# "created": dataset["create_time"], # Adjust the key based on the actual response
# "fileCount": dataset["doc_num"], # Adjust the key based on the actual response
# "name": dataset["name"]
# }
# for dataset in original_data
# ]
# }
return response.status_code, original_data
except HTTPError as http_err:
print(f"HTTP error occurred: {http_err}")
except Exception as err:
print(f"An error occurred: {err}")
response = requests.get(url=self.dataset_url, params=params, headers=self.authorization_header)
return response.json()
def get_dataset(self, dataset_id):
def get_dataset(self, dataset_name):
dataset_id = self.find_dataset_id_by_name(dataset_name)
endpoint = f"{self.dataset_url}/{dataset_id}"
response = requests.get(endpoint)
if response.status_code == 200:
return response.json()
else:
return None
response = requests.get(endpoint, headers=self.authorization_header)
return response.json()
def update_dataset(self, dataset_name, **params):
dataset_id = self.find_dataset_id_by_name(dataset_name)
def update_dataset(self, dataset_id, params):
endpoint = f"{self.dataset_url}/{dataset_id}"
response = requests.put(endpoint, json=params)
if response.status_code == 200:
return True
else:
return False
response = requests.put(endpoint, json=params, headers=self.authorization_header)
return response.json()

2
sdk/python/test/common.py
View File

@ -1,4 +1,4 @@
API_KEY = 'ImFmNWQ3YTY0Mjg5NjExZWZhNTdjMzA0M2Q3ZWU1MzdlIg.ZmldwA.9oP9pVtuEQSpg-Z18A2eOkWO-3E'
API_KEY = 'ImFhMmJhZmUwMmQxNzExZWZhZDdmMzA0M2Q3ZWU1MzdlIg.ZnDsIQ.u-0-_qCRU6a4WICxyAPsjaafyOo'
HOST_ADDRESS = 'http://127.0.0.1:9380'

203
sdk/python/test/test_dataset.py
View File

@ -1,3 +1,4 @@
from api.settings import RetCode
from test_sdkbase import TestSdk
from ragflow import RAGFlow
import pytest
@ -15,6 +16,19 @@ class TestDataset(TestSdk):
4. update the kb
5. delete the kb
"""
def setup_method(self):
"""
Delete all the datasets.
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
listed_data = ragflow.list_dataset()
listed_data = listed_data['data']
listed_names = {d['name'] for d in listed_data}
for name in listed_names:
ragflow.delete_dataset(name)
# -----------------------create_dataset---------------------------------
def test_create_dataset_with_success(self):
"""
@ -23,7 +37,7 @@ class TestDataset(TestSdk):
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
# create a kb
res = ragflow.create_dataset("kb1")
assert res['code'] == 0 and res['message'] == 'success'
assert res['code'] == RetCode.SUCCESS and res['message'] == 'success'
def test_create_dataset_with_empty_name(self):
"""
@ -31,7 +45,7 @@ class TestDataset(TestSdk):
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
res = ragflow.create_dataset("")
assert res['message'] == 'Empty dataset name' and res['code'] == 102
assert res['message'] == 'Empty dataset name' and res['code'] == RetCode.DATA_ERROR
def test_create_dataset_with_name_exceeding_limit(self):
"""
@ -41,7 +55,7 @@ class TestDataset(TestSdk):
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
res = ragflow.create_dataset(name)
assert (res['message'] == f"Dataset name: {name} with length {len(name)} exceeds {NAME_LENGTH_LIMIT}!"
and res['code'] == 102)
and res['code'] == RetCode.DATA_ERROR)
def test_create_dataset_name_with_space_in_the_middle(self):
"""
@ -50,7 +64,7 @@ class TestDataset(TestSdk):
name = "k b"
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
res = ragflow.create_dataset(name)
assert (res['code'] == 0 and res['message'] == 'success')
assert (res['code'] == RetCode.SUCCESS and res['message'] == 'success')
def test_create_dataset_name_with_space_in_the_head(self):
"""
@ -59,7 +73,7 @@ class TestDataset(TestSdk):
name = " kb"
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
res = ragflow.create_dataset(name)
assert (res['code'] == 0 and res['message'] == 'success')
assert (res['code'] == RetCode.SUCCESS and res['message'] == 'success')
def test_create_dataset_name_with_space_in_the_tail(self):
"""
@ -68,7 +82,7 @@ class TestDataset(TestSdk):
name = "kb "
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
res = ragflow.create_dataset(name)
assert (res['code'] == 0 and res['message'] == 'success')
assert (res['code'] == RetCode.SUCCESS and res['message'] == 'success')
def test_create_dataset_name_with_space_in_the_head_and_tail_and_length_exceed_limit(self):
"""
@ -78,7 +92,7 @@ class TestDataset(TestSdk):
name = " " + "k" * NAME_LENGTH_LIMIT + " "
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
res = ragflow.create_dataset(name)
assert (res['code'] == 0 and res['message'] == 'success')
assert (res['code'] == RetCode.SUCCESS and res['message'] == 'success')
def test_create_dataset_with_two_same_name(self):
"""
@ -86,9 +100,9 @@ class TestDataset(TestSdk):
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
res = ragflow.create_dataset("kb")
assert (res['code'] == 0 and res['message'] == 'success')
assert (res['code'] == RetCode.SUCCESS and res['message'] == 'success')
res = ragflow.create_dataset("kb")
assert (res['code'] == 0 and res['message'] == 'success')
assert (res['code'] == RetCode.SUCCESS and res['message'] == 'success')
def test_create_dataset_with_only_space_in_the_name(self):
"""
@ -96,7 +110,7 @@ class TestDataset(TestSdk):
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
res = ragflow.create_dataset(" ")
assert (res['code'] == 0 and res['message'] == 'success')
assert (res['code'] == RetCode.SUCCESS and res['message'] == 'success')
def test_create_dataset_with_space_number_exceeding_limit(self):
"""
@ -105,7 +119,7 @@ class TestDataset(TestSdk):
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
name = " " * NAME_LENGTH_LIMIT
res = ragflow.create_dataset(name)
assert (res['code'] == 0 and res['message'] == 'success')
assert (res['code'] == RetCode.SUCCESS and res['message'] == 'success')
def test_create_dataset_with_name_having_return(self):
"""
@ -114,7 +128,7 @@ class TestDataset(TestSdk):
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
name = "kb\n"
res = ragflow.create_dataset(name)
assert (res['code'] == 0 and res['message'] == 'success')
assert (res['code'] == RetCode.SUCCESS and res['message'] == 'success')
def test_create_dataset_with_name_having_the_null_character(self):
"""
@ -123,7 +137,7 @@ class TestDataset(TestSdk):
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
name = "kb\0"
res = ragflow.create_dataset(name)
assert (res['code'] == 0 and res['message'] == 'success')
assert (res['code'] == RetCode.SUCCESS and res['message'] == 'success')
# -----------------------list_dataset---------------------------------
def test_list_dataset_success(self):
@ -133,10 +147,7 @@ class TestDataset(TestSdk):
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
# Call the list_datasets method
response = ragflow.list_dataset()
code, datasets = response
assert code == 200
assert response['code'] == RetCode.SUCCESS
def test_list_dataset_with_checking_size_and_name(self):
"""
@ -152,12 +163,12 @@ class TestDataset(TestSdk):
dataset_name = response['data']['dataset_name']
real_name_to_create.add(dataset_name)
status_code, listed_data = ragflow.list_dataset(0, 3)
listed_data = listed_data['data']
response = ragflow.list_dataset(0, 3)
listed_data = response['data']
listed_names = {d['name'] for d in listed_data}
assert listed_names == real_name_to_create
assert status_code == 200
assert response['code'] == RetCode.SUCCESS
assert len(listed_data) == len(datasets_to_create)
def test_list_dataset_with_getting_empty_result(self):
@ -174,12 +185,13 @@ class TestDataset(TestSdk):
dataset_name = response['data']['dataset_name']
real_name_to_create.add(dataset_name)
status_code, listed_data = ragflow.list_dataset(0, 0)
listed_data = listed_data['data']
response = ragflow.list_dataset(0, 0)
listed_data = response['data']
listed_names = {d['name'] for d in listed_data}
assert listed_names == real_name_to_create
assert status_code == 200
assert response['code'] == RetCode.SUCCESS
assert len(listed_data) == 0
def test_list_dataset_with_creating_100_knowledge_bases(self):
@ -196,12 +208,12 @@ class TestDataset(TestSdk):
dataset_name = response['data']['dataset_name']
real_name_to_create.add(dataset_name)
status_code, listed_data = ragflow.list_dataset(0, 100)
listed_data = listed_data['data']
res = ragflow.list_dataset(0, 100)
listed_data = res['data']
listed_names = {d['name'] for d in listed_data}
assert listed_names == real_name_to_create
assert status_code == 200
assert res['code'] == RetCode.SUCCESS
assert len(listed_data) == 100
def test_list_dataset_with_showing_one_dataset(self):
@ -210,9 +222,8 @@ class TestDataset(TestSdk):
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
response = ragflow.list_dataset(0, 1)
code, response = response
datasets = response['data']
assert len(datasets) == 1
assert len(datasets) == 1 and response['code'] == RetCode.SUCCESS
def test_list_dataset_failure(self):
"""
@ -220,8 +231,7 @@ class TestDataset(TestSdk):
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
response = ragflow.list_dataset(-1, -1)
_, res = response
assert "IndexError" in res['message']
assert "IndexError" in response['message'] and response['code'] == RetCode.EXCEPTION_ERROR
def test_list_dataset_for_empty_datasets(self):
"""
@ -229,9 +239,8 @@ class TestDataset(TestSdk):
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
response = ragflow.list_dataset()
code, response = response
datasets = response['data']
assert len(datasets) == 0
assert len(datasets) == 0 and response['code'] == RetCode.SUCCESS
# TODO: have to set the limitation of the number of datasets
@ -245,8 +254,8 @@ class TestDataset(TestSdk):
res = ragflow.create_dataset("kb0")
real_dataset_name = res['data']['dataset_name']
# delete this dataset
result = ragflow.delete_dataset(real_dataset_name)
assert result["success"] is True
res = ragflow.delete_dataset(real_dataset_name)
assert res['code'] == RetCode.SUCCESS and 'successfully' in res['message']
def test_delete_dataset_with_not_existing_dataset(self):
"""
@ -254,7 +263,7 @@ class TestDataset(TestSdk):
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
res = ragflow.delete_dataset("weird_dataset")
assert res["success"] is False
assert res['code'] == RetCode.OPERATING_ERROR and res['message'] == 'The dataset cannot be found for your current account.'
def test_delete_dataset_with_creating_100_datasets_and_deleting_100_datasets(self):
"""
@ -273,40 +282,46 @@ class TestDataset(TestSdk):
for name in real_name_to_create:
res = ragflow.delete_dataset(name)
assert res["success"] is True
assert res['code'] == RetCode.SUCCESS and 'successfully' in res['message']
def test_delete_dataset_with_space_in_the_middle_of_the_name(self):
"""
Test deleting a dataset when its name has space in the middle.
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
ragflow.create_dataset("k b")
res = ragflow.delete_dataset("k b")
print(res)
assert res["success"] is True
assert res['code'] == RetCode.SUCCESS and 'successfully' in res['message']
def test_delete_dataset_with_space_in_the_head_of_the_name(self):
"""
Test deleting a dataset when its name has space in the head.
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
ragflow.create_dataset(" kb")
res = ragflow.delete_dataset(" kb")
assert res["success"] is False
assert (res['code'] == RetCode.OPERATING_ERROR
and res['message'] == 'The dataset cannot be found for your current account.')
def test_delete_dataset_with_space_in_the_tail_of_the_name(self):
"""
Test deleting a dataset when its name has space in the tail.
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
ragflow.create_dataset("kb ")
res = ragflow.delete_dataset("kb ")
assert res["success"] is False
assert (res['code'] == RetCode.OPERATING_ERROR
and res['message'] == 'The dataset cannot be found for your current account.')
def test_delete_dataset_with_only_space_in_the_name(self):
"""
Test deleting a dataset when its name only has space.
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
ragflow.create_dataset(" ")
res = ragflow.delete_dataset(" ")
assert res["success"] is False
assert (res['code'] == RetCode.OPERATING_ERROR
and res['message'] == 'The dataset cannot be found for your current account.')
def test_delete_dataset_with_only_exceeding_limit_space_in_the_name(self):
"""
@ -314,8 +329,10 @@ class TestDataset(TestSdk):
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
name = " " * (NAME_LENGTH_LIMIT + 1)
ragflow.create_dataset(name)
res = ragflow.delete_dataset(name)
assert res["success"] is False
assert (res['code'] == RetCode.OPERATING_ERROR
and res['message'] == 'The dataset cannot be found for your current account.')
def test_delete_dataset_with_name_with_space_in_the_head_and_tail_and_length_exceed_limit(self):
"""
@ -324,10 +341,102 @@ class TestDataset(TestSdk):
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
name = " " + "k" * NAME_LENGTH_LIMIT + " "
ragflow.create_dataset(name)
res = ragflow.delete_dataset(name)
assert res["success"] is False
assert (res['code'] == RetCode.OPERATING_ERROR
and res['message'] == 'The dataset cannot be found for your current account.')
# ---------------------------------get_dataset-----------------------------------------
def test_get_dataset_with_success(self):
"""
Test getting a dataset which exists.
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
response = ragflow.create_dataset("test")
dataset_name = response['data']['dataset_name']
res = ragflow.get_dataset(dataset_name)
assert res['code'] == RetCode.SUCCESS and res['data']['name'] == dataset_name
def test_get_dataset_with_failure(self):
"""
Test getting a dataset which does not exist.
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
res = ragflow.get_dataset("weird_dataset")
assert res['code'] == RetCode.DATA_ERROR and res['message'] == "Can't find this dataset!"
# ---------------------------------update a dataset-----------------------------------
def test_update_dataset_without_existing_dataset(self):
"""
Test updating a dataset which does not exist.
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
params = {
'name': 'new_name3',
'description': 'new_description',
"permission": 'me',
"parser_id": 'naive',
"language": 'English'
}
res = ragflow.update_dataset("weird_dataset", **params)
assert (res['code'] == RetCode.OPERATING_ERROR
and res['message'] == 'Only the owner of knowledgebase is authorized for this operation!')
def test_update_dataset_with_updating_six_parameters(self):
"""
Test updating a dataset when updating six parameters.
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
ragflow.create_dataset("new_name1")
params = {
'name': 'new_name',
'description': 'new_description1',
"permission": 'me',
"parser_id": 'naive',
"language": 'English'
}
res = ragflow.update_dataset("new_name1", **params)
assert res['code'] == RetCode.SUCCESS
assert (res['data']['description'] == 'new_description1'
and res['data']['name'] == 'new_name' and res['data']['permission'] == 'me'
and res['data']['language'] == 'English' and res['data']['parser_id'] == 'naive')
def test_update_dataset_with_updating_two_parameters(self):
"""
Test updating a dataset when updating two parameters.
"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
ragflow.create_dataset("new_name2")
params = {
"name": "new_name3",
"language": 'English'
}
res = ragflow.update_dataset("new_name2", **params)
assert (res['code'] == RetCode.SUCCESS and res['data']['name'] == "new_name3"
and res['data']['language'] == 'English')
def test_update_dataset_with_updating_layout_recognize(self):
"""Test updating a dataset with only updating the layout_recognize"""
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
ragflow.create_dataset("test_update_dataset_with_updating_layout_recognize")
params = {
"layout_recognize": False
}
res = ragflow.update_dataset("test_update_dataset_with_updating_layout_recognize", **params)
assert res['code'] == RetCode.SUCCESS and res['data']['parser_config']['layout_recognize'] is False
def test_update_dataset_with_empty_parameter(self):
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
ragflow.create_dataset("test_update_dataset_with_empty_parameter")
params = {}
res = ragflow.update_dataset("test_update_dataset_with_empty_parameter", **params)
assert (res['code'] == RetCode.DATA_ERROR
and res['message'] == 'Please input at least one parameter that you want to update!')
# ---------------------------------mix the different methods--------------------------
# ---------------------------------mix the different methods--------------------
def test_create_and_delete_dataset_together(self):
"""
Test creating 1 dataset, and then deleting 1 dataset.
@ -336,11 +445,11 @@ class TestDataset(TestSdk):
# create 1 dataset
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
res = ragflow.create_dataset("ddd")
assert res['code'] == 0 and res['message'] == 'success'
assert res['code'] == RetCode.SUCCESS and res['message'] == 'success'
# delete 1 dataset
res = ragflow.delete_dataset("ddd")
assert res["success"] is True
assert res["code"] == RetCode.SUCCESS
# create 10 datasets
datasets_to_create = ["dataset1"] * 10
@ -355,5 +464,5 @@ class TestDataset(TestSdk):
# delete 10 datasets
for name in real_name_to_create:
res = ragflow.delete_dataset(name)
assert res["success"] is True
assert res["code"] == RetCode.SUCCESS