You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659
  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. #
  16. import pathlib
  17. import datetime
  18. from api.db.services.dialog_service import keyword_extraction
  19. from rag.app.qa import rmPrefix, beAdoc
  20. from rag.nlp import rag_tokenizer
  21. from api.db import LLMType, ParserType
  22. from api.db.services.llm_service import TenantLLMService
  23. from api.settings import kg_retrievaler
  24. import hashlib
  25. import re
  26. from api.utils.api_utils import token_required
  27. from api.db.db_models import Task
  28. from api.db.services.task_service import TaskService, queue_tasks
  29. from api.utils.api_utils import server_error_response
  30. from api.utils.api_utils import get_result, get_error_data_result
  31. from io import BytesIO
  32. from elasticsearch_dsl import Q
  33. from flask import request, send_file
  34. from api.db import FileSource, TaskStatus, FileType
  35. from api.db.db_models import File
  36. from api.db.services.document_service import DocumentService
  37. from api.db.services.file2document_service import File2DocumentService
  38. from api.db.services.file_service import FileService
  39. from api.db.services.knowledgebase_service import KnowledgebaseService
  40. from api.settings import RetCode, retrievaler
  41. from api.utils.api_utils import construct_json_result,get_parser_config
  42. from rag.nlp import search
  43. from rag.utils import rmSpace
  44. from rag.utils.es_conn import ELASTICSEARCH
  45. from rag.utils.storage_factory import STORAGE_IMPL
  46. import os
  47. MAXIMUM_OF_UPLOADING_FILES = 256
  48. MAXIMUM_OF_UPLOADING_FILES = 256
  49. MAXIMUM_OF_UPLOADING_FILES = 256
  50. MAXIMUM_OF_UPLOADING_FILES = 256
  51. @manager.route('/dataset/<dataset_id>/document', methods=['POST'])
  52. @token_required
  53. def upload(dataset_id, tenant_id):
  54. if 'file' not in request.files:
  55. return get_error_data_result(
  56. retmsg='No file part!', retcode=RetCode.ARGUMENT_ERROR)
  57. file_objs = request.files.getlist('file')
  58. for file_obj in file_objs:
  59. if file_obj.filename == '':
  60. return get_result(
  61. retmsg='No file selected!', retcode=RetCode.ARGUMENT_ERROR)
  62. # total size
  63. total_size = 0
  64. for file_obj in file_objs:
  65. file_obj.seek(0, os.SEEK_END)
  66. total_size += file_obj.tell()
  67. file_obj.seek(0)
  68. MAX_TOTAL_FILE_SIZE=10*1024*1024
  69. if total_size > MAX_TOTAL_FILE_SIZE:
  70. return get_result(
  71. retmsg=f'Total file size exceeds 10MB limit! ({total_size / (1024 * 1024):.2f} MB)',
  72. retcode=RetCode.ARGUMENT_ERROR)
  73. e, kb = KnowledgebaseService.get_by_id(dataset_id)
  74. if not e:
  75. raise LookupError(f"Can't find the dataset with ID {dataset_id}!")
  76. err, files= FileService.upload_document(kb, file_objs, tenant_id)
  77. if err:
  78. return get_result(
  79. retmsg="\n".join(err), retcode=RetCode.SERVER_ERROR)
  80. # rename key's name
  81. renamed_doc_list = []
  82. for file in files:
  83. doc = file[0]
  84. key_mapping = {
  85. "chunk_num": "chunk_count",
  86. "kb_id": "dataset_id",
  87. "token_num": "token_count",
  88. "parser_id": "chunk_method"
  89. }
  90. renamed_doc = {}
  91. for key, value in doc.items():
  92. new_key = key_mapping.get(key, key)
  93. renamed_doc[new_key] = value
  94. renamed_doc["run"] = "UNSTART"
  95. renamed_doc_list.append(renamed_doc)
  96. return get_result(data=renamed_doc_list)
  97. @manager.route('/dataset/<dataset_id>/info/<document_id>', methods=['PUT'])
  98. @token_required
  99. def update_doc(tenant_id, dataset_id, document_id):
  100. req = request.json
  101. if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
  102. return get_error_data_result(retmsg="You don't own the dataset.")
  103. doc = DocumentService.query(kb_id=dataset_id, id=document_id)
  104. if not doc:
  105. return get_error_data_result(retmsg="The dataset doesn't own the document.")
  106. doc = doc[0]
  107. if "chunk_count" in req:
  108. if req["chunk_count"] != doc.chunk_num:
  109. return get_error_data_result(retmsg="Can't change `chunk_count`.")
  110. if "token_count" in req:
  111. if req["token_count"] != doc.token_num:
  112. return get_error_data_result(retmsg="Can't change `token_count`.")
  113. if "progress" in req:
  114. if req['progress'] != doc.progress:
  115. return get_error_data_result(retmsg="Can't change `progress`.")
  116. if "name" in req and req["name"] != doc.name:
  117. if pathlib.Path(req["name"].lower()).suffix != pathlib.Path(doc.name.lower()).suffix:
  118. return get_result(retmsg="The extension of file can't be changed", retcode=RetCode.ARGUMENT_ERROR)
  119. for d in DocumentService.query(name=req["name"], kb_id=doc.kb_id):
  120. if d.name == req["name"]:
  121. return get_error_data_result(
  122. retmsg="Duplicated document name in the same dataset.")
  123. if not DocumentService.update_by_id(
  124. document_id, {"name": req["name"]}):
  125. return get_error_data_result(
  126. retmsg="Database error (Document rename)!")
  127. informs = File2DocumentService.get_by_document_id(document_id)
  128. if informs:
  129. e, file = FileService.get_by_id(informs[0].file_id)
  130. FileService.update_by_id(file.id, {"name": req["name"]})
  131. if "parser_config" in req:
  132. DocumentService.update_parser_config(doc.id, req["parser_config"])
  133. if "chunk_method" in req:
  134. valid_chunk_method = {"naive","manual","qa","table","paper","book","laws","presentation","picture","one","knowledge_graph","email"}
  135. if req.get("chunk_method") not in valid_chunk_method:
  136. return get_error_data_result(f"`chunk_method` {req['chunk_method']} doesn't exist")
  137. if doc.parser_id.lower() == req["chunk_method"].lower():
  138. return get_result()
  139. if doc.type == FileType.VISUAL or re.search(
  140. r"\.(ppt|pptx|pages)$", doc.name):
  141. return get_error_data_result(retmsg="Not supported yet!")
  142. e = DocumentService.update_by_id(doc.id,
  143. {"parser_id": req["chunk_method"], "progress": 0, "progress_msg": "",
  144. "run": TaskStatus.UNSTART.value})
  145. if not e:
  146. return get_error_data_result(retmsg="Document not found!")
  147. req["parser_config"] = get_parser_config(req["chunk_method"], req.get("parser_config"))
  148. if doc.token_num > 0:
  149. e = DocumentService.increment_chunk_num(doc.id, doc.kb_id, doc.token_num * -1, doc.chunk_num * -1,
  150. doc.process_duation * -1)
  151. if not e:
  152. return get_error_data_result(retmsg="Document not found!")
  153. ELASTICSEARCH.deleteByQuery(
  154. Q("match", doc_id=doc.id), idxnm=search.index_name(tenant_id))
  155. return get_result()
  156. @manager.route('/dataset/<dataset_id>/document/<document_id>', methods=['GET'])
  157. @token_required
  158. def download(tenant_id, dataset_id, document_id):
  159. if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
  160. return get_error_data_result(retmsg=f'You do not own the dataset {dataset_id}.')
  161. doc = DocumentService.query(kb_id=dataset_id, id=document_id)
  162. if not doc:
  163. return get_error_data_result(retmsg=f'The dataset not own the document {document_id}.')
  164. # The process of downloading
  165. doc_id, doc_location = File2DocumentService.get_storage_address(doc_id=document_id) # minio address
  166. file_stream = STORAGE_IMPL.get(doc_id, doc_location)
  167. if not file_stream:
  168. return construct_json_result(message="This file is empty.", code=RetCode.DATA_ERROR)
  169. file = BytesIO(file_stream)
  170. # Use send_file with a proper filename and MIME type
  171. return send_file(
  172. file,
  173. as_attachment=True,
  174. download_name=doc[0].name,
  175. mimetype='application/octet-stream' # Set a default MIME type
  176. )
  177. @manager.route('/dataset/<dataset_id>/info', methods=['GET'])
  178. @token_required
  179. def list_docs(dataset_id, tenant_id):
  180. if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
  181. return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}. ")
  182. id = request.args.get("id")
  183. if not DocumentService.query(id=id,kb_id=dataset_id):
  184. return get_error_data_result(retmsg=f"You don't own the document {id}.")
  185. offset = int(request.args.get("offset", 1))
  186. keywords = request.args.get("keywords","")
  187. limit = int(request.args.get("limit", 1024))
  188. orderby = request.args.get("orderby", "create_time")
  189. if request.args.get("desc") == "False":
  190. desc = False
  191. else:
  192. desc = True
  193. docs, tol = DocumentService.get_list(dataset_id, offset, limit, orderby, desc, keywords, id)
  194. # rename key's name
  195. renamed_doc_list = []
  196. for doc in docs:
  197. key_mapping = {
  198. "chunk_num": "chunk_count",
  199. "kb_id": "dataset_id",
  200. "token_num": "token_count",
  201. "parser_id": "chunk_method"
  202. }
  203. run_mapping = {
  204. "0" :"UNSTART",
  205. "1":"RUNNING",
  206. "2":"CANCEL",
  207. "3":"DONE",
  208. "4":"FAIL"
  209. }
  210. renamed_doc = {}
  211. for key, value in doc.items():
  212. if key =="run":
  213. renamed_doc["run"]=run_mapping.get(str(value))
  214. new_key = key_mapping.get(key, key)
  215. renamed_doc[new_key] = value
  216. renamed_doc_list.append(renamed_doc)
  217. return get_result(data={"total": tol, "docs": renamed_doc_list})
  218. @manager.route('/dataset/<dataset_id>/document', methods=['DELETE'])
  219. @token_required
  220. def delete(tenant_id,dataset_id):
  221. if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
  222. return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}. ")
  223. req = request.json
  224. if not req:
  225. doc_ids=None
  226. else:
  227. doc_ids=req.get("ids")
  228. if not doc_ids:
  229. doc_list = []
  230. docs=DocumentService.query(kb_id=dataset_id)
  231. for doc in docs:
  232. doc_list.append(doc.id)
  233. else:
  234. doc_list=doc_ids
  235. root_folder = FileService.get_root_folder(tenant_id)
  236. pf_id = root_folder["id"]
  237. FileService.init_knowledgebase_docs(pf_id, tenant_id)
  238. errors = ""
  239. for doc_id in doc_list:
  240. try:
  241. e, doc = DocumentService.get_by_id(doc_id)
  242. if not e:
  243. return get_error_data_result(retmsg="Document not found!")
  244. tenant_id = DocumentService.get_tenant_id(doc_id)
  245. if not tenant_id:
  246. return get_error_data_result(retmsg="Tenant not found!")
  247. b, n = File2DocumentService.get_storage_address(doc_id=doc_id)
  248. if not DocumentService.remove_document(doc, tenant_id):
  249. return get_error_data_result(
  250. retmsg="Database error (Document removal)!")
  251. f2d = File2DocumentService.get_by_document_id(doc_id)
  252. FileService.filter_delete([File.source_type == FileSource.KNOWLEDGEBASE, File.id == f2d[0].file_id])
  253. File2DocumentService.delete_by_document_id(doc_id)
  254. STORAGE_IMPL.rm(b, n)
  255. except Exception as e:
  256. errors += str(e)
  257. if errors:
  258. return get_result(retmsg=errors, retcode=RetCode.SERVER_ERROR)
  259. return get_result()
  260. @manager.route('/dataset/<dataset_id>/chunk', methods=['POST'])
  261. @token_required
  262. def parse(tenant_id,dataset_id):
  263. if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
  264. return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}.")
  265. req = request.json
  266. if not req.get("document_ids"):
  267. return get_error_data_result("`document_ids` is required")
  268. for id in req["document_ids"]:
  269. doc = DocumentService.query(id=id,kb_id=dataset_id)
  270. if not doc:
  271. return get_error_data_result(retmsg=f"You don't own the document {id}.")
  272. if doc[0].progress != 0.0:
  273. return get_error_data_result("Can't stop parsing document with progress at 0 or 100")
  274. info = {"run": "1", "progress": 0}
  275. info["progress_msg"] = ""
  276. info["chunk_num"] = 0
  277. info["token_num"] = 0
  278. DocumentService.update_by_id(id, info)
  279. # if str(req["run"]) == TaskStatus.CANCEL.value:
  280. ELASTICSEARCH.deleteByQuery(
  281. Q("match", doc_id=id), idxnm=search.index_name(tenant_id))
  282. TaskService.filter_delete([Task.doc_id == id])
  283. e, doc = DocumentService.get_by_id(id)
  284. doc = doc.to_dict()
  285. doc["tenant_id"] = tenant_id
  286. bucket, name = File2DocumentService.get_storage_address(doc_id=doc["id"])
  287. queue_tasks(doc, bucket, name)
  288. return get_result()
  289. @manager.route('/dataset/<dataset_id>/chunk', methods=['DELETE'])
  290. @token_required
  291. def stop_parsing(tenant_id,dataset_id):
  292. if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
  293. return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}.")
  294. req = request.json
  295. if not req.get("document_ids"):
  296. return get_error_data_result("`document_ids` is required")
  297. for id in req["document_ids"]:
  298. doc = DocumentService.query(id=id, kb_id=dataset_id)
  299. if not doc:
  300. return get_error_data_result(retmsg=f"You don't own the document {id}.")
  301. if doc[0].progress == 100.0 or doc[0].progress == 0.0:
  302. return get_error_data_result("Can't stop parsing document with progress at 0 or 100")
  303. info = {"run": "2", "progress": 0}
  304. DocumentService.update_by_id(id, info)
  305. # if str(req["run"]) == TaskStatus.CANCEL.value:
  306. tenant_id = DocumentService.get_tenant_id(id)
  307. ELASTICSEARCH.deleteByQuery(
  308. Q("match", doc_id=id), idxnm=search.index_name(tenant_id))
  309. return get_result()
  310. @manager.route('/dataset/<dataset_id>/document/<document_id>/chunk', methods=['GET'])
  311. @token_required
  312. def list_chunks(tenant_id,dataset_id,document_id):
  313. if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
  314. return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}.")
  315. doc=DocumentService.query(id=document_id, kb_id=dataset_id)
  316. if not doc:
  317. return get_error_data_result(retmsg=f"You don't own the document {document_id}.")
  318. doc=doc[0]
  319. req = request.args
  320. doc_id = document_id
  321. page = int(req.get("offset", 1))
  322. size = int(req.get("limit", 30))
  323. question = req.get("keywords", "")
  324. query = {
  325. "doc_ids": [doc_id], "page": page, "size": size, "question": question, "sort": True
  326. }
  327. sres = retrievaler.search(query, search.index_name(tenant_id), highlight=True)
  328. key_mapping = {
  329. "chunk_num": "chunk_count",
  330. "kb_id": "dataset_id",
  331. "token_num": "token_count",
  332. "parser_id": "chunk_method"
  333. }
  334. run_mapping = {
  335. "0": "UNSTART",
  336. "1": "RUNNING",
  337. "2": "CANCEL",
  338. "3": "DONE",
  339. "4": "FAIL"
  340. }
  341. doc=doc.to_dict()
  342. renamed_doc = {}
  343. for key, value in doc.items():
  344. if key == "run":
  345. renamed_doc["run"] = run_mapping.get(str(value))
  346. new_key = key_mapping.get(key, key)
  347. renamed_doc[new_key] = value
  348. res = {"total": sres.total, "chunks": [], "doc": renamed_doc}
  349. origin_chunks = []
  350. sign = 0
  351. for id in sres.ids:
  352. d = {
  353. "chunk_id": id,
  354. "content_with_weight": rmSpace(sres.highlight[id]) if question and id in sres.highlight else sres.field[
  355. id].get(
  356. "content_with_weight", ""),
  357. "doc_id": sres.field[id]["doc_id"],
  358. "docnm_kwd": sres.field[id]["docnm_kwd"],
  359. "important_kwd": sres.field[id].get("important_kwd", []),
  360. "img_id": sres.field[id].get("img_id", ""),
  361. "available_int": sres.field[id].get("available_int", 1),
  362. "positions": sres.field[id].get("position_int", "").split("\t")
  363. }
  364. if len(d["positions"]) % 5 == 0:
  365. poss = []
  366. for i in range(0, len(d["positions"]), 5):
  367. poss.append([float(d["positions"][i]), float(d["positions"][i + 1]), float(d["positions"][i + 2]),
  368. float(d["positions"][i + 3]), float(d["positions"][i + 4])])
  369. d["positions"] = poss
  370. origin_chunks.append(d)
  371. if req.get("id"):
  372. if req.get("id") == id:
  373. origin_chunks.clear()
  374. origin_chunks.append(d)
  375. sign = 1
  376. break
  377. if req.get("id"):
  378. if sign == 0:
  379. return get_error_data_result(f"Can't find this chunk {req.get('id')}")
  380. for chunk in origin_chunks:
  381. key_mapping = {
  382. "chunk_id": "id",
  383. "content_with_weight": "content",
  384. "doc_id": "document_id",
  385. "important_kwd": "important_keywords",
  386. "img_id": "image_id"
  387. }
  388. renamed_chunk = {}
  389. for key, value in chunk.items():
  390. new_key = key_mapping.get(key, key)
  391. renamed_chunk[new_key] = value
  392. res["chunks"].append(renamed_chunk)
  393. return get_result(data=res)
  394. @manager.route('/dataset/<dataset_id>/document/<document_id>/chunk', methods=['POST'])
  395. @token_required
  396. def add_chunk(tenant_id,dataset_id,document_id):
  397. if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
  398. return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}.")
  399. doc = DocumentService.query(id=document_id, kb_id=dataset_id)
  400. if not doc:
  401. return get_error_data_result(retmsg=f"You don't own the document {document_id}.")
  402. doc = doc[0]
  403. req = request.json
  404. if not req.get("content"):
  405. return get_error_data_result(retmsg="`content` is required")
  406. if "important_keywords" in req:
  407. if type(req["important_keywords"]) != list:
  408. return get_error_data_result("`important_keywords` is required to be a list")
  409. md5 = hashlib.md5()
  410. md5.update((req["content"] + document_id).encode("utf-8"))
  411. chunk_id = md5.hexdigest()
  412. d = {"id": chunk_id, "content_ltks": rag_tokenizer.tokenize(req["content"]),
  413. "content_with_weight": req["content"]}
  414. d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
  415. d["important_kwd"] = req.get("important_keywords", [])
  416. d["important_tks"] = rag_tokenizer.tokenize(" ".join(req.get("important_keywords", [])))
  417. d["create_time"] = str(datetime.datetime.now()).replace("T", " ")[:19]
  418. d["create_timestamp_flt"] = datetime.datetime.now().timestamp()
  419. d["kb_id"] = [doc.kb_id]
  420. d["docnm_kwd"] = doc.name
  421. d["doc_id"] = doc.id
  422. embd_id = DocumentService.get_embd_id(document_id)
  423. embd_mdl = TenantLLMService.model_instance(
  424. tenant_id, LLMType.EMBEDDING.value, embd_id)
  425. v, c = embd_mdl.encode([doc.name, req["content"]])
  426. v = 0.1 * v[0] + 0.9 * v[1]
  427. d["q_%d_vec" % len(v)] = v.tolist()
  428. ELASTICSEARCH.upsert([d], search.index_name(tenant_id))
  429. DocumentService.increment_chunk_num(
  430. doc.id, doc.kb_id, c, 1, 0)
  431. d["chunk_id"] = chunk_id
  432. # rename keys
  433. key_mapping = {
  434. "chunk_id": "id",
  435. "content_with_weight": "content",
  436. "doc_id": "document_id",
  437. "important_kwd": "important_keywords",
  438. "kb_id": "dataset_id",
  439. "create_timestamp_flt": "create_timestamp",
  440. "create_time": "create_time",
  441. "document_keyword": "document",
  442. }
  443. renamed_chunk = {}
  444. for key, value in d.items():
  445. if key in key_mapping:
  446. new_key = key_mapping.get(key, key)
  447. renamed_chunk[new_key] = value
  448. return get_result(data={"chunk": renamed_chunk})
  449. # return get_result(data={"chunk_id": chunk_id})
  450. @manager.route('dataset/<dataset_id>/document/<document_id>/chunk', methods=['DELETE'])
  451. @token_required
  452. def rm_chunk(tenant_id,dataset_id,document_id):
  453. if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
  454. return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}.")
  455. doc = DocumentService.query(id=document_id, kb_id=dataset_id)
  456. if not doc:
  457. return get_error_data_result(retmsg=f"You don't own the document {document_id}.")
  458. doc = doc[0]
  459. req = request.json
  460. if not req.get("chunk_ids"):
  461. return get_error_data_result("`chunk_ids` is required")
  462. query = {
  463. "doc_ids": [doc.id], "page": 1, "size": 1024, "question": "", "sort": True}
  464. sres = retrievaler.search(query, search.index_name(tenant_id), highlight=True)
  465. for chunk_id in req.get("chunk_ids"):
  466. if chunk_id not in sres.ids:
  467. return get_error_data_result(f"Chunk {chunk_id} not found")
  468. if not ELASTICSEARCH.deleteByQuery(
  469. Q("ids", values=req["chunk_ids"]), search.index_name(tenant_id)):
  470. return get_error_data_result(retmsg="Index updating failure")
  471. deleted_chunk_ids = req["chunk_ids"]
  472. chunk_number = len(deleted_chunk_ids)
  473. DocumentService.decrement_chunk_num(doc.id, doc.kb_id, 1, chunk_number, 0)
  474. return get_result()
  475. @manager.route('/dataset/<dataset_id>/document/<document_id>/chunk/<chunk_id>', methods=['PUT'])
  476. @token_required
  477. def update_chunk(tenant_id,dataset_id,document_id,chunk_id):
  478. try:
  479. res = ELASTICSEARCH.get(
  480. chunk_id, search.index_name(
  481. tenant_id))
  482. except Exception as e:
  483. return get_error_data_result(f"Can't find this chunk {chunk_id}")
  484. if not KnowledgebaseService.query(id=dataset_id, tenant_id=tenant_id):
  485. return get_error_data_result(retmsg=f"You don't own the dataset {dataset_id}.")
  486. doc = DocumentService.query(id=document_id, kb_id=dataset_id)
  487. if not doc:
  488. return get_error_data_result(retmsg=f"You don't own the document {document_id}.")
  489. doc = doc[0]
  490. query = {
  491. "doc_ids": [document_id], "page": 1, "size": 1024, "question": "", "sort": True
  492. }
  493. sres = retrievaler.search(query, search.index_name(tenant_id), highlight=True)
  494. if chunk_id not in sres.ids:
  495. return get_error_data_result(f"You don't own the chunk {chunk_id}")
  496. req = request.json
  497. content=res["_source"].get("content_with_weight")
  498. d = {
  499. "id": chunk_id,
  500. "content_with_weight": req.get("content",content)}
  501. d["content_ltks"] = rag_tokenizer.tokenize(d["content_with_weight"])
  502. d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
  503. if "important_keywords" in req:
  504. if not isinstance(req["important_keywords"],list):
  505. return get_error_data_result("`important_keywords` should be a list")
  506. d["important_kwd"] = req.get("important_keywords")
  507. d["important_tks"] = rag_tokenizer.tokenize(" ".join(req["important_keywords"]))
  508. if "available" in req:
  509. d["available_int"] = int(req["available"])
  510. embd_id = DocumentService.get_embd_id(document_id)
  511. embd_mdl = TenantLLMService.model_instance(
  512. tenant_id, LLMType.EMBEDDING.value, embd_id)
  513. if doc.parser_id == ParserType.QA:
  514. arr = [
  515. t for t in re.split(
  516. r"[\n\t]",
  517. d["content_with_weight"]) if len(t) > 1]
  518. if len(arr) != 2:
  519. return get_error_data_result(
  520. retmsg="Q&A must be separated by TAB/ENTER key.")
  521. q, a = rmPrefix(arr[0]), rmPrefix(arr[1])
  522. d = beAdoc(d, arr[0], arr[1], not any(
  523. [rag_tokenizer.is_chinese(t) for t in q + a]))
  524. v, c = embd_mdl.encode([doc.name, d["content_with_weight"]])
  525. v = 0.1 * v[0] + 0.9 * v[1] if doc.parser_id != ParserType.QA else v[1]
  526. d["q_%d_vec" % len(v)] = v.tolist()
  527. ELASTICSEARCH.upsert([d], search.index_name(tenant_id))
  528. return get_result()
  529. @manager.route('/retrieval', methods=['POST'])
  530. @token_required
  531. def retrieval_test(tenant_id):
  532. req = request.json
  533. if not req.get("dataset_ids"):
  534. return get_error_data_result("`datasets` is required.")
  535. kb_ids = req["dataset_ids"]
  536. if not isinstance(kb_ids,list):
  537. return get_error_data_result("`datasets` should be a list")
  538. kbs = KnowledgebaseService.get_by_ids(kb_ids)
  539. for id in kb_ids:
  540. if not KnowledgebaseService.query(id=id,tenant_id=tenant_id):
  541. return get_error_data_result(f"You don't own the dataset {id}.")
  542. embd_nms = list(set([kb.embd_id for kb in kbs]))
  543. if len(embd_nms) != 1:
  544. return get_result(
  545. retmsg='Datasets use different embedding models."',
  546. retcode=RetCode.AUTHENTICATION_ERROR)
  547. if "question" not in req:
  548. return get_error_data_result("`question` is required.")
  549. page = int(req.get("offset", 1))
  550. size = int(req.get("limit", 1024))
  551. question = req["question"]
  552. doc_ids = req.get("document_ids", [])
  553. if not isinstance(doc_ids,list):
  554. return get_error_data_result("`documents` should be a list")
  555. doc_ids_list=KnowledgebaseService.list_documents_by_ids(kb_ids)
  556. for doc_id in doc_ids:
  557. if doc_id not in doc_ids_list:
  558. return get_error_data_result(f"The datasets don't own the document {doc_id}")
  559. similarity_threshold = float(req.get("similarity_threshold", 0.2))
  560. vector_similarity_weight = float(req.get("vector_similarity_weight", 0.3))
  561. top = int(req.get("top_k", 1024))
  562. if req.get("highlight")=="False" or req.get("highlight")=="false":
  563. highlight = False
  564. else:
  565. highlight = True
  566. try:
  567. e, kb = KnowledgebaseService.get_by_id(kb_ids[0])
  568. if not e:
  569. return get_error_data_result(retmsg="Dataset not found!")
  570. embd_mdl = TenantLLMService.model_instance(
  571. kb.tenant_id, LLMType.EMBEDDING.value, llm_name=kb.embd_id)
  572. rerank_mdl = None
  573. if req.get("rerank_id"):
  574. rerank_mdl = TenantLLMService.model_instance(
  575. kb.tenant_id, LLMType.RERANK.value, llm_name=req["rerank_id"])
  576. if req.get("keyword", False):
  577. chat_mdl = TenantLLMService.model_instance(kb.tenant_id, LLMType.CHAT)
  578. question += keyword_extraction(chat_mdl, question)
  579. retr = retrievaler if kb.parser_id != ParserType.KG else kg_retrievaler
  580. ranks = retr.retrieval(question, embd_mdl, kb.tenant_id, kb_ids, page, size,
  581. similarity_threshold, vector_similarity_weight, top,
  582. doc_ids, rerank_mdl=rerank_mdl, highlight=highlight)
  583. for c in ranks["chunks"]:
  584. if "vector" in c:
  585. del c["vector"]
  586. ##rename keys
  587. renamed_chunks = []
  588. for chunk in ranks["chunks"]:
  589. key_mapping = {
  590. "chunk_id": "id",
  591. "content_with_weight": "content",
  592. "doc_id": "document_id",
  593. "important_kwd": "important_keywords",
  594. "docnm_kwd": "document_keyword"
  595. }
  596. rename_chunk = {}
  597. for key, value in chunk.items():
  598. new_key = key_mapping.get(key, key)
  599. rename_chunk[new_key] = value
  600. renamed_chunks.append(rename_chunk)
  601. ranks["chunks"] = renamed_chunks
  602. return get_result(data=ranks)
  603. except Exception as e:
  604. if str(e).find("not_found") > 0:
  605. return get_result(retmsg=f'No chunk found! Check the chunk status please!',
  606. retcode=RetCode.DATA_ERROR)
  607. return server_error_response(e)