Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.

dataset_api.py 33KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766
  1. #
  2. # Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. import os
  16. import pathlib
  17. import re
  18. import warnings
  19. from io import BytesIO
  20. from elasticsearch_dsl import Q
  21. from flask import request, send_file
  22. from flask_login import login_required, current_user
  23. from httpx import HTTPError
  24. from api.contants import NAME_LENGTH_LIMIT
  25. from api.db import FileType, ParserType, FileSource, TaskStatus
  26. from api.db import StatusEnum
  27. from api.db.db_models import File, Task
  28. from api.db.services import duplicate_name
  29. from api.db.services.document_service import DocumentService
  30. from api.db.services.file2document_service import File2DocumentService
  31. from api.db.services.file_service import FileService
  32. from api.db.services.knowledgebase_service import KnowledgebaseService
  33. from api.db.services.task_service import TaskService
  34. from api.db.services.user_service import TenantService
  35. from api.settings import RetCode
  36. from api.utils import get_uuid
  37. from api.utils.api_utils import construct_json_result, construct_error_response
  38. from api.utils.api_utils import construct_result, validate_request
  39. from api.utils.file_utils import filename_type, thumbnail
  40. from rag.app import book, laws, manual, naive, one, paper, presentation, qa, resume, table, picture
  41. from rag.nlp import search
  42. from rag.utils.es_conn import ELASTICSEARCH
  43. from rag.utils.minio_conn import MINIO
  44. MAXIMUM_OF_UPLOADING_FILES = 256
  45. # ------------------------------ create a dataset ---------------------------------------
  46. @manager.route("/", methods=["POST"])
  47. @login_required # use login
  48. @validate_request("name") # check name key
  49. def create_dataset():
  50. # Check if Authorization header is present
  51. authorization_token = request.headers.get("Authorization")
  52. if not authorization_token:
  53. return construct_json_result(code=RetCode.AUTHENTICATION_ERROR, message="Authorization header is missing.")
  54. # TODO: Login or API key
  55. # objs = APIToken.query(token=authorization_token)
  56. #
  57. # # Authorization error
  58. # if not objs:
  59. # return construct_json_result(code=RetCode.AUTHENTICATION_ERROR, message="Token is invalid.")
  60. #
  61. # tenant_id = objs[0].tenant_id
  62. tenant_id = current_user.id
  63. request_body = request.json
  64. # In case that there's no name
  65. if "name" not in request_body:
  66. return construct_json_result(code=RetCode.DATA_ERROR, message="Expected 'name' field in request body")
  67. dataset_name = request_body["name"]
  68. # empty dataset_name
  69. if not dataset_name:
  70. return construct_json_result(code=RetCode.DATA_ERROR, message="Empty dataset name")
  71. # In case that there's space in the head or the tail
  72. dataset_name = dataset_name.strip()
  73. # In case that the length of the name exceeds the limit
  74. dataset_name_length = len(dataset_name)
  75. if dataset_name_length > NAME_LENGTH_LIMIT:
  76. return construct_json_result(
  77. code=RetCode.DATA_ERROR,
  78. message=f"Dataset name: {dataset_name} with length {dataset_name_length} exceeds {NAME_LENGTH_LIMIT}!")
  79. # In case that there are other fields in the data-binary
  80. if len(request_body.keys()) > 1:
  81. name_list = []
  82. for key_name in request_body.keys():
  83. if key_name != "name":
  84. name_list.append(key_name)
  85. return construct_json_result(code=RetCode.DATA_ERROR,
  86. message=f"fields: {name_list}, are not allowed in request body.")
  87. # If there is a duplicate name, it will modify it to make it unique
  88. request_body["name"] = duplicate_name(
  89. KnowledgebaseService.query,
  90. name=dataset_name,
  91. tenant_id=tenant_id,
  92. status=StatusEnum.VALID.value)
  93. try:
  94. request_body["id"] = get_uuid()
  95. request_body["tenant_id"] = tenant_id
  96. request_body["created_by"] = tenant_id
  97. exist, t = TenantService.get_by_id(tenant_id)
  98. if not exist:
  99. return construct_result(code=RetCode.AUTHENTICATION_ERROR, message="Tenant not found.")
  100. request_body["embd_id"] = t.embd_id
  101. if not KnowledgebaseService.save(**request_body):
  102. # failed to create new dataset
  103. return construct_result()
  104. return construct_json_result(code=RetCode.SUCCESS,
  105. data={"dataset_name": request_body["name"], "dataset_id": request_body["id"]})
  106. except Exception as e:
  107. return construct_error_response(e)
  108. # -----------------------------list datasets-------------------------------------------------------
  109. @manager.route("/", methods=["GET"])
  110. @login_required
  111. def list_datasets():
  112. offset = request.args.get("offset", 0)
  113. count = request.args.get("count", -1)
  114. orderby = request.args.get("orderby", "create_time")
  115. desc = request.args.get("desc", True)
  116. try:
  117. tenants = TenantService.get_joined_tenants_by_user_id(current_user.id)
  118. datasets = KnowledgebaseService.get_by_tenant_ids_by_offset(
  119. [m["tenant_id"] for m in tenants], current_user.id, int(offset), int(count), orderby, desc)
  120. return construct_json_result(data=datasets, code=RetCode.SUCCESS, message=f"List datasets successfully!")
  121. except Exception as e:
  122. return construct_error_response(e)
  123. except HTTPError as http_err:
  124. return construct_json_result(http_err)
  125. # ---------------------------------delete a dataset ----------------------------
  126. @manager.route("/<dataset_id>", methods=["DELETE"])
  127. @login_required
  128. def remove_dataset(dataset_id):
  129. try:
  130. datasets = KnowledgebaseService.query(created_by=current_user.id, id=dataset_id)
  131. # according to the id, searching for the dataset
  132. if not datasets:
  133. return construct_json_result(message=f"The dataset cannot be found for your current account.",
  134. code=RetCode.OPERATING_ERROR)
  135. # Iterating the documents inside the dataset
  136. for doc in DocumentService.query(kb_id=dataset_id):
  137. if not DocumentService.remove_document(doc, datasets[0].tenant_id):
  138. # the process of deleting failed
  139. return construct_json_result(code=RetCode.DATA_ERROR,
  140. message="There was an error during the document removal process. "
  141. "Please check the status of the RAGFlow server and try the removal again.")
  142. # delete the other files
  143. f2d = File2DocumentService.get_by_document_id(doc.id)
  144. FileService.filter_delete([File.source_type == FileSource.KNOWLEDGEBASE, File.id == f2d[0].file_id])
  145. File2DocumentService.delete_by_document_id(doc.id)
  146. # delete the dataset
  147. if not KnowledgebaseService.delete_by_id(dataset_id):
  148. return construct_json_result(code=RetCode.DATA_ERROR,
  149. message="There was an error during the dataset removal process. "
  150. "Please check the status of the RAGFlow server and try the removal again.")
  151. # success
  152. return construct_json_result(code=RetCode.SUCCESS, message=f"Remove dataset: {dataset_id} successfully")
  153. except Exception as e:
  154. return construct_error_response(e)
  155. # ------------------------------ get details of a dataset ----------------------------------------
  156. @manager.route("/<dataset_id>", methods=["GET"])
  157. @login_required
  158. def get_dataset(dataset_id):
  159. try:
  160. dataset = KnowledgebaseService.get_detail(dataset_id)
  161. if not dataset:
  162. return construct_json_result(code=RetCode.DATA_ERROR, message="Can't find this dataset!")
  163. return construct_json_result(data=dataset, code=RetCode.SUCCESS)
  164. except Exception as e:
  165. return construct_json_result(e)
  166. # ------------------------------ update a dataset --------------------------------------------
  167. @manager.route("/<dataset_id>", methods=["PUT"])
  168. @login_required
  169. def update_dataset(dataset_id):
  170. req = request.json
  171. try:
  172. # the request cannot be empty
  173. if not req:
  174. return construct_json_result(code=RetCode.DATA_ERROR, message="Please input at least one parameter that "
  175. "you want to update!")
  176. # check whether the dataset can be found
  177. if not KnowledgebaseService.query(created_by=current_user.id, id=dataset_id):
  178. return construct_json_result(message=f"Only the owner of knowledgebase is authorized for this operation!",
  179. code=RetCode.OPERATING_ERROR)
  180. exist, dataset = KnowledgebaseService.get_by_id(dataset_id)
  181. # check whether there is this dataset
  182. if not exist:
  183. return construct_json_result(code=RetCode.DATA_ERROR, message="This dataset cannot be found!")
  184. if "name" in req:
  185. name = req["name"].strip()
  186. # check whether there is duplicate name
  187. if name.lower() != dataset.name.lower() \
  188. and len(KnowledgebaseService.query(name=name, tenant_id=current_user.id,
  189. status=StatusEnum.VALID.value)) > 1:
  190. return construct_json_result(code=RetCode.DATA_ERROR,
  191. message=f"The name: {name.lower()} is already used by other "
  192. f"datasets. Please choose a different name.")
  193. dataset_updating_data = {}
  194. chunk_num = req.get("chunk_num")
  195. # modify the value of 11 parameters
  196. # 2 parameters: embedding id and chunk method
  197. # only if chunk_num is 0, the user can update the embedding id
  198. if req.get("embedding_model_id"):
  199. if chunk_num == 0:
  200. dataset_updating_data["embd_id"] = req["embedding_model_id"]
  201. else:
  202. return construct_json_result(code=RetCode.DATA_ERROR, message="You have already parsed the document in this "
  203. "dataset, so you cannot change the embedding "
  204. "model.")
  205. # only if chunk_num is 0, the user can update the chunk_method
  206. if "chunk_method" in req:
  207. type_value = req["chunk_method"]
  208. if is_illegal_value_for_enum(type_value, ParserType):
  209. return construct_json_result(message=f"Illegal value {type_value} for 'chunk_method' field.",
  210. code=RetCode.DATA_ERROR)
  211. if chunk_num != 0:
  212. construct_json_result(code=RetCode.DATA_ERROR, message="You have already parsed the document "
  213. "in this dataset, so you cannot "
  214. "change the chunk method.")
  215. dataset_updating_data["parser_id"] = req["template_type"]
  216. # convert the photo parameter to avatar
  217. if req.get("photo"):
  218. dataset_updating_data["avatar"] = req["photo"]
  219. # layout_recognize
  220. if "layout_recognize" in req:
  221. if "parser_config" not in dataset_updating_data:
  222. dataset_updating_data['parser_config'] = {}
  223. dataset_updating_data['parser_config']['layout_recognize'] = req['layout_recognize']
  224. # TODO: updating use_raptor needs to construct a class
  225. # 6 parameters
  226. for key in ["name", "language", "description", "permission", "id", "token_num"]:
  227. if key in req:
  228. dataset_updating_data[key] = req.get(key)
  229. # update
  230. if not KnowledgebaseService.update_by_id(dataset.id, dataset_updating_data):
  231. return construct_json_result(code=RetCode.OPERATING_ERROR, message="Failed to update! "
  232. "Please check the status of RAGFlow "
  233. "server and try again!")
  234. exist, dataset = KnowledgebaseService.get_by_id(dataset.id)
  235. if not exist:
  236. return construct_json_result(code=RetCode.DATA_ERROR, message="Failed to get the dataset "
  237. "using the dataset ID.")
  238. return construct_json_result(data=dataset.to_json(), code=RetCode.SUCCESS)
  239. except Exception as e:
  240. return construct_error_response(e)
  241. # --------------------------------content management ----------------------------------------------
  242. # ----------------------------upload files-----------------------------------------------------
  243. @manager.route("/<dataset_id>/documents/", methods=["POST"])
  244. @login_required
  245. def upload_documents(dataset_id):
  246. # no files
  247. if not request.files:
  248. return construct_json_result(
  249. message="There is no file!", code=RetCode.ARGUMENT_ERROR)
  250. # the number of uploading files exceeds the limit
  251. file_objs = request.files.getlist("file")
  252. num_file_objs = len(file_objs)
  253. if num_file_objs > MAXIMUM_OF_UPLOADING_FILES:
  254. return construct_json_result(code=RetCode.DATA_ERROR, message=f"You try to upload {num_file_objs} files, "
  255. f"which exceeds the maximum number of uploading files: {MAXIMUM_OF_UPLOADING_FILES}")
  256. # no dataset
  257. exist, dataset = KnowledgebaseService.get_by_id(dataset_id)
  258. if not exist:
  259. return construct_json_result(message="Can't find this dataset", code=RetCode.DATA_ERROR)
  260. for file_obj in file_objs:
  261. file_name = file_obj.filename
  262. # no name
  263. if not file_name:
  264. return construct_json_result(
  265. message="There is a file without name!", code=RetCode.ARGUMENT_ERROR)
  266. # TODO: support the remote files
  267. if 'http' in file_name:
  268. return construct_json_result(code=RetCode.ARGUMENT_ERROR, message="Remote files have not unsupported.")
  269. # get the root_folder
  270. root_folder = FileService.get_root_folder(current_user.id)
  271. # get the id of the root_folder
  272. parent_file_id = root_folder["id"] # document id
  273. # this is for the new user, create '.knowledgebase' file
  274. FileService.init_knowledgebase_docs(parent_file_id, current_user.id)
  275. # go inside this folder, get the kb_root_folder
  276. kb_root_folder = FileService.get_kb_folder(current_user.id)
  277. # link the file management to the kb_folder
  278. kb_folder = FileService.new_a_file_from_kb(dataset.tenant_id, dataset.name, kb_root_folder["id"])
  279. # grab all the errs
  280. err = []
  281. MAX_FILE_NUM_PER_USER = int(os.environ.get("MAX_FILE_NUM_PER_USER", 0))
  282. uploaded_docs_json = []
  283. for file in file_objs:
  284. try:
  285. # TODO: get this value from the database as some tenants have this limit while others don't
  286. if MAX_FILE_NUM_PER_USER > 0 and DocumentService.get_doc_count(dataset.tenant_id) >= MAX_FILE_NUM_PER_USER:
  287. return construct_json_result(code=RetCode.DATA_ERROR,
  288. message="Exceed the maximum file number of a free user!")
  289. # deal with the duplicate name
  290. filename = duplicate_name(
  291. DocumentService.query,
  292. name=file.filename,
  293. kb_id=dataset.id)
  294. # deal with the unsupported type
  295. filetype = filename_type(filename)
  296. if filetype == FileType.OTHER.value:
  297. return construct_json_result(code=RetCode.DATA_ERROR,
  298. message="This type of file has not been supported yet!")
  299. # upload to the minio
  300. location = filename
  301. while MINIO.obj_exist(dataset_id, location):
  302. location += "_"
  303. blob = file.read()
  304. # the content is empty, raising a warning
  305. if blob == b'':
  306. warnings.warn(f"[WARNING]: The content of the file {filename} is empty.")
  307. MINIO.put(dataset_id, location, blob)
  308. doc = {
  309. "id": get_uuid(),
  310. "kb_id": dataset.id,
  311. "parser_id": dataset.parser_id,
  312. "parser_config": dataset.parser_config,
  313. "created_by": current_user.id,
  314. "type": filetype,
  315. "name": filename,
  316. "location": location,
  317. "size": len(blob),
  318. "thumbnail": thumbnail(filename, blob)
  319. }
  320. if doc["type"] == FileType.VISUAL:
  321. doc["parser_id"] = ParserType.PICTURE.value
  322. if re.search(r"\.(ppt|pptx|pages)$", filename):
  323. doc["parser_id"] = ParserType.PRESENTATION.value
  324. DocumentService.insert(doc)
  325. FileService.add_file_from_kb(doc, kb_folder["id"], dataset.tenant_id)
  326. uploaded_docs_json.append(doc)
  327. except Exception as e:
  328. err.append(file.filename + ": " + str(e))
  329. if err:
  330. # return all the errors
  331. return construct_json_result(message="\n".join(err), code=RetCode.SERVER_ERROR)
  332. # success
  333. return construct_json_result(data=uploaded_docs_json, code=RetCode.SUCCESS)
  334. # ----------------------------delete a file-----------------------------------------------------
  335. @manager.route("/<dataset_id>/documents/<document_id>", methods=["DELETE"])
  336. @login_required
  337. def delete_document(document_id, dataset_id): # string
  338. # get the root folder
  339. root_folder = FileService.get_root_folder(current_user.id)
  340. # parent file's id
  341. parent_file_id = root_folder["id"]
  342. # consider the new user
  343. FileService.init_knowledgebase_docs(parent_file_id, current_user.id)
  344. # store all the errors that may have
  345. errors = ""
  346. try:
  347. # whether there is this document
  348. exist, doc = DocumentService.get_by_id(document_id)
  349. if not exist:
  350. return construct_json_result(message=f"Document {document_id} not found!", code=RetCode.DATA_ERROR)
  351. # whether this doc is authorized by this tenant
  352. tenant_id = DocumentService.get_tenant_id(document_id)
  353. if not tenant_id:
  354. return construct_json_result(
  355. message=f"You cannot delete this document {document_id} due to the authorization"
  356. f" reason!", code=RetCode.AUTHENTICATION_ERROR)
  357. # get the doc's id and location
  358. real_dataset_id, location = File2DocumentService.get_minio_address(doc_id=document_id)
  359. if real_dataset_id != dataset_id:
  360. return construct_json_result(message=f"The document {document_id} is not in the dataset: {dataset_id}, "
  361. f"but in the dataset: {real_dataset_id}.", code=RetCode.ARGUMENT_ERROR)
  362. # there is an issue when removing
  363. if not DocumentService.remove_document(doc, tenant_id):
  364. return construct_json_result(
  365. message="There was an error during the document removal process. Please check the status of the "
  366. "RAGFlow server and try the removal again.", code=RetCode.OPERATING_ERROR)
  367. # fetch the File2Document record associated with the provided document ID.
  368. file_to_doc = File2DocumentService.get_by_document_id(document_id)
  369. # delete the associated File record.
  370. FileService.filter_delete([File.source_type == FileSource.KNOWLEDGEBASE, File.id == file_to_doc[0].file_id])
  371. # delete the File2Document record itself using the document ID. This removes the
  372. # association between the document and the file after the File record has been deleted.
  373. File2DocumentService.delete_by_document_id(document_id)
  374. # delete it from minio
  375. MINIO.rm(dataset_id, location)
  376. except Exception as e:
  377. errors += str(e)
  378. if errors:
  379. return construct_json_result(data=False, message=errors, code=RetCode.SERVER_ERROR)
  380. return construct_json_result(data=True, code=RetCode.SUCCESS)
  381. # ----------------------------list files-----------------------------------------------------
  382. @manager.route('/<dataset_id>/documents/', methods=['GET'])
  383. @login_required
  384. def list_documents(dataset_id):
  385. if not dataset_id:
  386. return construct_json_result(
  387. data=False, message="Lack of 'dataset_id'", code=RetCode.ARGUMENT_ERROR)
  388. # searching keywords
  389. keywords = request.args.get("keywords", "")
  390. offset = request.args.get("offset", 0)
  391. count = request.args.get("count", -1)
  392. order_by = request.args.get("order_by", "create_time")
  393. descend = request.args.get("descend", True)
  394. try:
  395. docs, total = DocumentService.list_documents_in_dataset(dataset_id, int(offset), int(count), order_by,
  396. descend, keywords)
  397. return construct_json_result(data={"total": total, "docs": docs}, message=RetCode.SUCCESS)
  398. except Exception as e:
  399. return construct_error_response(e)
  400. # ----------------------------update: enable rename-----------------------------------------------------
  401. @manager.route("/<dataset_id>/documents/<document_id>", methods=["PUT"])
  402. @login_required
  403. def update_document(dataset_id, document_id):
  404. req = request.json
  405. try:
  406. legal_parameters = set()
  407. legal_parameters.add("name")
  408. legal_parameters.add("enable")
  409. legal_parameters.add("template_type")
  410. for key in req.keys():
  411. if key not in legal_parameters:
  412. return construct_json_result(code=RetCode.ARGUMENT_ERROR, message=f"{key} is an illegal parameter.")
  413. # The request body cannot be empty
  414. if not req:
  415. return construct_json_result(
  416. code=RetCode.DATA_ERROR,
  417. message="Please input at least one parameter that you want to update!")
  418. # Check whether there is this dataset
  419. exist, dataset = KnowledgebaseService.get_by_id(dataset_id)
  420. if not exist:
  421. return construct_json_result(code=RetCode.DATA_ERROR, message=f"This dataset {dataset_id} cannot be found!")
  422. # The document does not exist
  423. exist, document = DocumentService.get_by_id(document_id)
  424. if not exist:
  425. return construct_json_result(message=f"This document {document_id} cannot be found!",
  426. code=RetCode.ARGUMENT_ERROR)
  427. # Deal with the different keys
  428. updating_data = {}
  429. if "name" in req:
  430. new_name = req["name"]
  431. updating_data["name"] = new_name
  432. # Check whether the new_name is suitable
  433. # 1. no name value
  434. if not new_name:
  435. return construct_json_result(code=RetCode.DATA_ERROR, message="There is no new name.")
  436. # 2. In case that there's space in the head or the tail
  437. new_name = new_name.strip()
  438. # 3. Check whether the new_name has the same extension of file as before
  439. if pathlib.Path(new_name.lower()).suffix != pathlib.Path(
  440. document.name.lower()).suffix:
  441. return construct_json_result(
  442. data=False,
  443. message="The extension of file cannot be changed",
  444. code=RetCode.ARGUMENT_ERROR)
  445. # 4. Check whether the new name has already been occupied by other file
  446. for d in DocumentService.query(name=new_name, kb_id=document.kb_id):
  447. if d.name == new_name:
  448. return construct_json_result(
  449. message="Duplicated document name in the same dataset.",
  450. code=RetCode.ARGUMENT_ERROR)
  451. if "enable" in req:
  452. enable_value = req["enable"]
  453. if is_illegal_value_for_enum(enable_value, StatusEnum):
  454. return construct_json_result(message=f"Illegal value {enable_value} for 'enable' field.",
  455. code=RetCode.DATA_ERROR)
  456. updating_data["status"] = enable_value
  457. # TODO: Chunk-method - update parameters inside the json object parser_config
  458. if "template_type" in req:
  459. type_value = req["template_type"]
  460. if is_illegal_value_for_enum(type_value, ParserType):
  461. return construct_json_result(message=f"Illegal value {type_value} for 'template_type' field.",
  462. code=RetCode.DATA_ERROR)
  463. updating_data["parser_id"] = req["template_type"]
  464. # The process of updating
  465. if not DocumentService.update_by_id(document_id, updating_data):
  466. return construct_json_result(
  467. code=RetCode.OPERATING_ERROR,
  468. message="Failed to update document in the database! "
  469. "Please check the status of RAGFlow server and try again!")
  470. # name part: file service
  471. if "name" in req:
  472. # Get file by document id
  473. file_information = File2DocumentService.get_by_document_id(document_id)
  474. if file_information:
  475. exist, file = FileService.get_by_id(file_information[0].file_id)
  476. FileService.update_by_id(file.id, {"name": req["name"]})
  477. exist, document = DocumentService.get_by_id(document_id)
  478. # Success
  479. return construct_json_result(data=document.to_json(), message="Success", code=RetCode.SUCCESS)
  480. except Exception as e:
  481. return construct_error_response(e)
  482. # Helper method to judge whether it's an illegal value
  483. def is_illegal_value_for_enum(value, enum_class):
  484. return value not in enum_class.__members__.values()
  485. # ----------------------------download a file-----------------------------------------------------
  486. @manager.route("/<dataset_id>/documents/<document_id>", methods=["GET"])
  487. @login_required
  488. def download_document(dataset_id, document_id):
  489. try:
  490. # Check whether there is this dataset
  491. exist, _ = KnowledgebaseService.get_by_id(dataset_id)
  492. if not exist:
  493. return construct_json_result(code=RetCode.DATA_ERROR,
  494. message=f"This dataset '{dataset_id}' cannot be found!")
  495. # Check whether there is this document
  496. exist, document = DocumentService.get_by_id(document_id)
  497. if not exist:
  498. return construct_json_result(message=f"This document '{document_id}' cannot be found!",
  499. code=RetCode.ARGUMENT_ERROR)
  500. # The process of downloading
  501. doc_id, doc_location = File2DocumentService.get_minio_address(doc_id=document_id) # minio address
  502. file_stream = MINIO.get(doc_id, doc_location)
  503. if not file_stream:
  504. return construct_json_result(message="This file is empty.", code=RetCode.DATA_ERROR)
  505. file = BytesIO(file_stream)
  506. # Use send_file with a proper filename and MIME type
  507. return send_file(
  508. file,
  509. as_attachment=True,
  510. download_name=document.name,
  511. mimetype='application/octet-stream' # Set a default MIME type
  512. )
  513. # Error
  514. except Exception as e:
  515. return construct_error_response(e)
  516. # ----------------------------start parsing a document-----------------------------------------------------
  517. # helper method for parsing
  518. def dummy(prog=None, msg=""):
  519. pass
  520. def doc_parse(binary, doc_name, parser_name, tenant_id):
  521. match parser_name:
  522. case "book":
  523. book.chunk(doc_name, binary=binary, callback=dummy)
  524. case "laws":
  525. laws.chunk(doc_name, binary=binary, callback=dummy)
  526. case "manual":
  527. manual.chunk(doc_name, binary=binary, callback=dummy)
  528. case "naive":
  529. # It's the mode by default, which is general in the front-end
  530. naive.chunk(doc_name, binary=binary, callback=dummy)
  531. case "one":
  532. one.chunk(doc_name, binary=binary, callback=dummy)
  533. case "paper":
  534. paper.chunk(doc_name, binary=binary, callback=dummy)
  535. case "picture":
  536. picture.chunk(doc_name, binary=binary, tenant_id=tenant_id, lang="Chinese", callback=dummy)
  537. case "presentation":
  538. presentation.chunk(doc_name, binary=binary, callback=dummy)
  539. case "qa":
  540. qa.chunk(doc_name, binary=binary, callback=dummy)
  541. case "resume":
  542. resume.chunk(doc_name, binary=binary, callback=dummy)
  543. case "table":
  544. table.chunk(doc_name, binary=binary, callback=dummy)
  545. case _:
  546. return False
  547. return True
  548. @manager.route("/<dataset_id>/documents/<document_id>/status", methods=["POST"])
  549. @login_required
  550. def parse_document(dataset_id, document_id):
  551. try:
  552. # valid dataset
  553. exist, _ = KnowledgebaseService.get_by_id(dataset_id)
  554. if not exist:
  555. return construct_json_result(code=RetCode.DATA_ERROR,
  556. message=f"This dataset '{dataset_id}' cannot be found!")
  557. message = ""
  558. res = get_message_during_parsing_document(document_id, message)
  559. if isinstance(res, str):
  560. message += res
  561. return construct_json_result(code=RetCode.SUCCESS, message=message)
  562. else:
  563. return res
  564. except Exception as e:
  565. return construct_error_response(e)
  566. # ----------------------------start parsing documents-----------------------------------------------------
  567. @manager.route("/<dataset_id>/documents/status", methods=["POST"])
  568. @login_required
  569. def parse_documents(dataset_id):
  570. doc_ids = request.json["doc_ids"]
  571. try:
  572. exist, _ = KnowledgebaseService.get_by_id(dataset_id)
  573. if not exist:
  574. return construct_json_result(code=RetCode.DATA_ERROR,
  575. message=f"This dataset '{dataset_id}' cannot be found!")
  576. def process(doc_ids):
  577. message = ""
  578. # for loop
  579. for id in doc_ids:
  580. res = get_message_during_parsing_document(id, message)
  581. if isinstance(res, str):
  582. message += res
  583. else:
  584. return res
  585. return construct_json_result(data=True, code=RetCode.SUCCESS, message=message)
  586. # two conditions
  587. if doc_ids:
  588. return process(doc_ids)
  589. else:
  590. # documents inside the dataset
  591. docs, total = DocumentService.list_documents_in_dataset(dataset_id, 0, -1, "create_time",
  592. True, "")
  593. doc_ids = [doc["id"] for doc in docs]
  594. return process(doc_ids)
  595. except Exception as e:
  596. return construct_error_response(e)
  597. # helper method for getting message or response when parsing the document
  598. def get_message_during_parsing_document(id, message):
  599. try:
  600. # Check whether there is this document
  601. exist, document = DocumentService.get_by_id(id)
  602. if not exist:
  603. return construct_json_result(message=f"This document '{id}' cannot be found!",
  604. code=RetCode.ARGUMENT_ERROR)
  605. tenant_id = DocumentService.get_tenant_id(id)
  606. if not tenant_id:
  607. return construct_json_result(message="Tenant not found!", code=RetCode.AUTHENTICATION_ERROR)
  608. info = {"run": "1", "progress": 0}
  609. info["progress_msg"] = ""
  610. info["chunk_num"] = 0
  611. info["token_num"] = 0
  612. DocumentService.update_by_id(id, info)
  613. ELASTICSEARCH.deleteByQuery(Q("match", doc_id=id), idxnm=search.index_name(tenant_id))
  614. _, doc_attributes = DocumentService.get_by_id(id)
  615. doc_attributes = doc_attributes.to_dict()
  616. doc_id = doc_attributes["id"]
  617. bucket, doc_name = File2DocumentService.get_minio_address(doc_id=doc_id)
  618. binary = MINIO.get(bucket, doc_name)
  619. parser_name = doc_attributes["parser_id"]
  620. if binary:
  621. res = doc_parse(binary, doc_name, parser_name, tenant_id)
  622. if res is False:
  623. message += f"The parser id: {parser_name} of the document {doc_id} is not supported; "
  624. else:
  625. message += f"Empty data in the document: {doc_name}; "
  626. # failed in parsing
  627. if doc_attributes["status"] == TaskStatus.FAIL.value:
  628. message += f"Failed in parsing the document: {doc_id}; "
  629. return message
  630. except Exception as e:
  631. return construct_error_response(e)
  632. # ----------------------------stop parsing-----------------------------------------------------
  633. # ----------------------------show the status of the file-----------------------------------------------------
  634. # ----------------------------list the chunks of the file-----------------------------------------------------
  635. # -- --------------------------delete the chunk-----------------------------------------------------
  636. # ----------------------------edit the status of the chunk-----------------------------------------------------
  637. # ----------------------------insert a new chunk-----------------------------------------------------
  638. # ----------------------------upload a file-----------------------------------------------------
  639. # ----------------------------get a specific chunk-----------------------------------------------------
  640. # ----------------------------retrieval test-----------------------------------------------------