|
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735 |
- #
- # Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
- #
- # Licensed under the Apache License, Version 2.0 (the "License");
- # you may not use this file except in compliance with the License.
- # You may obtain a copy of the License at
- #
- # http://www.apache.org/licenses/LICENSE-2.0
- #
- # Unless required by applicable law or agreed to in writing, software
- # distributed under the License is distributed on an "AS IS" BASIS,
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- # See the License for the specific language governing permissions and
- # limitations under the License.
- #
- import json
- import os
- import re
- from datetime import datetime, timedelta
- from flask import request, Response
- from api.db.services.llm_service import TenantLLMService
- from flask_login import login_required, current_user
-
- from api.db import FileType, LLMType, ParserType, FileSource
- from api.db.db_models import APIToken, API4Conversation, Task, File
- from api.db.services import duplicate_name
- from api.db.services.api_service import APITokenService, API4ConversationService
- from api.db.services.dialog_service import DialogService, chat
- from api.db.services.document_service import DocumentService
- from api.db.services.file2document_service import File2DocumentService
- from api.db.services.file_service import FileService
- from api.db.services.knowledgebase_service import KnowledgebaseService
- from api.db.services.task_service import queue_tasks, TaskService
- from api.db.services.user_service import UserTenantService
- from api.settings import RetCode, retrievaler
- from api.utils import get_uuid, current_timestamp, datetime_format
- from api.utils.api_utils import server_error_response, get_data_error_result, get_json_result, validate_request
- from itsdangerous import URLSafeTimedSerializer
-
- from api.utils.file_utils import filename_type, thumbnail
- from rag.nlp import keyword_extraction
- from rag.utils.minio_conn import MINIO
-
- from api.db.services.canvas_service import CanvasTemplateService, UserCanvasService
- from agent.canvas import Canvas
- from functools import partial
-
-
- def generate_confirmation_token(tenent_id):
- serializer = URLSafeTimedSerializer(tenent_id)
- return "ragflow-" + serializer.dumps(get_uuid(), salt=tenent_id)[2:34]
-
-
- @manager.route('/new_token', methods=['POST'])
- @login_required
- def new_token():
- req = request.json
- try:
- tenants = UserTenantService.query(user_id=current_user.id)
- if not tenants:
- return get_data_error_result(retmsg="Tenant not found!")
-
- tenant_id = tenants[0].tenant_id
- obj = {"tenant_id": tenant_id, "token": generate_confirmation_token(tenant_id),
- "create_time": current_timestamp(),
- "create_date": datetime_format(datetime.now()),
- "update_time": None,
- "update_date": None
- }
- if req.get("canvas_id"):
- obj["dialog_id"] = req["canvas_id"]
- obj["source"] = "agent"
- else:
- obj["dialog_id"] = req["dialog_id"]
-
- if not APITokenService.save(**obj):
- return get_data_error_result(retmsg="Fail to new a dialog!")
-
- return get_json_result(data=obj)
- except Exception as e:
- return server_error_response(e)
-
-
- @manager.route('/token_list', methods=['GET'])
- @login_required
- def token_list():
- try:
- tenants = UserTenantService.query(user_id=current_user.id)
- if not tenants:
- return get_data_error_result(retmsg="Tenant not found!")
-
- id = request.args["dialog_id"] if "dialog_id" in request.args else request.args["canvas_id"]
- objs = APITokenService.query(tenant_id=tenants[0].tenant_id, dialog_id=id)
- return get_json_result(data=[o.to_dict() for o in objs])
- except Exception as e:
- return server_error_response(e)
-
-
- @manager.route('/rm', methods=['POST'])
- @validate_request("tokens", "tenant_id")
- @login_required
- def rm():
- req = request.json
- try:
- for token in req["tokens"]:
- APITokenService.filter_delete(
- [APIToken.tenant_id == req["tenant_id"], APIToken.token == token])
- return get_json_result(data=True)
- except Exception as e:
- return server_error_response(e)
-
-
- @manager.route('/stats', methods=['GET'])
- @login_required
- def stats():
- try:
- tenants = UserTenantService.query(user_id=current_user.id)
- if not tenants:
- return get_data_error_result(retmsg="Tenant not found!")
- objs = API4ConversationService.stats(
- tenants[0].tenant_id,
- request.args.get(
- "from_date",
- (datetime.now() -
- timedelta(
- days=7)).strftime("%Y-%m-%d 00:00:00")),
- request.args.get(
- "to_date",
- datetime.now().strftime("%Y-%m-%d %H:%M:%S")),
- "agent" if "canvas_id" in request.args else None)
- res = {
- "pv": [(o["dt"], o["pv"]) for o in objs],
- "uv": [(o["dt"], o["uv"]) for o in objs],
- "speed": [(o["dt"], float(o["tokens"]) / (float(o["duration"] + 0.1))) for o in objs],
- "tokens": [(o["dt"], float(o["tokens"]) / 1000.) for o in objs],
- "round": [(o["dt"], o["round"]) for o in objs],
- "thumb_up": [(o["dt"], o["thumb_up"]) for o in objs]
- }
- return get_json_result(data=res)
- except Exception as e:
- return server_error_response(e)
-
-
- @manager.route('/new_conversation', methods=['GET'])
- def set_conversation():
- token = request.headers.get('Authorization').split()[1]
- objs = APIToken.query(token=token)
- if not objs:
- return get_json_result(
- data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
- req = request.json
- try:
- if objs[0].source == "agent":
- e, c = UserCanvasService.get_by_id(objs[0].dialog_id)
- if not e:
- return server_error_response("canvas not found.")
- conv = {
- "id": get_uuid(),
- "dialog_id": c.id,
- "user_id": request.args.get("user_id", ""),
- "message": [{"role": "assistant", "content": "Hi there!"}],
- "source": "agent"
- }
- API4ConversationService.save(**conv)
- return get_json_result(data=conv)
- else:
- e, dia = DialogService.get_by_id(objs[0].dialog_id)
- if not e:
- return get_data_error_result(retmsg="Dialog not found")
- conv = {
- "id": get_uuid(),
- "dialog_id": dia.id,
- "user_id": request.args.get("user_id", ""),
- "message": [{"role": "assistant", "content": dia.prompt_config["prologue"]}]
- }
- API4ConversationService.save(**conv)
- return get_json_result(data=conv)
- except Exception as e:
- return server_error_response(e)
-
-
- @manager.route('/completion', methods=['POST'])
- @validate_request("conversation_id", "messages")
- def completion():
- token = request.headers.get('Authorization').split()[1]
- objs = APIToken.query(token=token)
- if not objs:
- return get_json_result(
- data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
- req = request.json
- e, conv = API4ConversationService.get_by_id(req["conversation_id"])
- if not e:
- return get_data_error_result(retmsg="Conversation not found!")
- if "quote" not in req: req["quote"] = False
-
- msg = []
- for m in req["messages"]:
- if m["role"] == "system":
- continue
- if m["role"] == "assistant" and not msg:
- continue
- msg.append({"role": m["role"], "content": m["content"]})
-
- def fillin_conv(ans):
- nonlocal conv
- if not conv.reference:
- conv.reference.append(ans["reference"])
- else:
- conv.reference[-1] = ans["reference"]
- conv.message[-1] = {"role": "assistant", "content": ans["answer"]}
-
- def rename_field(ans):
- reference = ans['reference']
- if not isinstance(reference, dict):
- return
- for chunk_i in reference.get('chunks', []):
- if 'docnm_kwd' in chunk_i:
- chunk_i['doc_name'] = chunk_i['docnm_kwd']
- chunk_i.pop('docnm_kwd')
-
- try:
- if conv.source == "agent":
- stream = req.get("stream", True)
- conv.message.append(msg[-1])
- e, cvs = UserCanvasService.get_by_id(conv.dialog_id)
- if not e:
- return server_error_response("canvas not found.")
- del req["conversation_id"]
- del req["messages"]
-
- if not isinstance(cvs.dsl, str):
- cvs.dsl = json.dumps(cvs.dsl, ensure_ascii=False)
-
- if not conv.reference:
- conv.reference = []
- conv.message.append({"role": "assistant", "content": ""})
- conv.reference.append({"chunks": [], "doc_aggs": []})
-
- final_ans = {"reference": [], "content": ""}
- canvas = Canvas(cvs.dsl, objs[0].tenant_id)
-
- canvas.messages.append(msg[-1])
- canvas.add_user_input(msg[-1]["content"])
- answer = canvas.run(stream=stream)
-
- assert answer is not None, "Nothing. Is it over?"
-
- if stream:
- assert isinstance(answer, partial), "Nothing. Is it over?"
-
- def sse():
- nonlocal answer, cvs, conv
- try:
- for ans in answer():
- for k in ans.keys():
- final_ans[k] = ans[k]
- ans = {"answer": ans["content"], "reference": ans.get("reference", [])}
- fillin_conv(ans)
- rename_field(ans)
- yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans},
- ensure_ascii=False) + "\n\n"
-
- canvas.messages.append({"role": "assistant", "content": final_ans["content"]})
- if final_ans.get("reference"):
- canvas.reference.append(final_ans["reference"])
- cvs.dsl = json.loads(str(canvas))
- API4ConversationService.append_message(conv.id, conv.to_dict())
- except Exception as e:
- yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
- "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
- ensure_ascii=False) + "\n\n"
- yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
-
- resp = Response(sse(), mimetype="text/event-stream")
- resp.headers.add_header("Cache-control", "no-cache")
- resp.headers.add_header("Connection", "keep-alive")
- resp.headers.add_header("X-Accel-Buffering", "no")
- resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
- return resp
-
- final_ans["content"] = "\n".join(answer["content"]) if "content" in answer else ""
- canvas.messages.append({"role": "assistant", "content": final_ans["content"]})
- if final_ans.get("reference"):
- canvas.reference.append(final_ans["reference"])
- cvs.dsl = json.loads(str(canvas))
-
- result = {"answer": final_ans["content"], "reference": final_ans.get("reference", [])}
- fillin_conv(result)
- API4ConversationService.append_message(conv.id, conv.to_dict())
- rename_field(result)
- return get_json_result(data=result)
-
- #******************For dialog******************
- conv.message.append(msg[-1])
- e, dia = DialogService.get_by_id(conv.dialog_id)
- if not e:
- return get_data_error_result(retmsg="Dialog not found!")
- del req["conversation_id"]
- del req["messages"]
-
- if not conv.reference:
- conv.reference = []
- conv.message.append({"role": "assistant", "content": ""})
- conv.reference.append({"chunks": [], "doc_aggs": []})
-
- def stream():
- nonlocal dia, msg, req, conv
- try:
- for ans in chat(dia, msg, True, **req):
- fillin_conv(ans)
- rename_field(ans)
- yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": ans},
- ensure_ascii=False) + "\n\n"
- API4ConversationService.append_message(conv.id, conv.to_dict())
- except Exception as e:
- yield "data:" + json.dumps({"retcode": 500, "retmsg": str(e),
- "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
- ensure_ascii=False) + "\n\n"
- yield "data:" + json.dumps({"retcode": 0, "retmsg": "", "data": True}, ensure_ascii=False) + "\n\n"
-
- if req.get("stream", True):
- resp = Response(stream(), mimetype="text/event-stream")
- resp.headers.add_header("Cache-control", "no-cache")
- resp.headers.add_header("Connection", "keep-alive")
- resp.headers.add_header("X-Accel-Buffering", "no")
- resp.headers.add_header("Content-Type", "text/event-stream; charset=utf-8")
- return resp
-
- answer = None
- for ans in chat(dia, msg, **req):
- answer = ans
- fillin_conv(ans)
- API4ConversationService.append_message(conv.id, conv.to_dict())
- break
- rename_field(answer)
- return get_json_result(data=answer)
-
- except Exception as e:
- return server_error_response(e)
-
-
- @manager.route('/conversation/<conversation_id>', methods=['GET'])
- # @login_required
- def get(conversation_id):
- try:
- e, conv = API4ConversationService.get_by_id(conversation_id)
- if not e:
- return get_data_error_result(retmsg="Conversation not found!")
-
- conv = conv.to_dict()
- for referenct_i in conv['reference']:
- if referenct_i is None or len(referenct_i) == 0:
- continue
- for chunk_i in referenct_i['chunks']:
- if 'docnm_kwd' in chunk_i.keys():
- chunk_i['doc_name'] = chunk_i['docnm_kwd']
- chunk_i.pop('docnm_kwd')
- return get_json_result(data=conv)
- except Exception as e:
- return server_error_response(e)
-
-
- @manager.route('/document/upload', methods=['POST'])
- @validate_request("kb_name")
- def upload():
- token = request.headers.get('Authorization').split()[1]
- objs = APIToken.query(token=token)
- if not objs:
- return get_json_result(
- data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
-
- kb_name = request.form.get("kb_name").strip()
- tenant_id = objs[0].tenant_id
-
- try:
- e, kb = KnowledgebaseService.get_by_name(kb_name, tenant_id)
- if not e:
- return get_data_error_result(
- retmsg="Can't find this knowledgebase!")
- kb_id = kb.id
- except Exception as e:
- return server_error_response(e)
-
- if 'file' not in request.files:
- return get_json_result(
- data=False, retmsg='No file part!', retcode=RetCode.ARGUMENT_ERROR)
-
- file = request.files['file']
- if file.filename == '':
- return get_json_result(
- data=False, retmsg='No file selected!', retcode=RetCode.ARGUMENT_ERROR)
-
- root_folder = FileService.get_root_folder(tenant_id)
- pf_id = root_folder["id"]
- FileService.init_knowledgebase_docs(pf_id, tenant_id)
- kb_root_folder = FileService.get_kb_folder(tenant_id)
- kb_folder = FileService.new_a_file_from_kb(kb.tenant_id, kb.name, kb_root_folder["id"])
-
- try:
- if DocumentService.get_doc_count(kb.tenant_id) >= int(os.environ.get('MAX_FILE_NUM_PER_USER', 8192)):
- return get_data_error_result(
- retmsg="Exceed the maximum file number of a free user!")
-
- filename = duplicate_name(
- DocumentService.query,
- name=file.filename,
- kb_id=kb_id)
- filetype = filename_type(filename)
- if not filetype:
- return get_data_error_result(
- retmsg="This type of file has not been supported yet!")
-
- location = filename
- while MINIO.obj_exist(kb_id, location):
- location += "_"
- blob = request.files['file'].read()
- MINIO.put(kb_id, location, blob)
- doc = {
- "id": get_uuid(),
- "kb_id": kb.id,
- "parser_id": kb.parser_id,
- "parser_config": kb.parser_config,
- "created_by": kb.tenant_id,
- "type": filetype,
- "name": filename,
- "location": location,
- "size": len(blob),
- "thumbnail": thumbnail(filename, blob)
- }
-
- form_data = request.form
- if "parser_id" in form_data.keys():
- if request.form.get("parser_id").strip() in list(vars(ParserType).values())[1:-3]:
- doc["parser_id"] = request.form.get("parser_id").strip()
- if doc["type"] == FileType.VISUAL:
- doc["parser_id"] = ParserType.PICTURE.value
- if doc["type"] == FileType.AURAL:
- doc["parser_id"] = ParserType.AUDIO.value
- if re.search(r"\.(ppt|pptx|pages)$", filename):
- doc["parser_id"] = ParserType.PRESENTATION.value
-
- doc_result = DocumentService.insert(doc)
- FileService.add_file_from_kb(doc, kb_folder["id"], kb.tenant_id)
- except Exception as e:
- return server_error_response(e)
-
- if "run" in form_data.keys():
- if request.form.get("run").strip() == "1":
- try:
- info = {"run": 1, "progress": 0}
- info["progress_msg"] = ""
- info["chunk_num"] = 0
- info["token_num"] = 0
- DocumentService.update_by_id(doc["id"], info)
- # if str(req["run"]) == TaskStatus.CANCEL.value:
- tenant_id = DocumentService.get_tenant_id(doc["id"])
- if not tenant_id:
- return get_data_error_result(retmsg="Tenant not found!")
-
- # e, doc = DocumentService.get_by_id(doc["id"])
- TaskService.filter_delete([Task.doc_id == doc["id"]])
- e, doc = DocumentService.get_by_id(doc["id"])
- doc = doc.to_dict()
- doc["tenant_id"] = tenant_id
- bucket, name = File2DocumentService.get_minio_address(doc_id=doc["id"])
- queue_tasks(doc, bucket, name)
- except Exception as e:
- return server_error_response(e)
-
- return get_json_result(data=doc_result.to_json())
-
-
- @manager.route('/list_chunks', methods=['POST'])
- # @login_required
- def list_chunks():
- token = request.headers.get('Authorization').split()[1]
- objs = APIToken.query(token=token)
- if not objs:
- return get_json_result(
- data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
-
- req = request.json
-
- try:
- if "doc_name" in req.keys():
- tenant_id = DocumentService.get_tenant_id_by_name(req['doc_name'])
- doc_id = DocumentService.get_doc_id_by_doc_name(req['doc_name'])
-
- elif "doc_id" in req.keys():
- tenant_id = DocumentService.get_tenant_id(req['doc_id'])
- doc_id = req['doc_id']
- else:
- return get_json_result(
- data=False, retmsg="Can't find doc_name or doc_id"
- )
-
- res = retrievaler.chunk_list(doc_id=doc_id, tenant_id=tenant_id)
- res = [
- {
- "content": res_item["content_with_weight"],
- "doc_name": res_item["docnm_kwd"],
- "img_id": res_item["img_id"]
- } for res_item in res
- ]
-
- except Exception as e:
- return server_error_response(e)
-
- return get_json_result(data=res)
-
-
- @manager.route('/list_kb_docs', methods=['POST'])
- # @login_required
- def list_kb_docs():
- token = request.headers.get('Authorization').split()[1]
- objs = APIToken.query(token=token)
- if not objs:
- return get_json_result(
- data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
-
- req = request.json
- tenant_id = objs[0].tenant_id
- kb_name = req.get("kb_name", "").strip()
-
- try:
- e, kb = KnowledgebaseService.get_by_name(kb_name, tenant_id)
- if not e:
- return get_data_error_result(
- retmsg="Can't find this knowledgebase!")
- kb_id = kb.id
-
- except Exception as e:
- return server_error_response(e)
-
- page_number = int(req.get("page", 1))
- items_per_page = int(req.get("page_size", 15))
- orderby = req.get("orderby", "create_time")
- desc = req.get("desc", True)
- keywords = req.get("keywords", "")
-
- try:
- docs, tol = DocumentService.get_by_kb_id(
- kb_id, page_number, items_per_page, orderby, desc, keywords)
- docs = [{"doc_id": doc['id'], "doc_name": doc['name']} for doc in docs]
-
- return get_json_result(data={"total": tol, "docs": docs})
-
- except Exception as e:
- return server_error_response(e)
-
-
- @manager.route('/document', methods=['DELETE'])
- # @login_required
- def document_rm():
- token = request.headers.get('Authorization').split()[1]
- objs = APIToken.query(token=token)
- if not objs:
- return get_json_result(
- data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
-
- tenant_id = objs[0].tenant_id
- req = request.json
- doc_ids = []
- try:
- doc_ids = [DocumentService.get_doc_id_by_doc_name(doc_name) for doc_name in req.get("doc_names", [])]
- for doc_id in req.get("doc_ids", []):
- if doc_id not in doc_ids:
- doc_ids.append(doc_id)
-
- if not doc_ids:
- return get_json_result(
- data=False, retmsg="Can't find doc_names or doc_ids"
- )
-
- except Exception as e:
- return server_error_response(e)
-
- root_folder = FileService.get_root_folder(tenant_id)
- pf_id = root_folder["id"]
- FileService.init_knowledgebase_docs(pf_id, tenant_id)
-
- errors = ""
- for doc_id in doc_ids:
- try:
- e, doc = DocumentService.get_by_id(doc_id)
- if not e:
- return get_data_error_result(retmsg="Document not found!")
- tenant_id = DocumentService.get_tenant_id(doc_id)
- if not tenant_id:
- return get_data_error_result(retmsg="Tenant not found!")
-
- b, n = File2DocumentService.get_minio_address(doc_id=doc_id)
-
- if not DocumentService.remove_document(doc, tenant_id):
- return get_data_error_result(
- retmsg="Database error (Document removal)!")
-
- f2d = File2DocumentService.get_by_document_id(doc_id)
- FileService.filter_delete([File.source_type == FileSource.KNOWLEDGEBASE, File.id == f2d[0].file_id])
- File2DocumentService.delete_by_document_id(doc_id)
-
- MINIO.rm(b, n)
- except Exception as e:
- errors += str(e)
-
- if errors:
- return get_json_result(data=False, retmsg=errors, retcode=RetCode.SERVER_ERROR)
-
- return get_json_result(data=True)
-
-
- @manager.route('/completion_aibotk', methods=['POST'])
- @validate_request("Authorization", "conversation_id", "word")
- def completion_faq():
- import base64
- req = request.json
-
- token = req["Authorization"]
- objs = APIToken.query(token=token)
- if not objs:
- return get_json_result(
- data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
-
- e, conv = API4ConversationService.get_by_id(req["conversation_id"])
- if not e:
- return get_data_error_result(retmsg="Conversation not found!")
- if "quote" not in req: req["quote"] = True
-
- msg = []
- msg.append({"role": "user", "content": req["word"]})
-
- try:
- conv.message.append(msg[-1])
- e, dia = DialogService.get_by_id(conv.dialog_id)
- if not e:
- return get_data_error_result(retmsg="Dialog not found!")
- del req["conversation_id"]
-
- if not conv.reference:
- conv.reference = []
- conv.message.append({"role": "assistant", "content": ""})
- conv.reference.append({"chunks": [], "doc_aggs": []})
-
- def fillin_conv(ans):
- nonlocal conv
- if not conv.reference:
- conv.reference.append(ans["reference"])
- else:
- conv.reference[-1] = ans["reference"]
- conv.message[-1] = {"role": "assistant", "content": ans["answer"]}
-
- data_type_picture = {
- "type": 3,
- "url": "base64 content"
- }
- data = [
- {
- "type": 1,
- "content": ""
- }
- ]
- ans = ""
- for a in chat(dia, msg, stream=False, **req):
- ans = a
- break
- data[0]["content"] += re.sub(r'##\d\$\$', '', ans["answer"])
- fillin_conv(ans)
- API4ConversationService.append_message(conv.id, conv.to_dict())
-
- chunk_idxs = [int(match[2]) for match in re.findall(r'##\d\$\$', ans["answer"])]
- for chunk_idx in chunk_idxs[:1]:
- if ans["reference"]["chunks"][chunk_idx]["img_id"]:
- try:
- bkt, nm = ans["reference"]["chunks"][chunk_idx]["img_id"].split("-")
- response = MINIO.get(bkt, nm)
- data_type_picture["url"] = base64.b64encode(response).decode('utf-8')
- data.append(data_type_picture)
- break
- except Exception as e:
- return server_error_response(e)
-
- response = {"code": 200, "msg": "success", "data": data}
- return response
-
- except Exception as e:
- return server_error_response(e)
-
-
- @manager.route('/retrieval', methods=['POST'])
- @validate_request("kb_id", "question")
- def retrieval():
- token = request.headers.get('Authorization').split()[1]
- objs = APIToken.query(token=token)
- if not objs:
- return get_json_result(
- data=False, retmsg='Token is not valid!"', retcode=RetCode.AUTHENTICATION_ERROR)
-
- req = request.json
- kb_ids = req.get("kb_id",[])
- doc_ids = req.get("doc_ids", [])
- question = req.get("question")
- page = int(req.get("page", 1))
- size = int(req.get("size", 30))
- similarity_threshold = float(req.get("similarity_threshold", 0.2))
- vector_similarity_weight = float(req.get("vector_similarity_weight", 0.3))
- top = int(req.get("top_k", 1024))
-
- try:
- kbs = KnowledgebaseService.get_by_ids(kb_ids)
- embd_nms = list(set([kb.embd_id for kb in kbs]))
- if len(embd_nms) != 1:
- return get_json_result(
- data=False, retmsg='Knowledge bases use different embedding models or does not exist."', retcode=RetCode.AUTHENTICATION_ERROR)
-
- embd_mdl = TenantLLMService.model_instance(
- kbs[0].tenant_id, LLMType.EMBEDDING.value, llm_name=kbs[0].embd_id)
- rerank_mdl = None
- if req.get("rerank_id"):
- rerank_mdl = TenantLLMService.model_instance(
- kbs[0].tenant_id, LLMType.RERANK.value, llm_name=req["rerank_id"])
- if req.get("keyword", False):
- chat_mdl = TenantLLMService.model_instance(kbs[0].tenant_id, LLMType.CHAT)
- question += keyword_extraction(chat_mdl, question)
- ranks = retrievaler.retrieval(question, embd_mdl, kbs[0].tenant_id, kb_ids, page, size,
- similarity_threshold, vector_similarity_weight, top,
- doc_ids, rerank_mdl=rerank_mdl)
- for c in ranks["chunks"]:
- if "vector" in c:
- del c["vector"]
- return get_json_result(data=ranks)
- except Exception as e:
- if str(e).find("not_found") > 0:
- return get_json_result(data=False, retmsg=f'No chunk found! Check the chunk status please!',
- retcode=RetCode.DATA_ERROR)
- return server_error_response(e)
|