Du kan inte välja fler än 25 ämnen Ämnen måste starta med en bokstav eller siffra, kan innehålla bindestreck ('-') och vara max 35 tecken långa.

doc.py 29KB

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