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

create and update dataset (#2110)

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

Added the ability to create and update dataset for SDK

### Type of change


- [x] New Feature (non-breaking change which adds functionality)

---------

Co-authored-by: root <root@xwg>
Co-authored-by: Kevin Hu <kevinhu.sh@gmail.com>
This commit is contained in:
LiuHua
2024-08-27 15:23:50 +08:00
committed by GitHub
parent 61d2a74b25
commit 0bea7f21ae
6 changed files with 149 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
sdk/python/ragflow/__init__.py
View File

@ -3,3 +3,4 @@ import importlib.metadata
__version__ = importlib.metadata.version("ragflow")
from .ragflow import RAGFlow
from .modules.dataset import DataSet

21
sdk/python/ragflow/modules/dataset.py
View File

@ -2,7 +2,7 @@ from .base import Base
class DataSet(Base):
class ParseConfig(Base):
class ParserConfig(Base):
def __init__(self, rag, res_dict):
self.chunk_token_count = 128
self.layout_recognize = True
@ -21,13 +21,18 @@ class DataSet(Base):
self.permission = "me"
self.document_count = 0
self.chunk_count = 0
self.parse_method = 0
self.parser_method = "naive"
self.parser_config = None
super().__init__(rag, res_dict)
def delete(self):
try:
self.post("/rm", {"kb_id": self.id})
return True
except Exception:
return False
def save(self):
res = self.post('/dataset/save',
{"id": self.id, "name": self.name, "avatar": self.avatar, "tenant_id": self.tenant_id,
"description": self.description, "language": self.language, "embd_id": self.embedding_model,
"permission": self.permission,
"doc_num": self.document_count, "chunk_num": self.chunk_count, "parser_id": self.parser_method,
"parser_config": self.parser_config.to_json()
})
res = res.json()
if not res.get("retmsg"): return True
raise Exception(res["retmsg"])

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

@ -21,180 +21,34 @@ from .modules.dataset import DataSet
class RAGFlow:
def __init__(self, user_key, base_url, version='v1'):
"""
api_url: http://<host_address>/v1
dataset_url: http://<host_address>/v1/kb
document_url: http://<host_address>/v1/dataset/{dataset_id}/documents
api_url: http://<host_address>/api/v1
"""
self.user_key = user_key
self.api_url = f"{base_url}/{version}"
self.dataset_url = f"{self.api_url}/kb"
self.authorization_header = {"Authorization": "{}".format(self.user_key)}
self.base_url = base_url
self.api_url = f"{base_url}/api/{version}"
self.authorization_header = {"Authorization": "{} {}".format("Bearer",self.user_key)}
def post(self, path, param):
res = requests.post(url=self.dataset_url + path, json=param, headers=self.authorization_header)
res = requests.post(url=self.api_url + path, json=param, headers=self.authorization_header)
return res
def get(self, path, params=''):
res = requests.get(self.dataset_url + path, params=params, headers=self.authorization_header)
res = requests.get(self.api_url + path, params=params, headers=self.authorization_header)
return res
def create_dataset(self, dataset_name):
"""
name: dataset name
"""
res_create = self.post("/create", {"name": dataset_name})
res_create_data = res_create.json()['data']
res_detail = self.get("/detail", {"kb_id": res_create_data["kb_id"]})
res_detail_data = res_detail.json()['data']
result = {}
result['id'] = res_detail_data['id']
result['name'] = res_detail_data['name']
result['avatar'] = res_detail_data['avatar']
result['description'] = res_detail_data['description']
result['language'] = res_detail_data['language']
result['embedding_model'] = res_detail_data['embd_id']
result['permission'] = res_detail_data['permission']
result['document_count'] = res_detail_data['doc_num']
result['chunk_count'] = res_detail_data['chunk_num']
result['parser_config'] = res_detail_data['parser_config']
dataset = DataSet(self, result)
return dataset
def create_dataset(self, name:str,avatar:str="",description:str="",language:str="English",permission:str="me",
document_count:int=0,chunk_count:int=0,parser_method:str="naive",
parser_config:DataSet.ParserConfig=None):
if parser_config is None:
parser_config = DataSet.ParserConfig(self, {"chunk_token_count":128,"layout_recognize": True, "delimiter":"\n!?。；！？","task_page_size":12})
parser_config=parser_config.to_json()
res=self.post("/dataset/save",{"name":name,"avatar":avatar,"description":description,"language":language,"permission":permission,
"doc_num": document_count,"chunk_num":chunk_count,"parser_id":parser_method,
"parser_config":parser_config
}
)
res = res.json()
if not res.get("retmsg"):
return DataSet(self, res["data"])
raise Exception(res["retmsg"])
"""
def delete_dataset(self, dataset_name):
dataset_id = self.find_dataset_id_by_name(dataset_name)
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)
for dataset in res.json()["data"]:
if dataset["name"] == dataset_name:
return dataset["id"]
return None
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, headers=self.authorization_header)
return response.json()
def update_dataset(self, dataset_name, **params):
dataset_id = self.find_dataset_id_by_name(dataset_name)
endpoint = f"{self.dataset_url}/{dataset_id}"
response = requests.put(endpoint, json=params, headers=self.authorization_header)
return response.json()
# ------------------------------- CONTENT MANAGEMENT -----------------------------------------------------
# ----------------------------upload local files-----------------------------------------------------
def upload_local_file(self, dataset_id, file_paths):
files = []
for file_path in file_paths:
if not isinstance(file_path, str):
return {"code": RetCode.ARGUMENT_ERROR, "message": f"{file_path} is not string."}
if "http" in file_path:
return {"code": RetCode.ARGUMENT_ERROR, "message": "Remote files have not unsupported."}
if os.path.isfile(file_path):
files.append(("file", open(file_path, "rb")))
else:
return {"code": RetCode.DATA_ERROR, "message": f"The file {file_path} does not exist"}
res = requests.request("POST", url=f"{self.dataset_url}/{dataset_id}/documents", files=files,
headers=self.authorization_header)
result_dict = json.loads(res.text)
return result_dict
# ----------------------------delete a file-----------------------------------------------------
def delete_files(self, document_id, dataset_id):
endpoint = f"{self.dataset_url}/{dataset_id}/documents/{document_id}"
res = requests.delete(endpoint, headers=self.authorization_header)
return res.json()
# ----------------------------list files-----------------------------------------------------
def list_files(self, dataset_id, offset=0, count=-1, order_by="create_time", descend=True, keywords=""):
params = {
"offset": offset,
"count": count,
"order_by": order_by,
"descend": descend,
"keywords": keywords
}
endpoint = f"{self.dataset_url}/{dataset_id}/documents/"
res = requests.get(endpoint, params=params, headers=self.authorization_header)
return res.json()
# ----------------------------update files: enable, rename, template_type-------------------------------------------
def update_file(self, dataset_id, document_id, **params):
endpoint = f"{self.dataset_url}/{dataset_id}/documents/{document_id}"
response = requests.put(endpoint, json=params, headers=self.authorization_header)
return response.json()
# ----------------------------download a file-----------------------------------------------------
def download_file(self, dataset_id, document_id):
endpoint = f"{self.dataset_url}/{dataset_id}/documents/{document_id}"
res = requests.get(endpoint, headers=self.authorization_header)
content = res.content # binary data
# decode the binary data
try:
decoded_content = content.decode("utf-8")
json_data = json.loads(decoded_content)
return json_data # message
except json.JSONDecodeError: # binary data
_, document = DocumentService.get_by_id(document_id)
file_path = os.path.join(os.getcwd(), document.name)
with open(file_path, "wb") as file:
file.write(content)
return {"code": RetCode.SUCCESS, "data": content}
# ----------------------------start parsing-----------------------------------------------------
def start_parsing_document(self, dataset_id, document_id):
endpoint = f"{self.dataset_url}/{dataset_id}/documents/{document_id}/status"
res = requests.post(endpoint, headers=self.authorization_header)
return res.json()
def start_parsing_documents(self, dataset_id, doc_ids=None):
endpoint = f"{self.dataset_url}/{dataset_id}/documents/status"
res = requests.post(endpoint, headers=self.authorization_header, json={"doc_ids": doc_ids})
return res.json()
# ----------------------------stop parsing-----------------------------------------------------
def stop_parsing_document(self, dataset_id, document_id):
endpoint = f"{self.dataset_url}/{dataset_id}/documents/{document_id}/status"
res = requests.delete(endpoint, headers=self.authorization_header)
return res.json()
def stop_parsing_documents(self, dataset_id, doc_ids=None):
endpoint = f"{self.dataset_url}/{dataset_id}/documents/status"
res = requests.delete(endpoint, headers=self.authorization_header, json={"doc_ids": doc_ids})
return res.json()
# ----------------------------show the status of the file-----------------------------------------------------
def show_parsing_status(self, dataset_id, document_id):
endpoint = f"{self.dataset_url}/{dataset_id}/documents/{document_id}/status"
res = requests.get(endpoint, headers=self.authorization_header)
return res.json()
# ----------------------------list the chunks of the file-----------------------------------------------------
# ----------------------------delete the chunk-----------------------------------------------------
# ----------------------------edit the status of the chunk-----------------------------------------------------
# ----------------------------insert a new chunk-----------------------------------------------------
# ----------------------------get a specific chunk-----------------------------------------------------
# ----------------------------retrieval test-----------------------------------------------------
"""

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

@ -1,4 +1,4 @@
API_KEY = 'IjUxNGM0MmM4NWY5MzExZWY5MDhhMDI0MmFjMTIwMDA2Ig.ZsWebA.mV1NKdSPPllgowiH-7vz36tMWyI'
API_KEY = 'ragflow-k0N2I1MzQwNjNhMzExZWY5ODg1MDI0Mm'
HOST_ADDRESS = 'http://127.0.0.1:9380'

27
sdk/python/test/t_dataset.py
View File

@ -1,4 +1,4 @@
from ragflow import RAGFlow
from ragflow import RAGFlow, DataSet
from common import API_KEY, HOST_ADDRESS
from test_sdkbase import TestSdk
@ -6,18 +6,27 @@ from test_sdkbase import TestSdk
class TestDataset(TestSdk):
def test_create_dataset_with_success(self):
"""
Test creating dataset with success
"""
rag = RAGFlow(API_KEY, HOST_ADDRESS)
ds = rag.create_dataset("God")
assert ds is not None, "The dataset creation failed, returned None."
assert ds.name == "God", "Dataset name does not match."
if isinstance(ds, DataSet):
assert ds.name == "God", "Name does not match."
else:
assert False, f"Failed to create dataset, error: {ds}"
def test_delete_one_file(self):
def test_update_dataset_with_success(self):
"""
Test deleting one file with success.
Test updating dataset with success.
"""
rag = RAGFlow(API_KEY, HOST_ADDRESS)
ds = rag.create_dataset("ABC")
assert ds is not None, "Failed to create dataset"
assert ds.name == "ABC", "Dataset name mismatch"
delete_result = ds.delete()
assert delete_result is True, "Failed to delete dataset"
if isinstance(ds, DataSet):
assert ds.name == "ABC", "Name does not match."
ds.name = 'DEF'
res = ds.save()
assert res is True, f"Failed to update dataset, error: {res}"
else:
assert False, f"Failed to create dataset, error: {ds}"