diff --git a/api/apps/sdk/dataset.py b/api/apps/sdk/dataset.py index 3d131f60748..bd14587c13c 100644 --- a/api/apps/sdk/dataset.py +++ b/api/apps/sdk/dataset.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # + from flask import request from api.db import StatusEnum, FileSource @@ -33,7 +34,7 @@ def save(tenant_id): req = request.json e, t = TenantService.get_by_id(tenant_id) if "id" not in req: - if "tenant_id" in req or "embd_id" in req: + if "tenant_id" in req or "embedding_model" in req: return get_data_error_result( retmsg="Tenant_id or embedding_model must not be provided") if "name" not in req: @@ -47,22 +48,39 @@ def save(tenant_id): if KnowledgebaseService.query(name=req["name"], tenant_id=tenant_id, status=StatusEnum.VALID.value): return get_data_error_result( retmsg="Duplicated knowledgebase name in creating dataset.") - req["tenant_id"] = tenant_id - req['created_by'] = tenant_id - req['embd_id'] = t.embd_id + req["tenant_id"] = req['created_by'] = tenant_id + req['embedding_model'] = t.embd_id + key_mapping = { + "chunk_num": "chunk_count", + "doc_num": "document_count", + "parser_id": "parse_method", + "embd_id": "embedding_model" + } + mapped_keys = {new_key: req[old_key] for new_key, old_key in key_mapping.items() if old_key in req} + req.update(mapped_keys) if not KnowledgebaseService.save(**req): return get_data_error_result(retmsg="Create dataset error.(Database error)") - return get_json_result(data=req) + renamed_data={} + e, k = KnowledgebaseService.get_by_id(req["id"]) + for key, value in k.to_dict().items(): + new_key = key_mapping.get(key, key) + renamed_data[new_key] = value + return get_json_result(data=renamed_data) else: + invalid_keys = {"embd_id", "chunk_num", "doc_num", "parser_id"} + if any(key in req for key in invalid_keys): + return get_data_error_result(retmsg="The input parameters are invalid.") + if "tenant_id" in req: if req["tenant_id"] != tenant_id: return get_data_error_result( retmsg="Can't change tenant_id.") - if "embd_id" in req: - if req["embd_id"] != t.embd_id: + if "embedding_model" in req: + if req["embedding_model"] != t.embd_id: return get_data_error_result( retmsg="Can't change embedding_model.") + req.pop("embedding_model") if not KnowledgebaseService.query( created_by=tenant_id, id=req["id"]): @@ -72,20 +90,23 @@ def save(tenant_id): e, kb = KnowledgebaseService.get_by_id(req["id"]) - if "chunk_num" in req: - if req["chunk_num"] != kb.chunk_num: + if "chunk_count" in req: + if req["chunk_count"] != kb.chunk_num: return get_data_error_result( retmsg="Can't change chunk_count.") + req.pop("chunk_count") - if "doc_num" in req: - if req['doc_num'] != kb.doc_num: + if "document_count" in req: + if req['document_count'] != kb.doc_num: return get_data_error_result( retmsg="Can't change document_count.") + req.pop("document_count") - if "parser_id" in req: - if kb.chunk_num > 0 and req['parser_id'] != kb.parser_id: + if "parse_method" in req: + if kb.chunk_num != 0 and req['parse_method'] != kb.parser_id: return get_data_error_result( - retmsg="if chunk count is not 0, parse method is not changable.") + retmsg="If chunk count is not 0, parse method is not changable.") + req['parser_id'] = req.pop('parse_method') if "name" in req: if req["name"].lower() != kb.name.lower() \ and len(KnowledgebaseService.query(name=req["name"], tenant_id=tenant_id, @@ -103,6 +124,9 @@ def save(tenant_id): @token_required def delete(tenant_id): req = request.args + if "id" not in req: + return get_data_error_result( + retmsg="id is required") kbs = KnowledgebaseService.query( created_by=tenant_id, id=req["id"]) if not kbs: @@ -120,7 +144,7 @@ def delete(tenant_id): if not KnowledgebaseService.delete_by_id(req["id"]): return get_data_error_result( - retmsg="Delete dataset error.(Database error)") + retmsg="Delete dataset error.(Database serror)") return get_json_result(data=True) @@ -134,37 +158,63 @@ def list_datasets(tenant_id): tenants = TenantService.get_joined_tenants_by_user_id(tenant_id) kbs = KnowledgebaseService.get_by_tenant_ids( [m["tenant_id"] for m in tenants], tenant_id, page_number, items_per_page, orderby, desc) - return get_json_result(data=kbs) + renamed_list = [] + for kb in kbs: + key_mapping = { + "chunk_num": "chunk_count", + "doc_num": "document_count", + "parser_id": "parse_method", + "embd_id": "embedding_model" + } + renamed_data = {} + for key, value in kb.items(): + new_key = key_mapping.get(key, key) + renamed_data[new_key] = value + renamed_list.append(renamed_data) + return get_json_result(data=renamed_list) @manager.route('/detail', methods=['GET']) @token_required def detail(tenant_id): req = request.args + key_mapping = { + "chunk_num": "chunk_count", + "doc_num": "document_count", + "parser_id": "parse_method", + "embd_id": "embedding_model" + } + renamed_data = {} if "id" in req: id = req["id"] kb = KnowledgebaseService.query(created_by=tenant_id, id=req["id"]) if not kb: return get_json_result( - data=False, retmsg='You do not own the dataset', + data=False, retmsg='You do not own the dataset.', retcode=RetCode.OPERATING_ERROR) if "name" in req: name = req["name"] if kb[0].name != name: return get_json_result( - data=False, retmsg='You do not own the dataset', + data=False, retmsg='You do not own the dataset.', retcode=RetCode.OPERATING_ERROR) e, k = KnowledgebaseService.get_by_id(id) - return get_json_result(data=k.to_dict()) + for key, value in k.to_dict().items(): + new_key = key_mapping.get(key, key) + renamed_data[new_key] = value + return get_json_result(data=renamed_data) else: if "name" in req: name = req["name"] e, k = KnowledgebaseService.get_by_name(kb_name=name, tenant_id=tenant_id) if not e: return get_json_result( - data=False, retmsg='You do not own the dataset', + data=False, retmsg='You do not own the dataset.', retcode=RetCode.OPERATING_ERROR) - return get_json_result(data=k.to_dict()) + for key, value in k.to_dict().items(): + new_key = key_mapping.get(key, key) + renamed_data[new_key] = value + return get_json_result(data=renamed_data) else: return get_data_error_result( retmsg="At least one of `id` or `name` must be provided.") diff --git a/sdk/python/ragflow/modules/dataset.py b/sdk/python/ragflow/modules/dataset.py index 753dbaa8b71..e9f7b3afce5 100644 --- a/sdk/python/ragflow/modules/dataset.py +++ b/sdk/python/ragflow/modules/dataset.py @@ -24,14 +24,6 @@ def __init__(self, rag, res_dict): self.parse_method = "naive" self.parser_config = None for k in list(res_dict.keys()): - if k == "embd_id": - res_dict["embedding_model"] = res_dict[k] - if k == "parser_id": - res_dict['parse_method'] = res_dict[k] - if k == "doc_num": - res_dict["document_count"] = res_dict[k] - if k == "chunk_num": - res_dict["chunk_count"] = res_dict[k] if k not in self.__dict__: res_dict.pop(k) super().__init__(rag, res_dict) @@ -39,9 +31,9 @@ def __init__(self, rag, res_dict): def save(self) -> bool: 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, + "description": self.description, "language": self.language, "embedding_model": self.embedding_model, "permission": self.permission, - "doc_num": self.document_count, "chunk_num": self.chunk_count, "parser_id": self.parse_method, + "document_count": self.document_count, "chunk_count": self.chunk_count, "parse_method": self.parse_method, "parser_config": self.parser_config.to_json() }) res = res.json() diff --git a/sdk/python/ragflow/ragflow.py b/sdk/python/ragflow/ragflow.py index f7a238834af..7a114861b03 100644 --- a/sdk/python/ragflow/ragflow.py +++ b/sdk/python/ragflow/ragflow.py @@ -52,7 +52,7 @@ def create_dataset(self, name: str, avatar: str = "", description: str = "", lan 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": parse_method, + "document_count": document_count, "chunk_count": chunk_count, "parse_method": parse_method, "parser_config": parser_config } ) @@ -61,7 +61,7 @@ def create_dataset(self, name: str, avatar: str = "", description: str = "", lan return DataSet(self, res["data"]) raise Exception(res["retmsg"]) - def list_datasets(self, page: int = 1, page_size: int = 150, orderby: str = "create_time", desc: bool = True) -> \ + def list_datasets(self, page: int = 1, page_size: int = 1024, orderby: str = "create_time", desc: bool = True) -> \ List[DataSet]: res = self.get("/dataset/list", {"page": page, "page_size": page_size, "orderby": orderby, "desc": desc}) res = res.json()