| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708 | 
							- import pathlib
 - import re
 - import datetime
 - import json
 - import traceback
 - 
 - from flask import request
 - from flask_login import login_required, current_user
 - from elasticsearch_dsl import Q
 - 
 - from rag.app.qa import rmPrefix, beAdoc
 - from rag.nlp import search, rag_tokenizer, keyword_extraction
 - from rag.utils.es_conn import ELASTICSEARCH
 - from rag.utils import rmSpace
 - from api.db import LLMType, ParserType
 - from api.db.services.knowledgebase_service import KnowledgebaseService
 - from api.db.services.llm_service import TenantLLMService
 - from api.db.services.user_service import UserTenantService
 - from api.utils.api_utils import server_error_response, get_data_error_result, validate_request
 - from api.db.services.document_service import DocumentService
 - from api.settings import RetCode, retrievaler, kg_retrievaler
 - from api.utils.api_utils import get_json_result
 - import hashlib
 - import re
 - from api.utils.api_utils import get_json_result, token_required, get_data_error_result
 - 
 - from api.db.db_models import Task, File
 - 
 - from api.db.services.task_service import TaskService, queue_tasks
 - from api.db.services.user_service import TenantService, UserTenantService
 - 
 - from api.utils.api_utils import server_error_response, get_data_error_result, validate_request
 - 
 - from api.utils.api_utils import get_json_result
 - 
 - from functools import partial
 - from io import BytesIO
 - 
 - from elasticsearch_dsl import Q
 - from flask import request, send_file
 - from flask_login import login_required
 - 
 - from api.db import FileSource, TaskStatus, FileType
 - from api.db.db_models import File
 - 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.settings import RetCode, retrievaler
 - from api.utils.api_utils import construct_json_result, construct_error_response
 - from rag.app import book, laws, manual, naive, one, paper, presentation, qa, resume, table, picture, audio, email
 - from rag.nlp import search
 - from rag.utils import rmSpace
 - from rag.utils.es_conn import ELASTICSEARCH
 - from rag.utils.storage_factory import STORAGE_IMPL
 - 
 - MAXIMUM_OF_UPLOADING_FILES = 256
 - 
 - MAXIMUM_OF_UPLOADING_FILES = 256
 - 
 - 
 - @manager.route('/dataset/<dataset_id>/documents/upload', methods=['POST'])
 - @token_required
 - def upload(dataset_id, tenant_id):
 -     if 'file' not in request.files:
 -         return get_json_result(
 -             data=False, retmsg='No file part!', retcode=RetCode.ARGUMENT_ERROR)
 -     file_objs = request.files.getlist('file')
 -     for file_obj in file_objs:
 -         if file_obj.filename == '':
 -             return get_json_result(
 -                 data=False, retmsg='No file selected!', retcode=RetCode.ARGUMENT_ERROR)
 -     e, kb = KnowledgebaseService.get_by_id(dataset_id)
 -     if not e:
 -         raise LookupError(f"Can't find the knowledgebase with ID {dataset_id}!")
 -     err, _ = FileService.upload_document(kb, file_objs, tenant_id)
 -     if err:
 -         return get_json_result(
 -             data=False, retmsg="\n".join(err), retcode=RetCode.SERVER_ERROR)
 -     return get_json_result(data=True)
 - 
 - 
 - @manager.route('/infos', methods=['GET'])
 - @token_required
 - def docinfos(tenant_id):
 -     req = request.args
 -     if "id" not in req and "name" not in req:
 -         return get_data_error_result(
 -             retmsg="Id or name should be provided")
 -     doc_id=None
 -     if "id" in req:
 -         doc_id = req["id"]
 -     if "name" in req:
 -         doc_name = req["name"]
 -         doc_id = DocumentService.get_doc_id_by_doc_name(doc_name)
 -     e, doc = DocumentService.get_by_id(doc_id)
 -     #rename key's name
 -     key_mapping = {
 -         "chunk_num": "chunk_count",
 -         "kb_id": "knowledgebase_id",
 -         "token_num": "token_count",
 -     }
 -     renamed_doc = {}
 -     for key, value in doc.to_dict().items():
 -         new_key = key_mapping.get(key, key)
 -         renamed_doc[new_key] = value
 - 
 -     return get_json_result(data=renamed_doc)
 - 
 - 
 - @manager.route('/save', methods=['POST'])
 - @token_required
 - def save_doc(tenant_id):
 -     req = request.json
 -     #get doc by id or name
 -     doc_id = None
 -     if "id" in req:
 -         doc_id = req["id"]
 -     elif "name" in req:
 -         doc_name = req["name"]
 -         doc_id = DocumentService.get_doc_id_by_doc_name(doc_name)
 -     if not doc_id:
 -         return get_json_result(retcode=400, retmsg="Document ID or name is required")
 -     e, doc = DocumentService.get_by_id(doc_id)
 -     if not e:
 -         return get_data_error_result(retmsg="Document not found!")
 -     #other value can't be changed
 -     if "chunk_num" in req:
 -         if req["chunk_num"] != doc.chunk_num:
 -             return get_data_error_result(
 -                 retmsg="Can't change chunk_count.")
 -     if "progress" in req:
 -         if req['progress'] != doc.progress:
 -             return get_data_error_result(
 -                 retmsg="Can't change progress.")
 -     #change name or parse_method
 -     if "name" in req and req["name"] != doc.name:
 -         try:
 -             if pathlib.Path(req["name"].lower()).suffix != pathlib.Path(
 -                     doc.name.lower()).suffix:
 -                 return get_json_result(
 -                     data=False,
 -                     retmsg="The extension of file can't be changed",
 -                     retcode=RetCode.ARGUMENT_ERROR)
 -             for d in DocumentService.query(name=req["name"], kb_id=doc.kb_id):
 -                 if d.name == req["name"]:
 -                     return get_data_error_result(
 -                         retmsg="Duplicated document name in the same knowledgebase.")
 - 
 -             if not DocumentService.update_by_id(
 -                     doc_id, {"name": req["name"]}):
 -                 return get_data_error_result(
 -                     retmsg="Database error (Document rename)!")
 - 
 -             informs = File2DocumentService.get_by_document_id(doc_id)
 -             if informs:
 -                 e, file = FileService.get_by_id(informs[0].file_id)
 -                 FileService.update_by_id(file.id, {"name": req["name"]})
 -         except Exception as e:
 -             return server_error_response(e)
 -     if "parser_id" in req:
 -         try:
 -             if doc.parser_id.lower() == req["parser_id"].lower():
 -                 if "parser_config" in req:
 -                     if req["parser_config"] == doc.parser_config:
 -                         return get_json_result(data=True)
 -                 else:
 -                     return get_json_result(data=True)
 - 
 -             if doc.type == FileType.VISUAL or re.search(
 -                     r"\.(ppt|pptx|pages)$", doc.name):
 -                 return get_data_error_result(retmsg="Not supported yet!")
 - 
 -             e = DocumentService.update_by_id(doc.id,
 -                                              {"parser_id": req["parser_id"], "progress": 0, "progress_msg": "",
 -                                               "run": TaskStatus.UNSTART.value})
 -             if not e:
 -                 return get_data_error_result(retmsg="Document not found!")
 -             if "parser_config" in req:
 -                 DocumentService.update_parser_config(doc.id, req["parser_config"])
 -             if doc.token_num > 0:
 -                 e = DocumentService.increment_chunk_num(doc.id, doc.kb_id, doc.token_num * -1, doc.chunk_num * -1,
 -                                                         doc.process_duation * -1)
 -                 if not e:
 -                     return get_data_error_result(retmsg="Document not found!")
 -                 tenant_id = DocumentService.get_tenant_id(req["doc_id"])
 -                 if not tenant_id:
 -                     return get_data_error_result(retmsg="Tenant not found!")
 -                 ELASTICSEARCH.deleteByQuery(
 -                     Q("match", doc_id=doc.id), idxnm=search.index_name(tenant_id))
 -         except Exception as e:
 -             return server_error_response(e)
 -     return get_json_result(data=True)
 - 
 - 
 - 
 - @manager.route('/change_parser', methods=['POST'])
 - @token_required
 - def change_parser(tenant_id):
 -     req = request.json
 -     try:
 -         e, doc = DocumentService.get_by_id(req["doc_id"])
 -         if not e:
 -             return get_data_error_result(retmsg="Document not found!")
 -         if doc.parser_id.lower() == req["parser_id"].lower():
 -             if "parser_config" in req:
 -                 if req["parser_config"] == doc.parser_config:
 -                     return get_json_result(data=True)
 -             else:
 -                 return get_json_result(data=True)
 - 
 -         if doc.type == FileType.VISUAL or re.search(
 -                 r"\.(ppt|pptx|pages)$", doc.name):
 -             return get_data_error_result(retmsg="Not supported yet!")
 - 
 -         e = DocumentService.update_by_id(doc.id,
 -                                          {"parser_id": req["parser_id"], "progress": 0, "progress_msg": "",
 -                                           "run": TaskStatus.UNSTART.value})
 -         if not e:
 -             return get_data_error_result(retmsg="Document not found!")
 -         if "parser_config" in req:
 -             DocumentService.update_parser_config(doc.id, req["parser_config"])
 -         if doc.token_num > 0:
 -             e = DocumentService.increment_chunk_num(doc.id, doc.kb_id, doc.token_num * -1, doc.chunk_num * -1,
 -                                                     doc.process_duation * -1)
 -             if not e:
 -                 return get_data_error_result(retmsg="Document not found!")
 -             tenant_id = DocumentService.get_tenant_id(req["doc_id"])
 -             if not tenant_id:
 -                 return get_data_error_result(retmsg="Tenant not found!")
 -             ELASTICSEARCH.deleteByQuery(
 -                 Q("match", doc_id=doc.id), idxnm=search.index_name(tenant_id))
 - 
 -         return get_json_result(data=True)
 -     except Exception as e:
 -         return server_error_response(e)
 - 
 - @manager.route('/rename', methods=['POST'])
 - @login_required
 - @validate_request("doc_id", "name")
 - def rename():
 -     req = request.json
 -     try:
 -         e, doc = DocumentService.get_by_id(req["doc_id"])
 -         if not e:
 -             return get_data_error_result(retmsg="Document not found!")
 -         if pathlib.Path(req["name"].lower()).suffix != pathlib.Path(
 -                 doc.name.lower()).suffix:
 -             return get_json_result(
 -                 data=False,
 -                 retmsg="The extension of file can't be changed",
 -                 retcode=RetCode.ARGUMENT_ERROR)
 -         for d in DocumentService.query(name=req["name"], kb_id=doc.kb_id):
 -             if d.name == req["name"]:
 -                 return get_data_error_result(
 -                     retmsg="Duplicated document name in the same knowledgebase.")
 - 
 -         if not DocumentService.update_by_id(
 -                 req["doc_id"], {"name": req["name"]}):
 -             return get_data_error_result(
 -                 retmsg="Database error (Document rename)!")
 -           
 -         informs = File2DocumentService.get_by_document_id(req["doc_id"])
 -         if informs:
 -             e, file = FileService.get_by_id(informs[0].file_id)
 -             FileService.update_by_id(file.id, {"name": req["name"]})
 - 
 -         return get_json_result(data=True)
 -     except Exception as e:
 -         return server_error_response(e)
 - 
 - 
 - @manager.route("/<document_id>", methods=["GET"])
 - @token_required
 - def download_document(dataset_id, document_id,tenant_id):
 -     try:
 -         # Check whether there is this document
 -         exist, document = DocumentService.get_by_id(document_id)
 -         if not exist:
 -             return construct_json_result(message=f"This document '{document_id}' cannot be found!",
 -                                          code=RetCode.ARGUMENT_ERROR)
 - 
 -         # The process of downloading
 -         doc_id, doc_location = File2DocumentService.get_minio_address(doc_id=document_id)  # minio address
 -         file_stream = STORAGE_IMPL.get(doc_id, doc_location)
 -         if not file_stream:
 -             return construct_json_result(message="This file is empty.", code=RetCode.DATA_ERROR)
 - 
 -         file = BytesIO(file_stream)
 - 
 -         # Use send_file with a proper filename and MIME type
 -         return send_file(
 -             file,
 -             as_attachment=True,
 -             download_name=document.name,
 -             mimetype='application/octet-stream'  # Set a default MIME type
 -         )
 - 
 -     # Error
 -     except Exception as e:
 -         return construct_error_response(e)
 - 
 - 
 - @manager.route('/dataset/<dataset_id>/documents', methods=['GET'])
 - @token_required
 - def list_docs(dataset_id, tenant_id):
 -     kb_id = request.args.get("kb_id")
 -     if not kb_id:
 -         return get_json_result(
 -             data=False, retmsg='Lack of "KB ID"', retcode=RetCode.ARGUMENT_ERROR)
 -     tenants = UserTenantService.query(user_id=tenant_id)
 -     for tenant in tenants:
 -         if KnowledgebaseService.query(
 -                 tenant_id=tenant.tenant_id, id=kb_id):
 -             break
 -     else:
 -         return get_json_result(
 -             data=False, retmsg=f'Only owner of knowledgebase authorized for this operation.',
 -             retcode=RetCode.OPERATING_ERROR)
 -     keywords = request.args.get("keywords", "")
 - 
 -     page_number = int(request.args.get("page", 1))
 -     items_per_page = int(request.args.get("page_size", 15))
 -     orderby = request.args.get("orderby", "create_time")
 -     desc = request.args.get("desc", True)
 -     try:
 -         docs, tol = DocumentService.get_by_kb_id(
 -             kb_id, page_number, items_per_page, orderby, desc, keywords)
 - 
 -         # rename key's name
 -         renamed_doc_list = []
 -         for doc in docs:
 -             key_mapping = {
 -                 "chunk_num": "chunk_count",
 -                 "kb_id": "knowledgebase_id",
 -                 "token_num": "token_count",
 -             }
 -             renamed_doc = {}
 -             for key, value in doc.items():
 -                 new_key = key_mapping.get(key, key)
 -                 renamed_doc[new_key] = value
 -             renamed_doc_list.append(renamed_doc)
 -         return get_json_result(data={"total": tol, "docs": renamed_doc_list})
 -     except Exception as e:
 -         return server_error_response(e)
 - 
 - 
 - @manager.route('/delete', methods=['DELETE'])
 - @token_required
 - def rm(tenant_id):
 -     req = request.args
 -     if "doc_id" not in req:
 -         return get_data_error_result(
 -             retmsg="doc_id is required")
 -     doc_ids = req["doc_id"]
 -     if isinstance(doc_ids, str): doc_ids = [doc_ids]
 -     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)
 - 
 -             STORAGE_IMPL.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, retmsg="success")
 - 
 - @manager.route("/<document_id>/status", methods=["GET"])
 - @token_required
 - def show_parsing_status(tenant_id, document_id):
 -     try:
 -         # valid document
 -         exist, _ = DocumentService.get_by_id(document_id)
 -         if not exist:
 -             return construct_json_result(code=RetCode.DATA_ERROR,
 -                                          message=f"This document: '{document_id}' is not a valid document.")
 - 
 -         _, doc = DocumentService.get_by_id(document_id)  # get doc object
 -         doc_attributes = doc.to_dict()
 - 
 -         return construct_json_result(
 -             data={"progress": doc_attributes["progress"], "status": TaskStatus(doc_attributes["status"]).name},
 -             code=RetCode.SUCCESS
 -         )
 -     except Exception as e:
 -         return construct_error_response(e)
 - 
 - 
 - 
 - @manager.route('/run', methods=['POST'])
 - @token_required
 - def run(tenant_id):
 -     req = request.json
 -     try:
 -         for id in req["doc_ids"]:
 -             info = {"run": str(req["run"]), "progress": 0}
 -             if str(req["run"]) == TaskStatus.RUNNING.value:
 -                 info["progress_msg"] = ""
 -                 info["chunk_num"] = 0
 -                 info["token_num"] = 0
 -             DocumentService.update_by_id(id, info)
 -             # if str(req["run"]) == TaskStatus.CANCEL.value:
 -             tenant_id = DocumentService.get_tenant_id(id)
 -             if not tenant_id:
 -                 return get_data_error_result(retmsg="Tenant not found!")
 -             ELASTICSEARCH.deleteByQuery(
 -                 Q("match", doc_id=id), idxnm=search.index_name(tenant_id))
 - 
 -             if str(req["run"]) == TaskStatus.RUNNING.value:
 -                 TaskService.filter_delete([Task.doc_id == id])
 -                 e, doc = DocumentService.get_by_id(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)
 - 
 -         return get_json_result(data=True)
 -     except Exception as e:
 -         return server_error_response(e)
 - 
 - 
 - @manager.route('/chunk/list', methods=['POST'])
 - @token_required
 - @validate_request("doc_id")
 - def list_chunk(tenant_id):
 -     req = request.json
 -     doc_id = req["doc_id"]
 -     page = int(req.get("page", 1))
 -     size = int(req.get("size", 30))
 -     question = req.get("keywords", "")
 -     try:
 -         tenant_id = DocumentService.get_tenant_id(req["doc_id"])
 -         if not tenant_id:
 -             return get_data_error_result(retmsg="Tenant not found!")
 -         e, doc = DocumentService.get_by_id(doc_id)
 -         if not e:
 -             return get_data_error_result(retmsg="Document not found!")
 -         query = {
 -             "doc_ids": [doc_id], "page": page, "size": size, "question": question, "sort": True
 -         }
 -         if "available_int" in req:
 -             query["available_int"] = int(req["available_int"])
 -         sres = retrievaler.search(query, search.index_name(tenant_id), highlight=True)
 -         res = {"total": sres.total, "chunks": [], "doc": doc.to_dict()}
 - 
 -         origin_chunks=[]
 -         for id in sres.ids:
 -             d = {
 -                 "chunk_id": id,
 -                 "content_with_weight": rmSpace(sres.highlight[id]) if question and id in sres.highlight else sres.field[
 -                     id].get(
 -                     "content_with_weight", ""),
 -                 "doc_id": sres.field[id]["doc_id"],
 -                 "docnm_kwd": sres.field[id]["docnm_kwd"],
 -                 "important_kwd": sres.field[id].get("important_kwd", []),
 -                 "img_id": sres.field[id].get("img_id", ""),
 -                 "available_int": sres.field[id].get("available_int", 1),
 -                 "positions": sres.field[id].get("position_int", "").split("\t")
 -             }
 -             if len(d["positions"]) % 5 == 0:
 -                 poss = []
 -                 for i in range(0, len(d["positions"]), 5):
 -                     poss.append([float(d["positions"][i]), float(d["positions"][i + 1]), float(d["positions"][i + 2]),
 -                                  float(d["positions"][i + 3]), float(d["positions"][i + 4])])
 -                 d["positions"] = poss
 -                 
 -             origin_chunks.append(d)
 -             ##rename keys
 -             for chunk in origin_chunks:
 -                 key_mapping = {
 -                     "chunk_id": "id",
 -                     "content_with_weight": "content",
 -                     "doc_id": "document_id",
 -                     "important_kwd": "important_keywords",
 -                 }
 -                 renamed_chunk = {}
 -                 for key, value in chunk.items():
 -                     new_key = key_mapping.get(key, key)
 -                     renamed_chunk[new_key] = value
 -                 res["chunks"].append(renamed_chunk)
 -         return get_json_result(data=res)
 -     except Exception as e:
 -         if str(e).find("not_found") > 0:
 -             return get_json_result(data=False, retmsg=f'No chunk found!',
 -                                    retcode=RetCode.DATA_ERROR)
 -         return server_error_response(e)
 - 
 - 
 - @manager.route('/chunk/create', methods=['POST'])
 - @token_required
 - @validate_request("doc_id", "content_with_weight")
 - def create(tenant_id):
 -     req = request.json
 -     md5 = hashlib.md5()
 -     md5.update((req["content_with_weight"] + req["doc_id"]).encode("utf-8"))
 - 
 -     chunk_id = md5.hexdigest()
 -     d = {"id": chunk_id, "content_ltks": rag_tokenizer.tokenize(req["content_with_weight"]),
 -          "content_with_weight": req["content_with_weight"]}
 -     d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
 -     d["important_kwd"] = req.get("important_kwd", [])
 -     d["important_tks"] = rag_tokenizer.tokenize(" ".join(req.get("important_kwd", [])))
 -     d["create_time"] = str(datetime.datetime.now()).replace("T", " ")[:19]
 -     d["create_timestamp_flt"] = datetime.datetime.now().timestamp()
 - 
 -     try:
 -         e, doc = DocumentService.get_by_id(req["doc_id"])
 -         if not e:
 -             return get_data_error_result(retmsg="Document not found!")
 -         d["kb_id"] = [doc.kb_id]
 -         d["docnm_kwd"] = doc.name
 -         d["doc_id"] = doc.id
 - 
 -         tenant_id = DocumentService.get_tenant_id(req["doc_id"])
 -         if not tenant_id:
 -             return get_data_error_result(retmsg="Tenant not found!")
 - 
 -         embd_id = DocumentService.get_embd_id(req["doc_id"])
 -         embd_mdl = TenantLLMService.model_instance(
 -             tenant_id, LLMType.EMBEDDING.value, embd_id)
 - 
 -         v, c = embd_mdl.encode([doc.name, req["content_with_weight"]])
 -         v = 0.1 * v[0] + 0.9 * v[1]
 -         d["q_%d_vec" % len(v)] = v.tolist()
 -         ELASTICSEARCH.upsert([d], search.index_name(tenant_id))
 - 
 -         DocumentService.increment_chunk_num(
 -             doc.id, doc.kb_id, c, 1, 0)
 -         d["chunk_id"] = chunk_id
 -         #rename keys
 -         key_mapping = {
 -             "chunk_id": "id",
 -             "content_with_weight": "content",
 -             "doc_id": "document_id",
 -             "important_kwd": "important_keywords",
 -             "kb_id":"knowledge_base_id",
 -         }
 -         renamed_chunk = {}
 -         for key, value in d.items():
 -             new_key = key_mapping.get(key, key)
 -             renamed_chunk[new_key] = value
 - 
 -         return get_json_result(data={"chunk": renamed_chunk})
 -         # return get_json_result(data={"chunk_id": chunk_id})
 -     except Exception as e:
 -         return server_error_response(e)
 - 
 - @manager.route('/chunk/rm', methods=['POST'])
 - @token_required
 - @validate_request("chunk_ids", "doc_id")
 - def rm_chunk(tenant_id):
 -     req = request.json
 -     try:
 -         if not ELASTICSEARCH.deleteByQuery(
 -                 Q("ids", values=req["chunk_ids"]), search.index_name(tenant_id)):
 -             return get_data_error_result(retmsg="Index updating failure")
 -         e, doc = DocumentService.get_by_id(req["doc_id"])
 -         if not e:
 -             return get_data_error_result(retmsg="Document not found!")
 -         deleted_chunk_ids = req["chunk_ids"]
 -         chunk_number = len(deleted_chunk_ids)
 -         DocumentService.decrement_chunk_num(doc.id, doc.kb_id, 1, chunk_number, 0)
 -         return get_json_result(data=True)
 -     except Exception as e:
 -         return server_error_response(e)
 - 
 - @manager.route('/chunk/set', methods=['POST'])
 - @token_required
 - @validate_request("doc_id", "chunk_id", "content_with_weight",
 -                   "important_kwd")
 - def set(tenant_id):
 -     req = request.json
 -     d = {
 -         "id": req["chunk_id"],
 -         "content_with_weight": req["content_with_weight"]}
 -     d["content_ltks"] = rag_tokenizer.tokenize(req["content_with_weight"])
 -     d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
 -     d["important_kwd"] = req["important_kwd"]
 -     d["important_tks"] = rag_tokenizer.tokenize(" ".join(req["important_kwd"]))
 -     if "available_int" in req:
 -         d["available_int"] = req["available_int"]
 - 
 -     try:
 -         tenant_id = DocumentService.get_tenant_id(req["doc_id"])
 -         if not tenant_id:
 -             return get_data_error_result(retmsg="Tenant not found!")
 - 
 -         embd_id = DocumentService.get_embd_id(req["doc_id"])
 -         embd_mdl = TenantLLMService.model_instance(
 -             tenant_id, LLMType.EMBEDDING.value, embd_id)
 - 
 -         e, doc = DocumentService.get_by_id(req["doc_id"])
 -         if not e:
 -             return get_data_error_result(retmsg="Document not found!")
 - 
 -         if doc.parser_id == ParserType.QA:
 -             arr = [
 -                 t for t in re.split(
 -                     r"[\n\t]",
 -                     req["content_with_weight"]) if len(t) > 1]
 -             if len(arr) != 2:
 -                 return get_data_error_result(
 -                     retmsg="Q&A must be separated by TAB/ENTER key.")
 -             q, a = rmPrefix(arr[0]), rmPrefix(arr[1])
 -             d = beAdoc(d, arr[0], arr[1], not any(
 -                 [rag_tokenizer.is_chinese(t) for t in q + a]))
 - 
 -         v, c = embd_mdl.encode([doc.name, req["content_with_weight"]])
 -         v = 0.1 * v[0] + 0.9 * v[1] if doc.parser_id != ParserType.QA else v[1]
 -         d["q_%d_vec" % len(v)] = v.tolist()
 -         ELASTICSEARCH.upsert([d], search.index_name(tenant_id))
 -         return get_json_result(data=True)
 -     except Exception as e:
 -         return server_error_response(e)
 - 
 - @manager.route('/retrieval_test', methods=['POST'])
 - @token_required
 - @validate_request("kb_id", "question")
 - def retrieval_test(tenant_id):
 -     req = request.json
 -     page = int(req.get("page", 1))
 -     size = int(req.get("size", 30))
 -     question = req["question"]
 -     kb_id = req["kb_id"]
 -     if isinstance(kb_id, str): kb_id = [kb_id]
 -     doc_ids = req.get("doc_ids", [])
 -     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:
 -         tenants = UserTenantService.query(user_id=tenant_id)
 -         for kid in kb_id:
 -             for tenant in tenants:
 -                 if KnowledgebaseService.query(
 -                         tenant_id=tenant.tenant_id, id=kid):
 -                     break
 -             else:
 -                 return get_json_result(
 -                     data=False, retmsg=f'Only owner of knowledgebase authorized for this operation.',
 -                     retcode=RetCode.OPERATING_ERROR)
 - 
 -         e, kb = KnowledgebaseService.get_by_id(kb_id[0])
 -         if not e:
 -             return get_data_error_result(retmsg="Knowledgebase not found!")
 - 
 -         embd_mdl = TenantLLMService.model_instance(
 -             kb.tenant_id, LLMType.EMBEDDING.value, llm_name=kb.embd_id)
 - 
 -         rerank_mdl = None
 -         if req.get("rerank_id"):
 -             rerank_mdl = TenantLLMService.model_instance(
 -                 kb.tenant_id, LLMType.RERANK.value, llm_name=req["rerank_id"])
 - 
 -         if req.get("keyword", False):
 -             chat_mdl = TenantLLMService.model_instance(kb.tenant_id, LLMType.CHAT)
 -             question += keyword_extraction(chat_mdl, question)
 - 
 -         retr = retrievaler if kb.parser_id != ParserType.KG else kg_retrievaler
 -         ranks = retr.retrieval(question, embd_mdl, kb.tenant_id, kb_id, page, size,
 -                                similarity_threshold, vector_similarity_weight, top,
 -                                doc_ids, rerank_mdl=rerank_mdl, highlight=req.get("highlight"))
 -         for c in ranks["chunks"]:
 -             if "vector" in c:
 -                 del c["vector"]
 - 
 -         ##rename keys
 -         renamed_chunks=[]
 -         for chunk in ranks["chunks"]:
 -             key_mapping = {
 -                 "chunk_id": "id",
 -                 "content_with_weight": "content",
 -                 "doc_id": "document_id",
 -                 "important_kwd": "important_keywords",
 -             }
 -             rename_chunk={}
 -             for key, value in chunk.items():
 -                 new_key = key_mapping.get(key, key)
 -                 rename_chunk[new_key] = value
 -                 renamed_chunks.append(rename_chunk)
 -         ranks["chunks"] = renamed_chunks
 -         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)
 
 
  |