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

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735
  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 api.db.services.llm_service import TenantLLMService
  22. from flask_login import login_required, current_user
  23. from api.db import FileType, LLMType, ParserType, FileSource
  24. from api.db.db_models import APIToken, API4Conversation, Task, File
  25. from api.db.services import duplicate_name
  26. from api.db.services.api_service import APITokenService, API4ConversationService
  27. from api.db.services.dialog_service import DialogService, chat
  28. from api.db.services.document_service import DocumentService
  29. from api.db.services.file2document_service import File2DocumentService
  30. from api.db.services.file_service import FileService
  31. from api.db.services.knowledgebase_service import KnowledgebaseService
  32. from api.db.services.task_service import queue_tasks, TaskService
  33. from api.db.services.user_service import UserTenantService
  34. from api.settings import RetCode, retrievaler
  35. from api.utils import get_uuid, current_timestamp, datetime_format
  36. from api.utils.api_utils import server_error_response, get_data_error_result, get_json_result, validate_request
  37. from itsdangerous import URLSafeTimedSerializer
  38. from api.utils.file_utils import filename_type, thumbnail
  39. from rag.nlp import keyword_extraction
  40. from rag.utils.minio_conn import MINIO
  41. from api.db.services.canvas_service import CanvasTemplateService, UserCanvasService
  42. from agent.canvas import Canvas
  43. from functools import partial
  44. def generate_confirmation_token(tenent_id):
  45. serializer = URLSafeTimedSerializer(tenent_id)
  46. return "ragflow-" + serializer.dumps(get_uuid(), salt=tenent_id)[2:34]
  47. @manager.route('/new_token', methods=['POST'])
  48. @login_required
  49. def new_token():
  50. req = request.json
  51. try:
  52. tenants = UserTenantService.query(user_id=current_user.id)
  53. if not tenants:
  54. return get_data_error_result(retmsg="Tenant not found!")
  55. tenant_id = tenants[0].tenant_id
  56. obj = {"tenant_id": tenant_id, "token": generate_confirmation_token(tenant_id),
  57. "create_time": current_timestamp(),
  58. "create_date": datetime_format(datetime.now()),
  59. "update_time": None,
  60. "update_date": None
  61. }
  62. if req.get("canvas_id"):
  63. obj["dialog_id"] = req["canvas_id"]
  64. obj["source"] = "agent"
  65. else:
  66. obj["dialog_id"] = req["dialog_id"]
  67. if not APITokenService.save(**obj):
  68. return get_data_error_result(retmsg="Fail to new a dialog!")
  69. return get_json_result(data=obj)
  70. except Exception as e:
  71. return server_error_response(e)
  72. @manager.route('/token_list', methods=['GET'])
  73. @login_required
  74. def token_list():
  75. try:
  76. tenants = UserTenantService.query(user_id=current_user.id)
  77. if not tenants:
  78. return get_data_error_result(retmsg="Tenant not found!")
  79. id = request.args["dialog_id"] if "dialog_id" in request.args else request.args["canvas_id"]
  80. objs = APITokenService.query(tenant_id=tenants[0].tenant_id, dialog_id=id)
  81. return get_json_result(data=[o.to_dict() for o in objs])
  82. except Exception as e:
  83. return server_error_response(e)
  84. @manager.route('/rm', methods=['POST'])
  85. @validate_request("tokens", "tenant_id")
  86. @login_required
  87. def rm():
  88. req = request.json
  89. try:
  90. for token in req["tokens"]:
  91. APITokenService.filter_delete(
  92. [APIToken.tenant_id == req["tenant_id"], APIToken.token == token])
  93. return get_json_result(data=True)
  94. except Exception as e:
  95. return server_error_response(e)
  96. @manager.route('/stats', methods=['GET'])
  97. @login_required
  98. def stats():
  99. try:
  100. tenants = UserTenantService.query(user_id=current_user.id)
  101. if not tenants:
  102. return get_data_error_result(retmsg="Tenant not found!")
  103. objs = API4ConversationService.stats(
  104. tenants[0].tenant_id,
  105. request.args.get(
  106. "from_date",
  107. (datetime.now() -
  108. timedelta(
  109. days=7)).strftime("%Y-%m-%d 00:00:00")),
  110. request.args.get(
  111. "to_date",
  112. datetime.now().strftime("%Y-%m-%d %H:%M:%S")),
  113. "agent" if "canvas_id" in request.args else None)
  114. res = {
  115. "pv": [(o["dt"], o["pv"]) for o in objs],
  116. "uv": [(o["dt"], o["uv"]) for o in objs],
  117. "speed": [(o["dt"], float(o["tokens"]) / (float(o["duration"] + 0.1))) for o in objs],
  118. "tokens": [(o["dt"], float(o["tokens"]) / 1000.) for o in objs],
  119. "round": [(o["dt"], o["round"]) for o in objs],
  120. "thumb_up": [(o["dt"], o["thumb_up"]) for o in objs]
  121. }
  122. return get_json_result(data=res)
  123. except Exception as e:
  124. return server_error_response(e)
  125. @manager.route('/new_conversation', methods=['GET'])
  126. def set_conversation():
  127. token = request.headers.get('Authorization').split()[1]
  128. objs = APIToken.query(token=token)
  129. if not objs:
  130. return get_json_result(
  131. data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
  132. req = request.json
  133. try:
  134. if objs[0].source == "agent":
  135. e, c = UserCanvasService.get_by_id(objs[0].dialog_id)
  136. if not e:
  137. return server_error_response("canvas not found.")
  138. conv = {
  139. "id": get_uuid(),
  140. "dialog_id": c.id,
  141. "user_id": request.args.get("user_id", ""),
  142. "message": [{"role": "assistant", "content": "Hi there!"}],
  143. "source": "agent"
  144. }
  145. API4ConversationService.save(**conv)
  146. return get_json_result(data=conv)
  147. else:
  148. e, dia = DialogService.get_by_id(objs[0].dialog_id)
  149. if not e:
  150. return get_data_error_result(retmsg="Dialog not found")
  151. conv = {
  152. "id": get_uuid(),
  153. "dialog_id": dia.id,
  154. "user_id": request.args.get("user_id", ""),
  155. "message": [{"role": "assistant", "content": dia.prompt_config["prologue"]}]
  156. }
  157. API4ConversationService.save(**conv)
  158. return get_json_result(data=conv)
  159. except Exception as e:
  160. return server_error_response(e)
  161. @manager.route('/completion', methods=['POST'])
  162. @validate_request("conversation_id", "messages")
  163. def completion():
  164. token = request.headers.get('Authorization').split()[1]
  165. objs = APIToken.query(token=token)
  166. if not objs:
  167. return get_json_result(
  168. data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
  169. req = request.json
  170. e, conv = API4ConversationService.get_by_id(req["conversation_id"])
  171. if not e:
  172. return get_data_error_result(retmsg="Conversation not found!")
  173. if "quote" not in req: req["quote"] = False
  174. msg = []
  175. for m in req["messages"]:
  176. if m["role"] == "system":
  177. continue
  178. if m["role"] == "assistant" and not msg:
  179. continue
  180. msg.append({"role": m["role"], "content": m["content"]})
  181. def fillin_conv(ans):
  182. nonlocal conv
  183. if not conv.reference:
  184. conv.reference.append(ans["reference"])
  185. else:
  186. conv.reference[-1] = ans["reference"]
  187. conv.message[-1] = {"role": "assistant", "content": ans["answer"]}
  188. def rename_field(ans):
  189. reference = ans['reference']
  190. if not isinstance(reference, dict):
  191. return
  192. for chunk_i in reference.get('chunks', []):
  193. if 'docnm_kwd' in chunk_i:
  194. chunk_i['doc_name'] = chunk_i['docnm_kwd']
  195. chunk_i.pop('docnm_kwd')
  196. try:
  197. if conv.source == "agent":
  198. stream = req.get("stream", True)
  199. conv.message.append(msg[-1])
  200. e, cvs = UserCanvasService.get_by_id(conv.dialog_id)
  201. if not e:
  202. return server_error_response("canvas not found.")
  203. del req["conversation_id"]
  204. del req["messages"]
  205. if not isinstance(cvs.dsl, str):
  206. cvs.dsl = json.dumps(cvs.dsl, ensure_ascii=False)
  207. if not conv.reference:
  208. conv.reference = []
  209. conv.message.append({"role": "assistant", "content": ""})
  210. conv.reference.append({"chunks": [], "doc_aggs": []})
  211. final_ans = {"reference": [], "content": ""}
  212. canvas = Canvas(cvs.dsl, objs[0].tenant_id)
  213. canvas.messages.append(msg[-1])
  214. canvas.add_user_input(msg[-1]["content"])
  215. answer = canvas.run(stream=stream)
  216. assert answer is not None, "Nothing. Is it over?"
  217. if stream:
  218. assert isinstance(answer, partial), "Nothing. Is it over?"
  219. def sse():
  220. nonlocal answer, cvs, conv
  221. try:
  222. for ans in answer():
  223. for k in ans.keys():
  224. final_ans[k] = ans[k]
  225. ans = {"answer": ans["content"], "reference": ans.get("reference", [])}
  226. fillin_conv(ans)
  227. rename_field(ans)
  228. yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans},
  229. ensure_ascii=False) + "\n\n"
  230. canvas.messages.append({"role": "assistant", "content": final_ans["content"]})
  231. if final_ans.get("reference"):
  232. canvas.reference.append(final_ans["reference"])
  233. cvs.dsl = json.loads(str(canvas))
  234. API4ConversationService.append_message(conv.id, conv.to_dict())
  235. except Exception as e:
  236. yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
  237. "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
  238. ensure_ascii=False) + "\n\n"
  239. yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
  240. resp = Response(sse(), mimetype="text/event-stream")
  241. resp.headers.add_header("Cache-control", "no-cache")
  242. resp.headers.add_header("Connection", "keep-alive")
  243. resp.headers.add_header("X-Accel-Buffering", "no")
  244. resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
  245. return resp
  246. final_ans["content"] = "\n".join(answer["content"]) if "content" in answer else ""
  247. canvas.messages.append({"role": "assistant", "content": final_ans["content"]})
  248. if final_ans.get("reference"):
  249. canvas.reference.append(final_ans["reference"])
  250. cvs.dsl = json.loads(str(canvas))
  251. result = {"answer": final_ans["content"], "reference": final_ans.get("reference", [])}
  252. fillin_conv(result)
  253. API4ConversationService.append_message(conv.id, conv.to_dict())
  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_ids = 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. kbs = KnowledgebaseService.get_by_ids(kb_ids)
  604. embd_nms = list(set([kb.embd_id for kb in kbs]))
  605. if len(embd_nms) != 1:
  606. return get_json_result(
  607. data=False, retmsg='Knowledge bases use different embedding models or does not exist."', retcode=RetCode.AUTHENTICATION_ERROR)
  608. embd_mdl = TenantLLMService.model_instance(
  609. kbs[0].tenant_id, LLMType.EMBEDDING.value, llm_name=kbs[0].embd_id)
  610. rerank_mdl = None
  611. if req.get("rerank_id"):
  612. rerank_mdl = TenantLLMService.model_instance(
  613. kbs[0].tenant_id, LLMType.RERANK.value, llm_name=req["rerank_id"])
  614. if req.get("keyword", False):
  615. chat_mdl = TenantLLMService.model_instance(kbs[0].tenant_id, LLMType.CHAT)
  616. question += keyword_extraction(chat_mdl, question)
  617. ranks = retrievaler.retrieval(question, embd_mdl, kbs[0].tenant_id, kb_ids, 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)