# # 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, FileSource from api.db.db_models import File from api.db.services.document_service import DocumentService from api.db.services.file2document_service import File2DocumentService from api.db.services.file_service import FileService 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_json_result, token_required, get_data_error_result @manager.route('/save', methods=['POST']) @token_required 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: return get_data_error_result( retmsg="Tenant_id or embedding_model must not be provided") if "name" not in req: return get_data_error_result( retmsg="Name is not empty!") req['id'] = get_uuid() req["name"] = req["name"].strip() if req["name"] == "": return get_data_error_result( retmsg="Name is not empty string!") 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 if not KnowledgebaseService.save(**req): return get_data_error_result(retmsg="Create dataset error.(Database error)") return get_json_result(data=req) else: 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: return get_data_error_result( retmsg="Can't change embedding_model.") if not KnowledgebaseService.query( created_by=tenant_id, id=req["id"]): return get_json_result( data=False, retmsg='You do not own the dataset.', retcode=RetCode.OPERATING_ERROR) e, kb = KnowledgebaseService.get_by_id(req["id"]) if "chunk_num" in req: if req["chunk_num"] != kb.chunk_num: return get_data_error_result( retmsg="Can't change chunk_count.") if "doc_num" in req: if req['doc_num'] != kb.doc_num: return get_data_error_result( retmsg="Can't change document_count.") if "parser_id" in req: if kb.chunk_num > 0 and req['parser_id'] != kb.parser_id: return get_data_error_result( retmsg="if chunk count is not 0, parse method is not changable.") if "name" in req: if req["name"].lower() != kb.name.lower() \ and len(KnowledgebaseService.query(name=req["name"], tenant_id=tenant_id, status=StatusEnum.VALID.value)) > 0: return get_data_error_result( retmsg="Duplicated knowledgebase name in updating dataset.") del req["id"] if not KnowledgebaseService.update_by_id(kb.id, req): return get_data_error_result(retmsg="Update dataset error.(Database error)") return get_json_result(data=True) @manager.route('/delete', methods=['DELETE']) @token_required def delete(tenant_id): req = request.args kbs = KnowledgebaseService.query( created_by=tenant_id, id=req["id"]) if not kbs: return get_json_result( data=False, retmsg='You do not own the dataset', retcode=RetCode.OPERATING_ERROR) for doc in DocumentService.query(kb_id=req["id"]): if not DocumentService.remove_document(doc, kbs[0].tenant_id): return get_data_error_result( retmsg="Remove document error.(Database error)") 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["id"]): return get_data_error_result( retmsg="Delete dataset error.(Database error)") return get_json_result(data=True) @manager.route('/list', methods=['GET']) @token_required def list_datasets(tenant_id): page_number = int(request.args.get("page", 1)) items_per_page = int(request.args.get("page_size", 1024)) orderby = request.args.get("orderby", "create_time") desc = bool(request.args.get("desc", True)) 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) @manager.route('/detail', methods=['GET']) @token_required def detail(tenant_id): req = request.args 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', 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', retcode=RetCode.OPERATING_ERROR) e, k = KnowledgebaseService.get_by_id(id) return get_json_result(data=k.to_dict()) 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', retcode=RetCode.OPERATING_ERROR) return get_json_result(data=k.to_dict()) else: return get_data_error_result( retmsg="At least one of `id` or `name` must be provided.")