AI/ragflow
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git synced 2025-08-12 03:49:06 +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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 280 additions and 123 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

135
api/apps/dataset_api.py
View File

@ -20,6 +20,7 @@ import re
from datetime import datetime, timedelta
from flask import request, Response
from flask_login import login_required, current_user
from httpx import HTTPError
from api.db import FileType, ParserType, FileSource, StatusEnum
from api.db.db_models import APIToken, API4Conversation, Task, File
@ -45,6 +46,7 @@ from api.utils.api_utils import construct_json_result, construct_result, constru
from api.contants import NAME_LENGTH_LIMIT
# ------------------------------ create a dataset ---------------------------------------
@manager.route('/', methods=['POST'])
@login_required # use login
@validate_request("name") # check name key
@ -104,18 +106,20 @@ def create_dataset():
request_body["id"] = get_uuid()
request_body["tenant_id"] = tenant_id
request_body["created_by"] = tenant_id
e, t = TenantService.get_by_id(tenant_id)
if not e:
exist, t = TenantService.get_by_id(tenant_id)
if not exist:
return construct_result(code=RetCode.AUTHENTICATION_ERROR, message="Tenant not found.")
request_body["embd_id"] = t.embd_id
if not KnowledgebaseService.save(**request_body):
# failed to create new dataset
return construct_result()
return construct_json_result(data={"dataset_name": request_body["name"]})
return construct_json_result(code=RetCode.SUCCESS,
data={"dataset_name": request_body["name"], "dataset_id": request_body["id"]})
except Exception as e:
return construct_error_response(e)
# -----------------------------list datasets-------------------------------------------------------
@manager.route('/', methods=['GET'])
@login_required
def list_datasets():
@ -125,67 +129,140 @@ def list_datasets():
desc = request.args.get("desc", True)
try:
tenants = TenantService.get_joined_tenants_by_user_id(current_user.id)
kbs = KnowledgebaseService.get_by_tenant_ids_by_offset(
datasets = KnowledgebaseService.get_by_tenant_ids_by_offset(
[m["tenant_id"] for m in tenants], current_user.id, int(offset), int(count), orderby, desc)
return construct_json_result(data=kbs, code=RetCode.DATA_ERROR, message=f"attempt to list datasets")
return construct_json_result(data=datasets, code=RetCode.SUCCESS, message=f"List datasets successfully!")
except Exception as e:
return construct_error_response(e)
except HTTPError as http_err:
return construct_json_result(http_err)
# ---------------------------------delete a dataset ----------------------------
@manager.route('/<dataset_id>', methods=['DELETE'])
@login_required
@validate_request("dataset_id")
def 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)
datasets = KnowledgebaseService.query(created_by=current_user.id, id=dataset_id)
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)!")
# according to the id, searching for the dataset
if not datasets:
return construct_json_result(message=f'The dataset cannot be found for your current account.',
code=RetCode.OPERATING_ERROR)
# Iterating the documents inside the dataset
for doc in DocumentService.query(kb_id=dataset_id):
if not DocumentService.remove_document(doc, datasets[0].tenant_id):
# the process of deleting failed
return construct_json_result(code=RetCode.DATA_ERROR,
message="There was an error during the document removal process. "
"Please check the status of the RAGFlow server and try the removal again.")
# delete the other files
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}")
# delete the dataset
if not KnowledgebaseService.delete_by_id(dataset_id):
return construct_json_result(code=RetCode.DATA_ERROR, message="There was an error during the dataset removal process. "
"Please check the status of the RAGFlow server and try the removal again.")
# success
return construct_json_result(code=RetCode.SUCCESS, message=f"Remove dataset: {dataset_id} successfully")
except Exception as e:
return construct_error_response(e)
# ------------------------------ get details of a dataset ----------------------------------------
@manager.route('/<dataset_id>', methods=['GET'])
@login_required
@validate_request("dataset_id")
def get_dataset():
dataset_id = request.args["dataset_id"]
def get_dataset(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}")
return construct_json_result(code=RetCode.DATA_ERROR, message="Can't find this dataset!")
return construct_json_result(data=dataset, code=RetCode.SUCCESS)
except Exception as e:
return construct_json_result(e)
# ------------------------------ update a dataset --------------------------------------------
@manager.route('/<dataset_id>', methods=['PUT'])
@login_required
@validate_request("name")
def update_dataset(dataset_id):
return construct_json_result(code=RetCode.DATA_ERROR, message=f"attempt to update dataset: {dataset_id}")
req = request.json
try:
# the request cannot be empty
if not req:
return construct_json_result(code=RetCode.DATA_ERROR, message="Please input at least one parameter that "
"you want to update!")
# check whether the dataset can be found
if not KnowledgebaseService.query(created_by=current_user.id, id=dataset_id):
return construct_json_result(message=f'Only the owner of knowledgebase is authorized for this operation!',
code=RetCode.OPERATING_ERROR)
exist, dataset = KnowledgebaseService.get_by_id(dataset_id)
# check whether there is this dataset
if not exist:
return construct_json_result(code=RetCode.DATA_ERROR, message="This dataset cannot be found!")
if 'name' in req:
name = req["name"].strip()
# check whether there is duplicate name
if name.lower() != dataset.name.lower() \
and len(KnowledgebaseService.query(name=name, tenant_id=current_user.id,
status=StatusEnum.VALID.value)) > 1:
return construct_json_result(code=RetCode.DATA_ERROR, message=f"The name: {name.lower()} is already used by other "
f"datasets. Please choose a different name.")
dataset_updating_data = {}
chunk_num = req.get("chunk_num")
# modify the value of 11 parameters
# 2 parameters: embedding id and chunk method
# only if chunk_num is 0, the user can update the embedding id
if req.get('embedding_model_id'):
if chunk_num == 0:
dataset_updating_data['embd_id'] = req['embedding_model_id']
else:
construct_json_result(code=RetCode.DATA_ERROR, message="You have already parsed the document in this "
"dataset, so you cannot change the embedding "
"model.")
# only if chunk_num is 0, the user can update the chunk_method
if req.get("chunk_method"):
if chunk_num == 0:
dataset_updating_data['parser_id'] = req["chunk_method"]
else:
construct_json_result(code=RetCode.DATA_ERROR, message="You have already parsed the document "
"in this dataset, so you cannot "
"change the chunk method.")
# convert the photo parameter to avatar
if req.get("photo"):
dataset_updating_data['avatar'] = req["photo"]
# layout_recognize
if 'layout_recognize' in req:
if 'parser_config' not in dataset_updating_data:
dataset_updating_data['parser_config'] = {}
dataset_updating_data['parser_config']['layout_recognize'] = req['layout_recognize']
# TODO: updating use_raptor needs to construct a class
# 6 parameters
for key in ['name', 'language', 'description', 'permission', 'id', 'token_num']:
if key in req:
dataset_updating_data[key] = req.get(key)
# update
if not KnowledgebaseService.update_by_id(dataset.id, dataset_updating_data):
return construct_json_result(code=RetCode.OPERATING_ERROR, message="Failed to update! "
"Please check the status of RAGFlow "
"server and try again!")
exist, dataset = KnowledgebaseService.get_by_id(dataset.id)
if not exist:
return construct_json_result(code=RetCode.DATA_ERROR, message="Failed to get the dataset "
"using the dataset ID.")
return construct_json_result(data=dataset.to_json(), code=RetCode.SUCCESS)
except Exception as e:
return construct_error_response(e)

3
docs/references/ragflow_api.md
View File

@ -55,7 +55,8 @@ You are *required* to save the `data.id` value returned in the response data, wh
{
"code": 0,
"data": {
"dataset_name": "kb1"
"dataset_name": "kb1",
"dataset_id": "375e8ada2d3c11ef98f93043d7ee537e"
},
"message": "success"
}

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