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 29KB

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