create and update dataset (#2110)

### 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 <[email protected]>

api/apps/sdk/dataset.py

@@ -0,0 +1,96 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+from flask import request
+
+from api.db import StatusEnum
+from api.db.db_models import APIToken
+from api.db.services.knowledgebase_service import KnowledgebaseService
+from api.db.services.user_service import TenantService
+from api.settings import RetCode
+from api.utils import get_uuid
+from api.utils.api_utils import get_data_error_result
+from api.utils.api_utils import get_json_result
+
+
+@manager.route('/save', methods=['POST'])
+def save():
+    req = request.json
+    token = request.headers.get('Authorization').split()[1]
+    objs = APIToken.query(token=token)
+    if not objs:
+        return get_json_result(
+            data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
+    tenant_id = objs[0].tenant_id
+    e, t = TenantService.get_by_id(tenant_id)
+    if not e:
+        return get_data_error_result(retmsg="Tenant not found.")
+    if "id" not in req:
+        req['id'] = get_uuid()
+        req["name"] = req["name"].strip()
+        if req["name"] == "":
+            return get_data_error_result(
+                retmsg="Name is not empty")
+        if KnowledgebaseService.query(name=req["name"]):
+            return get_data_error_result(
+                retmsg="Duplicated knowledgebase name")
+        req["tenant_id"] = tenant_id
+        req['created_by'] = tenant_id
+        req['embd_id'] = t.embd_id
+        if not KnowledgebaseService.save(**req):
+            return get_data_error_result(retmsg="Data saving error")
+        req.pop('created_by')
+        keys_to_rename = {'embd_id': "embedding_model", 'parser_id': 'parser_method',
+                          'chunk_num': 'chunk_count', 'doc_num': 'document_count'}
+        for old_key,new_key in keys_to_rename.items():
+            if old_key in req:
+                req[new_key]=req.pop(old_key)
+        return get_json_result(data=req)
+    else:
+        if req["tenant_id"] != tenant_id or req["embd_id"] != t.embd_id:
+            return get_data_error_result(
+                retmsg="Can't change tenant_id or embedding_model")
+
+        e, kb = KnowledgebaseService.get_by_id(req["id"])
+        if not e:
+            return get_data_error_result(
+                retmsg="Can't find this knowledgebase!")
+
+        if not KnowledgebaseService.query(
+                created_by=tenant_id, id=req["id"]):
+            return get_json_result(
+                data=False, retmsg=f'Only owner of knowledgebase authorized for this operation.',
+                retcode=RetCode.OPERATING_ERROR)
+
+        if req["chunk_num"] != kb.chunk_num or req['doc_num'] != kb.doc_num:
+            return get_data_error_result(
+                retmsg="Can't change document_count or chunk_count ")
+
+        if kb.chunk_num > 0 and req['parser_id'] != kb.parser_id:
+            return get_data_error_result(
+                retmsg="if chunk count is not 0, parser method is not changable. ")
+
+
+        if req["name"].lower() != kb.name.lower() \
+                and len(KnowledgebaseService.query(name=req["name"], tenant_id=req['tenant_id'],
+                                                   status=StatusEnum.VALID.value)) > 0:
+            return get_data_error_result(
+                retmsg="Duplicated knowledgebase name.")
+
+        del req["id"]
+        req['created_by'] = tenant_id
+        if not KnowledgebaseService.update_by_id(kb.id, req):
+            return get_data_error_result(retmsg="Data update error ")
+        return get_json_result(data=True)

sdk/python/ragflow/__init__.py

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

sdk/python/ragflow/modules/dataset.py

@@ -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"])

sdk/python/ragflow/ragflow.py

@@ -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 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-----------------------------------------------------
+    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"])
 
-    # ----------------------------get a specific chunk-----------------------------------------------------
 
-    # ----------------------------retrieval test-----------------------------------------------------
-"""

sdk/python/test/common.py

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

sdk/python/test/t_dataset.py

@@ -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}"