Vous ne pouvez pas sélectionner plus de 25 sujets Les noms de sujets doivent commencer par une lettre ou un nombre, peuvent contenir des tirets ('-') et peuvent comporter jusqu'à 35 caractères.

api_app.py 28KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737
  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 json
  17. import os
  18. import re
  19. from datetime import datetime, timedelta
  20. from flask import request, Response
  21. from flask_login import login_required, current_user
  22. from api.db import FileType, ParserType, FileSource
  23. from api.db.db_models import APIToken, API4Conversation, Task, File
  24. from api.db.services import duplicate_name
  25. from api.db.services.api_service import APITokenService, API4ConversationService
  26. from api.db.services.dialog_service import DialogService, chat
  27. from api.db.services.document_service import DocumentService
  28. from api.db.services.file2document_service import File2DocumentService
  29. from api.db.services.file_service import FileService
  30. from api.db.services.knowledgebase_service import KnowledgebaseService
  31. from api.db.services.task_service import queue_tasks, TaskService
  32. from api.db.services.user_service import UserTenantService
  33. from api.settings import RetCode, retrievaler
  34. from api.utils import get_uuid, current_timestamp, datetime_format
  35. from api.utils.api_utils import server_error_response, get_data_error_result, get_json_result, validate_request
  36. from itsdangerous import URLSafeTimedSerializer
  37. from api.utils.file_utils import filename_type, thumbnail
  38. from rag.utils.minio_conn import MINIO
  39. from api.db.services.canvas_service import CanvasTemplateService, UserCanvasService
  40. from agent.canvas import Canvas
  41. from functools import partial
  42. def generate_confirmation_token(tenent_id):
  43. serializer = URLSafeTimedSerializer(tenent_id)
  44. return "ragflow-" + serializer.dumps(get_uuid(), salt=tenent_id)[2:34]
  45. @manager.route('/new_token', methods=['POST'])
  46. @login_required
  47. def new_token():
  48. req = request.json
  49. try:
  50. tenants = UserTenantService.query(user_id=current_user.id)
  51. if not tenants:
  52. return get_data_error_result(retmsg="Tenant not found!")
  53. tenant_id = tenants[0].tenant_id
  54. obj = {"tenant_id": tenant_id, "token": generate_confirmation_token(tenant_id),
  55. "create_time": current_timestamp(),
  56. "create_date": datetime_format(datetime.now()),
  57. "update_time": None,
  58. "update_date": None
  59. }
  60. if req.get("canvas_id"):
  61. obj["dialog_id"] = req["canvas_id"]
  62. obj["source"] = "agent"
  63. else:
  64. obj["dialog_id"] = req["dialog_id"]
  65. if not APITokenService.save(**obj):
  66. return get_data_error_result(retmsg="Fail to new a dialog!")
  67. return get_json_result(data=obj)
  68. except Exception as e:
  69. return server_error_response(e)
  70. @manager.route('/token_list', methods=['GET'])
  71. @login_required
  72. def token_list():
  73. try:
  74. tenants = UserTenantService.query(user_id=current_user.id)
  75. if not tenants:
  76. return get_data_error_result(retmsg="Tenant not found!")
  77. objs = APITokenService.query(tenant_id=tenants[0].tenant_id, dialog_id=request.args["dialog_id"])
  78. return get_json_result(data=[o.to_dict() for o in objs])
  79. except Exception as e:
  80. return server_error_response(e)
  81. @manager.route('/rm', methods=['POST'])
  82. @validate_request("tokens", "tenant_id")
  83. @login_required
  84. def rm():
  85. req = request.json
  86. try:
  87. for token in req["tokens"]:
  88. APITokenService.filter_delete(
  89. [APIToken.tenant_id == req["tenant_id"], APIToken.token == token])
  90. return get_json_result(data=True)
  91. except Exception as e:
  92. return server_error_response(e)
  93. @manager.route('/stats', methods=['GET'])
  94. @login_required
  95. def stats():
  96. try:
  97. tenants = UserTenantService.query(user_id=current_user.id)
  98. if not tenants:
  99. return get_data_error_result(retmsg="Tenant not found!")
  100. objs = API4ConversationService.stats(
  101. tenants[0].tenant_id,
  102. request.args.get(
  103. "from_date",
  104. (datetime.now() -
  105. timedelta(
  106. days=7)).strftime("%Y-%m-%d 24:00:00")),
  107. request.args.get(
  108. "to_date",
  109. datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
  110. res = {
  111. "pv": [(o["dt"], o["pv"]) for o in objs],
  112. "uv": [(o["dt"], o["uv"]) for o in objs],
  113. "speed": [(o["dt"], float(o["tokens"]) / (float(o["duration"] + 0.1))) for o in objs],
  114. "tokens": [(o["dt"], float(o["tokens"]) / 1000.) for o in objs],
  115. "round": [(o["dt"], o["round"]) for o in objs],
  116. "thumb_up": [(o["dt"], o["thumb_up"]) for o in objs]
  117. }
  118. return get_json_result(data=res)
  119. except Exception as e:
  120. return server_error_response(e)
  121. @manager.route('/new_conversation', methods=['GET'])
  122. def set_conversation():
  123. token = request.headers.get('Authorization').split()[1]
  124. objs = APIToken.query(token=token)
  125. if not objs:
  126. return get_json_result(
  127. data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
  128. req = request.json
  129. try:
  130. if objs[0].source == "agent":
  131. e, c = UserCanvasService.get_by_id(objs[0].dialog_id)
  132. if not e:
  133. return server_error_response("canvas not found.")
  134. conv = {
  135. "id": get_uuid(),
  136. "dialog_id": c.id,
  137. "user_id": request.args.get("user_id", ""),
  138. "message": [{"role": "assistant", "content": "Hi there!"}],
  139. "source": "agent"
  140. }
  141. API4ConversationService.save(**conv)
  142. return get_json_result(data=conv)
  143. else:
  144. e, dia = DialogService.get_by_id(objs[0].dialog_id)
  145. if not e:
  146. return get_data_error_result(retmsg="Dialog not found")
  147. conv = {
  148. "id": get_uuid(),
  149. "dialog_id": dia.id,
  150. "user_id": request.args.get("user_id", ""),
  151. "message": [{"role": "assistant", "content": dia.prompt_config["prologue"]}]
  152. }
  153. API4ConversationService.save(**conv)
  154. return get_json_result(data=conv)
  155. except Exception as e:
  156. return server_error_response(e)
  157. @manager.route('/completion', methods=['POST'])
  158. @validate_request("conversation_id", "messages")
  159. def completion():
  160. token = request.headers.get('Authorization').split()[1]
  161. objs = APIToken.query(token=token)
  162. if not objs:
  163. return get_json_result(
  164. data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
  165. req = request.json
  166. e, conv = API4ConversationService.get_by_id(req["conversation_id"])
  167. if not e:
  168. return get_data_error_result(retmsg="Conversation not found!")
  169. if "quote" not in req: req["quote"] = False
  170. msg = []
  171. for m in req["messages"]:
  172. if m["role"] == "system":
  173. continue
  174. if m["role"] == "assistant" and not msg:
  175. continue
  176. msg.append({"role": m["role"], "content": m["content"]})
  177. def fillin_conv(ans):
  178. nonlocal conv
  179. if not conv.reference:
  180. conv.reference.append(ans["reference"])
  181. else:
  182. conv.reference[-1] = ans["reference"]
  183. conv.message[-1] = {"role": "assistant", "content": ans["answer"]}
  184. def rename_field(ans):
  185. reference = ans['reference']
  186. if not isinstance(reference, dict):
  187. return
  188. for chunk_i in reference.get('chunks', []):
  189. if 'docnm_kwd' in chunk_i:
  190. chunk_i['doc_name'] = chunk_i['docnm_kwd']
  191. chunk_i.pop('docnm_kwd')
  192. try:
  193. if conv.source == "agent":
  194. stream = req.get("stream", True)
  195. conv.message.append(msg[-1])
  196. e, cvs = UserCanvasService.get_by_id(conv.dialog_id)
  197. if not e:
  198. return server_error_response("canvas not found.")
  199. del req["conversation_id"]
  200. del req["messages"]
  201. if not isinstance(cvs.dsl, str):
  202. cvs.dsl = json.dumps(cvs.dsl, ensure_ascii=False)
  203. if not conv.reference:
  204. conv.reference = []
  205. conv.message.append({"role": "assistant", "content": ""})
  206. conv.reference.append({"chunks": [], "doc_aggs": []})
  207. final_ans = {"reference": [], "content": ""}
  208. canvas = Canvas(cvs.dsl, objs[0].tenant_id)
  209. canvas.messages.append(msg[-1])
  210. canvas.add_user_input(msg[-1]["content"])
  211. answer = canvas.run(stream=stream)
  212. assert answer is not None, "Nothing. Is it over?"
  213. if stream:
  214. assert isinstance(answer, partial), "Nothing. Is it over?"
  215. def sse():
  216. nonlocal answer, cvs, conv
  217. try:
  218. for ans in answer():
  219. for k in ans.keys():
  220. final_ans[k] = ans[k]
  221. ans = {"answer": ans["content"], "reference": ans.get("reference", [])}
  222. fillin_conv(ans)
  223. rename_field(ans)
  224. yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans},
  225. ensure_ascii=False) + "\n\n"
  226. canvas.messages.append({"role": "assistant", "content": final_ans["content"]})
  227. if final_ans.get("reference"):
  228. canvas.reference.append(final_ans["reference"])
  229. cvs.dsl = json.loads(str(canvas))
  230. API4ConversationService.append_message(conv.id, conv.to_dict())
  231. except Exception as e:
  232. yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
  233. "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
  234. ensure_ascii=False) + "\n\n"
  235. yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
  236. resp = Response(sse(), mimetype="text/event-stream")
  237. resp.headers.add_header("Cache-control", "no-cache")
  238. resp.headers.add_header("Connection", "keep-alive")
  239. resp.headers.add_header("X-Accel-Buffering", "no")
  240. resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
  241. return resp
  242. final_ans["content"] = "\n".join(answer["content"]) if "content" in answer else ""
  243. canvas.messages.append({"role": "assistant", "content": final_ans["content"]})
  244. if final_ans.get("reference"):
  245. canvas.reference.append(final_ans["reference"])
  246. cvs.dsl = json.loads(str(canvas))
  247. result = None
  248. for ans in answer():
  249. ans = {"answer": ans["content"], "reference": ans.get("reference", [])}
  250. result = ans
  251. fillin_conv(ans)
  252. API4ConversationService.append_message(conv.id, conv.to_dict())
  253. break
  254. rename_field(result)
  255. return get_json_result(data=result)
  256. #******************For dialog******************
  257. conv.message.append(msg[-1])
  258. e, dia = DialogService.get_by_id(conv.dialog_id)
  259. if not e:
  260. return get_data_error_result(retmsg="Dialog not found!")
  261. del req["conversation_id"]
  262. del req["messages"]
  263. if not conv.reference:
  264. conv.reference = []
  265. conv.message.append({"role": "assistant", "content": ""})
  266. conv.reference.append({"chunks": [], "doc_aggs": []})
  267. def stream():
  268. nonlocal dia, msg, req, conv
  269. try:
  270. for ans in chat(dia, msg, True, **req):
  271. fillin_conv(ans)
  272. rename_field(ans)
  273. yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans},
  274. ensure_ascii=False) + "\n\n"
  275. API4ConversationService.append_message(conv.id, conv.to_dict())
  276. except Exception as e:
  277. yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
  278. "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
  279. ensure_ascii=False) + "\n\n"
  280. yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
  281. if req.get("stream", True):
  282. resp = Response(stream(), mimetype="text/event-stream")
  283. resp.headers.add_header("Cache-control", "no-cache")
  284. resp.headers.add_header("Connection", "keep-alive")
  285. resp.headers.add_header("X-Accel-Buffering", "no")
  286. resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
  287. return resp
  288. answer = None
  289. for ans in chat(dia, msg, **req):
  290. answer = ans
  291. fillin_conv(ans)
  292. API4ConversationService.append_message(conv.id, conv.to_dict())
  293. break
  294. rename_field(answer)
  295. return get_json_result(data=answer)
  296. except Exception as e:
  297. return server_error_response(e)
  298. @manager.route('/conversation/<conversation_id>', methods=['GET'])
  299. # @login_required
  300. def get(conversation_id):
  301. try:
  302. e, conv = API4ConversationService.get_by_id(conversation_id)
  303. if not e:
  304. return get_data_error_result(retmsg="Conversation not found!")
  305. conv = conv.to_dict()
  306. for referenct_i in conv['reference']:
  307. if referenct_i is None or len(referenct_i) == 0:
  308. continue
  309. for chunk_i in referenct_i['chunks']:
  310. if 'docnm_kwd' in chunk_i.keys():
  311. chunk_i['doc_name'] = chunk_i['docnm_kwd']
  312. chunk_i.pop('docnm_kwd')
  313. return get_json_result(data=conv)
  314. except Exception as e:
  315. return server_error_response(e)
  316. @manager.route('/document/upload', methods=['POST'])
  317. @validate_request("kb_name")
  318. def upload():
  319. token = request.headers.get('Authorization').split()[1]
  320. objs = APIToken.query(token=token)
  321. if not objs:
  322. return get_json_result(
  323. data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
  324. kb_name = request.form.get("kb_name").strip()
  325. tenant_id = objs[0].tenant_id
  326. try:
  327. e, kb = KnowledgebaseService.get_by_name(kb_name, tenant_id)
  328. if not e:
  329. return get_data_error_result(
  330. retmsg="Can't find this knowledgebase!")
  331. kb_id = kb.id
  332. except Exception as e:
  333. return server_error_response(e)
  334. if 'file' not in request.files:
  335. return get_json_result(
  336. data=False, retmsg='No file part!', retcode=RetCode.ARGUMENT_ERROR)
  337. file = request.files['file']
  338. if file.filename == '':
  339. return get_json_result(
  340. data=False, retmsg='No file selected!', retcode=RetCode.ARGUMENT_ERROR)
  341. root_folder = FileService.get_root_folder(tenant_id)
  342. pf_id = root_folder["id"]
  343. FileService.init_knowledgebase_docs(pf_id, tenant_id)
  344. kb_root_folder = FileService.get_kb_folder(tenant_id)
  345. kb_folder = FileService.new_a_file_from_kb(kb.tenant_id, kb.name, kb_root_folder["id"])
  346. try:
  347. if DocumentService.get_doc_count(kb.tenant_id) >= int(os.environ.get('MAX_FILE_NUM_PER_USER', 8192)):
  348. return get_data_error_result(
  349. retmsg="Exceed the maximum file number of a free user!")
  350. filename = duplicate_name(
  351. DocumentService.query,
  352. name=file.filename,
  353. kb_id=kb_id)
  354. filetype = filename_type(filename)
  355. if not filetype:
  356. return get_data_error_result(
  357. retmsg="This type of file has not been supported yet!")
  358. location = filename
  359. while MINIO.obj_exist(kb_id, location):
  360. location += "_"
  361. blob = request.files['file'].read()
  362. MINIO.put(kb_id, location, blob)
  363. doc = {
  364. "id": get_uuid(),
  365. "kb_id": kb.id,
  366. "parser_id": kb.parser_id,
  367. "parser_config": kb.parser_config,
  368. "created_by": kb.tenant_id,
  369. "type": filetype,
  370. "name": filename,
  371. "location": location,
  372. "size": len(blob),
  373. "thumbnail": thumbnail(filename, blob)
  374. }
  375. form_data = request.form
  376. if "parser_id" in form_data.keys():
  377. if request.form.get("parser_id").strip() in list(vars(ParserType).values())[1:-3]:
  378. doc["parser_id"] = request.form.get("parser_id").strip()
  379. if doc["type"] == FileType.VISUAL:
  380. doc["parser_id"] = ParserType.PICTURE.value
  381. if doc["type"] == FileType.AURAL:
  382. doc["parser_id"] = ParserType.AUDIO.value
  383. if re.search(r"\.(ppt|pptx|pages)$", filename):
  384. doc["parser_id"] = ParserType.PRESENTATION.value
  385. doc_result = DocumentService.insert(doc)
  386. FileService.add_file_from_kb(doc, kb_folder["id"], kb.tenant_id)
  387. except Exception as e:
  388. return server_error_response(e)
  389. if "run" in form_data.keys():
  390. if request.form.get("run").strip() == "1":
  391. try:
  392. info = {"run": 1, "progress": 0}
  393. info["progress_msg"] = ""
  394. info["chunk_num"] = 0
  395. info["token_num"] = 0
  396. DocumentService.update_by_id(doc["id"], info)
  397. # if str(req["run"]) == TaskStatus.CANCEL.value:
  398. tenant_id = DocumentService.get_tenant_id(doc["id"])
  399. if not tenant_id:
  400. return get_data_error_result(retmsg="Tenant not found!")
  401. # e, doc = DocumentService.get_by_id(doc["id"])
  402. TaskService.filter_delete([Task.doc_id == doc["id"]])
  403. e, doc = DocumentService.get_by_id(doc["id"])
  404. doc = doc.to_dict()
  405. doc["tenant_id"] = tenant_id
  406. bucket, name = File2DocumentService.get_minio_address(doc_id=doc["id"])
  407. queue_tasks(doc, bucket, name)
  408. except Exception as e:
  409. return server_error_response(e)
  410. return get_json_result(data=doc_result.to_json())
  411. @manager.route('/list_chunks', methods=['POST'])
  412. # @login_required
  413. def list_chunks():
  414. token = request.headers.get('Authorization').split()[1]
  415. objs = APIToken.query(token=token)
  416. if not objs:
  417. return get_json_result(
  418. data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
  419. req = request.json
  420. try:
  421. if "doc_name" in req.keys():
  422. tenant_id = DocumentService.get_tenant_id_by_name(req['doc_name'])
  423. doc_id = DocumentService.get_doc_id_by_doc_name(req['doc_name'])
  424. elif "doc_id" in req.keys():
  425. tenant_id = DocumentService.get_tenant_id(req['doc_id'])
  426. doc_id = req['doc_id']
  427. else:
  428. return get_json_result(
  429. data=False, retmsg="Can't find doc_name or doc_id"
  430. )
  431. res = retrievaler.chunk_list(doc_id=doc_id, tenant_id=tenant_id)
  432. res = [
  433. {
  434. "content": res_item["content_with_weight"],
  435. "doc_name": res_item["docnm_kwd"],
  436. "img_id": res_item["img_id"]
  437. } for res_item in res
  438. ]
  439. except Exception as e:
  440. return server_error_response(e)
  441. return get_json_result(data=res)
  442. @manager.route('/list_kb_docs', methods=['POST'])
  443. # @login_required
  444. def list_kb_docs():
  445. token = request.headers.get('Authorization').split()[1]
  446. objs = APIToken.query(token=token)
  447. if not objs:
  448. return get_json_result(
  449. data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
  450. req = request.json
  451. tenant_id = objs[0].tenant_id
  452. kb_name = req.get("kb_name", "").strip()
  453. try:
  454. e, kb = KnowledgebaseService.get_by_name(kb_name, tenant_id)
  455. if not e:
  456. return get_data_error_result(
  457. retmsg="Can't find this knowledgebase!")
  458. kb_id = kb.id
  459. except Exception as e:
  460. return server_error_response(e)
  461. page_number = int(req.get("page", 1))
  462. items_per_page = int(req.get("page_size", 15))
  463. orderby = req.get("orderby", "create_time")
  464. desc = req.get("desc", True)
  465. keywords = req.get("keywords", "")
  466. try:
  467. docs, tol = DocumentService.get_by_kb_id(
  468. kb_id, page_number, items_per_page, orderby, desc, keywords)
  469. docs = [{"doc_id": doc['id'], "doc_name": doc['name']} for doc in docs]
  470. return get_json_result(data={"total": tol, "docs": docs})
  471. except Exception as e:
  472. return server_error_response(e)
  473. @manager.route('/document', methods=['DELETE'])
  474. # @login_required
  475. def document_rm():
  476. token = request.headers.get('Authorization').split()[1]
  477. objs = APIToken.query(token=token)
  478. if not objs:
  479. return get_json_result(
  480. data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
  481. tenant_id = objs[0].tenant_id
  482. req = request.json
  483. doc_ids = []
  484. try:
  485. doc_ids = [DocumentService.get_doc_id_by_doc_name(doc_name) for doc_name in req.get("doc_names", [])]
  486. for doc_id in req.get("doc_ids", []):
  487. if doc_id not in doc_ids:
  488. doc_ids.append(doc_id)
  489. if not doc_ids:
  490. return get_json_result(
  491. data=False, retmsg="Can't find doc_names or doc_ids"
  492. )
  493. except Exception as e:
  494. return server_error_response(e)
  495. root_folder = FileService.get_root_folder(tenant_id)
  496. pf_id = root_folder["id"]
  497. FileService.init_knowledgebase_docs(pf_id, tenant_id)
  498. errors = ""
  499. for doc_id in doc_ids:
  500. try:
  501. e, doc = DocumentService.get_by_id(doc_id)
  502. if not e:
  503. return get_data_error_result(retmsg="Document not found!")
  504. tenant_id = DocumentService.get_tenant_id(doc_id)
  505. if not tenant_id:
  506. return get_data_error_result(retmsg="Tenant not found!")
  507. b, n = File2DocumentService.get_minio_address(doc_id=doc_id)
  508. if not DocumentService.remove_document(doc, tenant_id):
  509. return get_data_error_result(
  510. retmsg="Database error (Document removal)!")
  511. f2d = File2DocumentService.get_by_document_id(doc_id)
  512. FileService.filter_delete([File.source_type == FileSource.KNOWLEDGEBASE, File.id == f2d[0].file_id])
  513. File2DocumentService.delete_by_document_id(doc_id)
  514. MINIO.rm(b, n)
  515. except Exception as e:
  516. errors += str(e)
  517. if errors:
  518. return get_json_result(data=False, retmsg=errors, retcode=RetCode.SERVER_ERROR)
  519. return get_json_result(data=True)
  520. @manager.route('/completion_aibotk', methods=['POST'])
  521. @validate_request("Authorization", "conversation_id", "word")
  522. def completion_faq():
  523. import base64
  524. req = request.json
  525. token = req["Authorization"]
  526. objs = APIToken.query(token=token)
  527. if not objs:
  528. return get_json_result(
  529. data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
  530. e, conv = API4ConversationService.get_by_id(req["conversation_id"])
  531. if not e:
  532. return get_data_error_result(retmsg="Conversation not found!")
  533. if "quote" not in req: req["quote"] = True
  534. msg = []
  535. msg.append({"role": "user", "content": req["word"]})
  536. try:
  537. conv.message.append(msg[-1])
  538. e, dia = DialogService.get_by_id(conv.dialog_id)
  539. if not e:
  540. return get_data_error_result(retmsg="Dialog not found!")
  541. del req["conversation_id"]
  542. if not conv.reference:
  543. conv.reference = []
  544. conv.message.append({"role": "assistant", "content": ""})
  545. conv.reference.append({"chunks": [], "doc_aggs": []})
  546. def fillin_conv(ans):
  547. nonlocal conv
  548. if not conv.reference:
  549. conv.reference.append(ans["reference"])
  550. else:
  551. conv.reference[-1] = ans["reference"]
  552. conv.message[-1] = {"role": "assistant", "content": ans["answer"]}
  553. data_type_picture = {
  554. "type": 3,
  555. "url": "base64 content"
  556. }
  557. data = [
  558. {
  559. "type": 1,
  560. "content": ""
  561. }
  562. ]
  563. ans = ""
  564. for a in chat(dia, msg, stream=False, **req):
  565. ans = a
  566. break
  567. data[0]["content"] += re.sub(r'##\d\$\$', '', ans["answer"])
  568. fillin_conv(ans)
  569. API4ConversationService.append_message(conv.id, conv.to_dict())
  570. chunk_idxs = [int(match[2]) for match in re.findall(r'##\d\$\$', ans["answer"])]
  571. for chunk_idx in chunk_idxs[:1]:
  572. if ans["reference"]["chunks"][chunk_idx]["img_id"]:
  573. try:
  574. bkt, nm = ans["reference"]["chunks"][chunk_idx]["img_id"].split("-")
  575. response = MINIO.get(bkt, nm)
  576. data_type_picture["url"] = base64.b64encode(response).decode('utf-8')
  577. data.append(data_type_picture)
  578. break
  579. except Exception as e:
  580. return server_error_response(e)
  581. response = {"code": 200, "msg": "success", "data": data}
  582. return response
  583. except Exception as e:
  584. return server_error_response(e)
  585. @manager.route('/retrieval', methods=['POST'])
  586. @validate_request("kb_id", "question")
  587. def retrieval():
  588. token = request.headers.get('Authorization').split()[1]
  589. objs = APIToken.query(token=token)
  590. if not objs:
  591. return get_json_result(
  592. data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
  593. req = request.json
  594. kb_id = req.get("kb_id")
  595. doc_ids = req.get("doc_ids", [])
  596. question = req.get("question")
  597. page = int(req.get("page", 1))
  598. size = int(req.get("size", 30))
  599. similarity_threshold = float(req.get("similarity_threshold", 0.2))
  600. vector_similarity_weight = float(req.get("vector_similarity_weight", 0.3))
  601. top = int(req.get("top_k", 1024))
  602. try:
  603. e, kb = KnowledgebaseService.get_by_id(kb_id)
  604. if not e:
  605. return get_data_error_result(retmsg="Knowledgebase not found!")
  606. embd_mdl = TenantLLMService.model_instance(
  607. kb.tenant_id, LLMType.EMBEDDING.value, llm_name=kb.embd_id)
  608. rerank_mdl = None
  609. if req.get("rerank_id"):
  610. rerank_mdl = TenantLLMService.model_instance(
  611. kb.tenant_id, LLMType.RERANK.value, llm_name=req["rerank_id"])
  612. if req.get("keyword", False):
  613. chat_mdl = TenantLLMService.model_instance(kb.tenant_id, LLMType.CHAT)
  614. question += keyword_extraction(chat_mdl, question)
  615. ranks = retrievaler.retrieval(question, embd_mdl, kb.tenant_id, [kb_id], page, size,
  616. similarity_threshold, vector_similarity_weight, top,
  617. doc_ids, rerank_mdl=rerank_mdl)
  618. for c in ranks["chunks"]:
  619. if "vector" in c:
  620. del c["vector"]
  621. return get_json_result(data=ranks)
  622. except Exception as e:
  623. if str(e).find("not_found") > 0:
  624. return get_json_result(data=False, retmsg=f'No chunk found! Check the chunk status please!',
  625. retcode=RetCode.DATA_ERROR)
  626. return server_error_response(e)