zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-09 12:54:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refa: migrate chunk APIs to RESTful routes (#14291)

Browse Source 
### What problem does this PR solve?

migrate chunk APIs to RESTful routes

### Type of change
- [x] Refactoring
This commit is contained in:
buua436
2026-04-23 14:17:23 +08:00
committed by GitHub
parent 76b017ca32
commit 7817b0d779
30 changed files with 1594 additions and 2237 deletions
Download Patch File Download Diff File Expand all files Collapse all files

394
api/apps/chunk_app.py
View File

@@ -13,401 +13,35 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
import base64
import datetime
import json
import logging
import re
import xxhash
from quart import request
from api.db.services.document_service import DocumentService
from api.apps import current_user, login_required
from api.db.joint_services.tenant_model_service import (
get_model_config_by_id,
get_model_config_by_type_and_name,
get_tenant_default_model_by_type,
)
from api.db.services.doc_metadata_service import DocMetadataService
from api.utils.image_utils import store_chunk_image
from api.db.services.document_service import DocumentService
from api.db.services.knowledgebase_service import KnowledgebaseService
from api.db.services.llm_service import LLMBundle
from common.metadata_utils import apply_meta_data_filter
from api.db.services.search_service import SearchService
from api.db.services.user_service import UserTenantService
from api.db.joint_services.tenant_model_service import get_model_config_by_id, get_tenant_default_model_by_type, get_model_config_by_type_and_name
from api.utils.api_utils import (
get_data_error_result,
get_json_result,
get_request_json,
server_error_response,
validate_request,
get_request_json,
)
from common.misc_utils import thread_pool_exec
from common.tag_feature_utils import validate_tag_features
from rag.app.qa import beAdoc, rmPrefix
from rag.app.tag import label_question
from rag.nlp import rag_tokenizer, search
from rag.prompts.generator import cross_languages, keyword_extraction
from common.string_utils import is_content_empty, remove_redundant_spaces
from common.constants import RetCode, LLMType, ParserType, PAGERANK_FLD
from common import settings
from api.apps import login_required, current_user
@manager.route('/list', methods=['POST']) # noqa: F821
@login_required
@validate_request("doc_id")
async def list_chunk():
req = await get_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(message="Tenant not found!")
e, doc = DocumentService.get_by_id(doc_id)
if not e:
return get_data_error_result(message="Document not found!")
kb_ids = KnowledgebaseService.get_kb_ids(tenant_id)
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 = await settings.retriever.search(query, search.index_name(tenant_id), kb_ids, highlight=["content_ltks"])
res = {"total": sres.total, "chunks": [], "doc": doc.to_dict()}
for id in sres.ids:
d = {
"chunk_id": id,
"content_with_weight": remove_redundant_spaces(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", []),
"question_kwd": sres.field[id].get("question_kwd", []),
"image_id": sres.field[id].get("img_id", ""),
"available_int": int(sres.field[id].get("available_int", 1)),
"positions": sres.field[id].get("position_int", []),
"doc_type_kwd": sres.field[id].get("doc_type_kwd")
}
assert isinstance(d["positions"], list)
assert len(d["positions"]) == 0 or (isinstance(d["positions"][0], list) and len(d["positions"][0]) == 5)
res["chunks"].append(d)
return get_json_result(data=res)
except Exception as e:
if str(e).find("not_found") > 0:
return get_json_result(data=False, message='No chunk found!',
code=RetCode.DATA_ERROR)
return server_error_response(e)
@manager.route('/get', methods=['GET']) # noqa: F821
@login_required
def get():
chunk_id = request.args["chunk_id"]
try:
chunk = None
tenants = UserTenantService.query(user_id=current_user.id)
if not tenants:
return get_data_error_result(message="Tenant not found!")
for tenant in tenants:
kb_ids = KnowledgebaseService.get_kb_ids(tenant.tenant_id)
chunk = settings.docStoreConn.get(chunk_id, search.index_name(tenant.tenant_id), kb_ids)
if chunk:
break
if chunk is None:
return server_error_response(Exception("Chunk not found"))
k = []
for n in chunk.keys():
if re.search(r"(_vec$|_sm_|_tks|_ltks)", n):
k.append(n)
for n in k:
del chunk[n]
return get_json_result(data=chunk)
except Exception as e:
if str(e).find("NotFoundError") >= 0:
return get_json_result(data=False, message='Chunk not found!',
code=RetCode.DATA_ERROR)
return server_error_response(e)
@manager.route('/set', methods=['POST']) # noqa: F821
@login_required
@validate_request("doc_id", "chunk_id", "content_with_weight")
async def set():
req = await get_request_json()
content_with_weight = req["content_with_weight"]
if not isinstance(content_with_weight, (str, bytes)):
raise TypeError("expected string or bytes-like object")
if isinstance(content_with_weight, bytes):
content_with_weight = content_with_weight.decode("utf-8", errors="ignore")
if is_content_empty(content_with_weight):
return get_data_error_result(message="`content_with_weight` is required")
d = {
"id": req["chunk_id"],
"content_with_weight": content_with_weight}
d["content_ltks"] = rag_tokenizer.tokenize(content_with_weight)
d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
if "important_kwd" in req:
if not isinstance(req["important_kwd"], list):
return get_data_error_result(message="`important_kwd` should be a list")
d["important_kwd"] = req["important_kwd"]
d["important_tks"] = rag_tokenizer.tokenize(" ".join(req["important_kwd"]))
if "question_kwd" in req:
if not isinstance(req["question_kwd"], list):
return get_data_error_result(message="`question_kwd` should be a list")
d["question_kwd"] = req["question_kwd"]
d["question_tks"] = rag_tokenizer.tokenize("\n".join(req["question_kwd"]))
if "tag_kwd" in req:
if not isinstance(req["tag_kwd"], list):
return get_data_error_result(message="`tag_kwd` should be a list")
if not all(isinstance(t, str) for t in req["tag_kwd"]):
return get_data_error_result(message="`tag_kwd` must be a list of strings")
d["tag_kwd"] = req["tag_kwd"]
if "tag_feas" in req:
try:
d["tag_feas"] = validate_tag_features(req["tag_feas"])
except ValueError as exc:
return get_data_error_result(message=f"`tag_feas` {exc}")
if "available_int" in req:
d["available_int"] = req["available_int"]
try:
def _set_sync():
tenant_id = DocumentService.get_tenant_id(req["doc_id"])
if not tenant_id:
return get_data_error_result(message="Tenant not found!")
e, doc = DocumentService.get_by_id(req["doc_id"])
if not e:
return get_data_error_result(message="Document not found!")
tenant_embd_id = DocumentService.get_tenant_embd_id(req["doc_id"])
if tenant_embd_id:
embd_model_config = get_model_config_by_id(tenant_embd_id)
else:
embd_id = DocumentService.get_embd_id(req["doc_id"])
if embd_id:
embd_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.EMBEDDING, embd_id)
else:
embd_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.EMBEDDING)
embd_mdl = LLMBundle(tenant_id, embd_model_config)
_d = d
if doc.parser_id == ParserType.QA:
arr = [
t for t in re.split(
r"[\n\t]",
req["content_with_weight"]) if len(t) > 1]
q, a = rmPrefix(arr[0]), rmPrefix("\n".join(arr[1:]))
_d = beAdoc(d, q, a, not any(
[rag_tokenizer.is_chinese(t) for t in q + a]))
v, c = embd_mdl.encode([doc.name, content_with_weight if not _d.get("question_kwd") else "\n".join(_d["question_kwd"])])
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()
settings.docStoreConn.update({"id": req["chunk_id"]}, _d, search.index_name(tenant_id), doc.kb_id)
# update image
image_base64 = req.get("image_base64", None)
img_id = req.get("img_id", "")
if image_base64 and img_id and "-" in img_id:
bkt, name = img_id.split("-", 1)
image_binary = base64.b64decode(image_base64)
settings.STORAGE_IMPL.put(bkt, name, image_binary)
return get_json_result(data=True)
return await thread_pool_exec(_set_sync)
except Exception as e:
return server_error_response(e)
@manager.route('/switch', methods=['POST']) # noqa: F821
@login_required
@validate_request("chunk_ids", "available_int", "doc_id")
async def switch():
req = await get_request_json()
try:
def _switch_sync():
e, doc = DocumentService.get_by_id(req["doc_id"])
if not e:
return get_data_error_result(message="Document not found!")
for cid in req["chunk_ids"]:
if not settings.docStoreConn.update({"id": cid},
{"available_int": int(req["available_int"])},
search.index_name(DocumentService.get_tenant_id(req["doc_id"])),
doc.kb_id):
return get_data_error_result(message="Index updating failure")
return get_json_result(data=True)
return await thread_pool_exec(_switch_sync)
except Exception as e:
return server_error_response(e)
@manager.route('/rm', methods=['POST']) # noqa: F821
@login_required
@validate_request("doc_id")
async def rm():
req = await get_request_json()
try:
def _rm_sync():
deleted_chunk_ids = req.get("chunk_ids")
if isinstance(deleted_chunk_ids, list):
unique_chunk_ids = list(dict.fromkeys(deleted_chunk_ids))
has_ids = len(unique_chunk_ids) > 0
elif deleted_chunk_ids is not None:
unique_chunk_ids = [deleted_chunk_ids]
has_ids = deleted_chunk_ids not in (None, "")
else:
unique_chunk_ids = []
has_ids = False
if not has_ids:
if req.get("delete_all") is True:
e, doc = DocumentService.get_by_id(req["doc_id"])
if not e:
return get_data_error_result(message="Document not found!")
tenant_id = DocumentService.get_tenant_id(req["doc_id"])
# Clean up storage assets while index rows still exist for discovery
DocumentService.delete_chunk_images(doc, tenant_id)
condition = {"doc_id": req["doc_id"]}
try:
deleted_count = settings.docStoreConn.delete(condition, search.index_name(tenant_id), doc.kb_id)
except Exception:
return get_data_error_result(message="Chunk deleting failure")
if deleted_count > 0:
DocumentService.decrement_chunk_num(doc.id, doc.kb_id, 1, deleted_count, 0)
return get_json_result(data=True)
return get_json_result(data=True)
e, doc = DocumentService.get_by_id(req["doc_id"])
if not e:
return get_data_error_result(message="Document not found!")
condition = {"id": req["chunk_ids"], "doc_id": req["doc_id"]}
try:
deleted_count = settings.docStoreConn.delete(condition,
search.index_name(DocumentService.get_tenant_id(req["doc_id"])),
doc.kb_id)
except Exception:
return get_data_error_result(message="Chunk deleting failure")
if has_ids and deleted_count == 0:
return get_data_error_result(message="Index updating failure")
if deleted_count > 0 and deleted_count < len(unique_chunk_ids):
deleted_count += settings.docStoreConn.delete({"doc_id": req["doc_id"]},
search.index_name(DocumentService.get_tenant_id(req["doc_id"])),
doc.kb_id)
chunk_number = deleted_count
DocumentService.decrement_chunk_num(doc.id, doc.kb_id, 1, chunk_number, 0)
for cid in deleted_chunk_ids:
if settings.STORAGE_IMPL.obj_exist(doc.kb_id, cid):
settings.STORAGE_IMPL.rm(doc.kb_id, cid)
return get_json_result(data=True)
return await thread_pool_exec(_rm_sync)
except Exception as e:
return server_error_response(e)
@manager.route('/create', methods=['POST']) # noqa: F821
@login_required
@validate_request("doc_id", "content_with_weight")
async def create():
req = await get_request_json()
req_id = request.headers.get("X-Request-ID")
chunck_id = xxhash.xxh64((req["content_with_weight"] + req["doc_id"]).encode("utf-8")).hexdigest()
d = {"id": chunck_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", [])
if not isinstance(d["important_kwd"], list):
return get_data_error_result(message="`important_kwd` is required to be a list")
d["important_tks"] = rag_tokenizer.tokenize(" ".join(d["important_kwd"]))
d["question_kwd"] = req.get("question_kwd", [])
if not isinstance(d["question_kwd"], list):
return get_data_error_result(message="`question_kwd` is required to be a list")
d["question_tks"] = rag_tokenizer.tokenize("\n".join(d["question_kwd"]))
d["create_time"] = str(datetime.datetime.now()).replace("T", " ")[:19]
d["create_timestamp_flt"] = datetime.datetime.now().timestamp()
if "tag_kwd" in req:
if not isinstance(req["tag_kwd"], list):
return get_data_error_result(message="`tag_kwd` is required to be a list")
if not all(isinstance(t, str) for t in req["tag_kwd"]):
return get_data_error_result(message="`tag_kwd` must be a list of strings")
d["tag_kwd"] = req["tag_kwd"]
if "tag_feas" in req:
try:
d["tag_feas"] = validate_tag_features(req["tag_feas"])
except ValueError as exc:
return get_data_error_result(message=f"`tag_feas` {exc}")
image_base64 = req.get("image_base64", None)
try:
def _log_response(resp, code, message):
logging.info(
"chunk_create response req_id=%s status=%s code=%s message=%s",
req_id,
getattr(resp, "status_code", None),
code,
message,
)
def _create_sync():
e, doc = DocumentService.get_by_id(req["doc_id"])
if not e:
resp = get_data_error_result(message="Document not found!")
_log_response(resp, RetCode.DATA_ERROR, "Document not found!")
return resp
d["kb_id"] = [doc.kb_id]
d["docnm_kwd"] = doc.name
d["title_tks"] = rag_tokenizer.tokenize(doc.name)
d["doc_id"] = doc.id
tenant_id = DocumentService.get_tenant_id(req["doc_id"])
if not tenant_id:
resp = get_data_error_result(message="Tenant not found!")
_log_response(resp, RetCode.DATA_ERROR, "Tenant not found!")
return resp
e, kb = KnowledgebaseService.get_by_id(doc.kb_id)
if not e:
resp = get_data_error_result(message="Knowledgebase not found!")
_log_response(resp, RetCode.DATA_ERROR, "Knowledgebase not found!")
return resp
if kb.pagerank:
d[PAGERANK_FLD] = kb.pagerank
tenant_embd_id = DocumentService.get_tenant_embd_id(req["doc_id"])
if tenant_embd_id:
embd_model_config = get_model_config_by_id(tenant_embd_id)
else:
embd_id = DocumentService.get_embd_id(req["doc_id"])
if embd_id:
embd_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.EMBEDDING, embd_id)
else:
embd_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.EMBEDDING)
embd_mdl = LLMBundle(tenant_id, embd_model_config)
if image_base64:
d["img_id"] = "{}-{}".format(doc.kb_id, chunck_id)
d["doc_type_kwd"] = "image"
v, c = embd_mdl.encode([doc.name, req["content_with_weight"] if not d["question_kwd"] else "\n".join(d["question_kwd"])])
v = 0.1 * v[0] + 0.9 * v[1]
d["q_%d_vec" % len(v)] = v.tolist()
settings.docStoreConn.insert([d], search.index_name(tenant_id), doc.kb_id)
if image_base64:
store_chunk_image(doc.kb_id, chunck_id, base64.b64decode(image_base64))
DocumentService.increment_chunk_num(
doc.id, doc.kb_id, c, 1, 0)
resp = get_json_result(data={"chunk_id": chunck_id, "image_id": d.get("img_id", "")})
_log_response(resp, RetCode.SUCCESS, "success")
return resp
return await thread_pool_exec(_create_sync)
except Exception as e:
logging.info("chunk_create exception req_id=%s error=%r", req_id, e)
return server_error_response(e)
from common.constants import LLMType, RetCode
from common.metadata_utils import apply_meta_data_filter
from rag.app.tag import label_question
from rag.nlp import search
from rag.prompts.generator import cross_languages, keyword_extraction
@manager.route('/retrieval_test', methods=['POST']) # noqa: F821

445
api/apps/restful_apis/chunk_api.py Normal file
View File

@@ -0,0 +1,445 @@
#
# Copyright 2026 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 base64
import datetime
import re
import xxhash
from pydantic import BaseModel, Field, validator
from quart import request
from api.apps import login_required
from api.db.joint_services.tenant_model_service import (
get_model_config_by_id,
get_model_config_by_type_and_name,
)
from api.db.services.document_service import DocumentService
from api.db.services.knowledgebase_service import KnowledgebaseService
from api.db.services.tenant_llm_service import TenantLLMService
from api.utils.api_utils import (
add_tenant_id_to_kwargs,
check_duplicate_ids,
get_error_data_result,
get_request_json,
get_result,
server_error_response,
)
from api.utils.image_utils import store_chunk_image
from common import settings
from common.constants import LLMType, ParserType, RetCode
from common.misc_utils import thread_pool_exec
from common.string_utils import is_content_empty, remove_redundant_spaces
from common.tag_feature_utils import validate_tag_features
from rag.app.qa import beAdoc, rmPrefix
from rag.nlp import rag_tokenizer, search
class Chunk(BaseModel):
id: str = ""
content: str = ""
document_id: str = ""
docnm_kwd: str = ""
important_keywords: list = Field(default_factory=list)
tag_kwd: list = Field(default_factory=list)
questions: list = Field(default_factory=list)
question_tks: str = ""
image_id: str = ""
available: bool = True
positions: list[list[int]] = Field(default_factory=list)
@validator("positions")
def validate_positions(cls, value):
for sublist in value:
if len(sublist) != 5:
raise ValueError("Each sublist in positions must have a length of 5")
return value
def _map_doc(doc):
key_mapping = {
"chunk_num": "chunk_count",
"kb_id": "dataset_id",
"token_num": "token_count",
"parser_id": "chunk_method",
}
run_mapping = {
"0": "UNSTART",
"1": "RUNNING",
"2": "CANCEL",
"3": "DONE",
"4": "FAIL",
}
renamed_doc = {}
for key, value in doc.to_dict().items():
renamed_doc[key_mapping.get(key, key)] = value
if key == "run":
renamed_doc["run"] = run_mapping.get(str(value))
return renamed_doc
def _strip_chunk_runtime_fields(chunk):
for name in [name for name in chunk.keys() if re.search(r"(_vec$|_sm_|_tks|_ltks)", name)]:
del chunk[name]
return chunk
@manager.route("/datasets/<dataset_id>/documents/<document_id>/chunks", methods=["GET"]) # noqa: F821
@login_required
@add_tenant_id_to_kwargs
async def list_chunks(tenant_id, dataset_id, document_id):
if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id):
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
doc = DocumentService.query(id=document_id, kb_id=dataset_id)
if not doc:
return get_error_data_result(message=f"You don't own the document {document_id}.")
doc = doc[0]
req = request.args
page = int(req.get("page", 1))
size = int(req.get("page_size", 30))
question = req.get("keywords", "")
query = {
"doc_ids": [document_id],
"page": page,
"size": size,
"question": question,
"sort": True,
}
if "available" in req:
query["available_int"] = 1 if req["available"] == "true" else 0
res = {"total": 0, "chunks": [], "doc": _map_doc(doc)}
if req.get("id"):
chunk = settings.docStoreConn.get(req.get("id"), search.index_name(tenant_id), [dataset_id])
if not chunk:
return get_result(message=f"Chunk not found: {dataset_id}/{req.get('id')}", code=RetCode.DATA_ERROR)
if str(chunk.get("doc_id", chunk.get("document_id"))) != str(document_id):
return get_result(message=f"Chunk not found: {dataset_id}/{req.get('id')}", code=RetCode.DATA_ERROR)
_strip_chunk_runtime_fields(chunk)
res["total"] = 1
final_chunk = {
"id": chunk.get("id", chunk.get("chunk_id")),
"content": chunk["content_with_weight"],
"document_id": chunk.get("doc_id", chunk.get("document_id")),
"docnm_kwd": chunk["docnm_kwd"],
"important_keywords": chunk.get("important_kwd", []),
"questions": chunk.get("question_kwd", []),
"dataset_id": chunk.get("kb_id", chunk.get("dataset_id")),
"image_id": chunk.get("img_id", ""),
"available": bool(chunk.get("available_int", 1)),
"positions": chunk.get("position_int", []),
"tag_kwd": chunk.get("tag_kwd", []),
"tag_feas": chunk.get("tag_feas", {}),
}
res["chunks"].append(final_chunk)
_ = Chunk(**final_chunk)
elif settings.docStoreConn.index_exist(search.index_name(tenant_id), dataset_id):
sres = await settings.retriever.search(
query,
search.index_name(tenant_id),
[dataset_id],
emb_mdl=None,
highlight=True,
)
res["total"] = sres.total
for chunk_id in sres.ids:
d = {
"id": chunk_id,
"content": (
remove_redundant_spaces(sres.highlight[chunk_id])
if question and chunk_id in sres.highlight
else sres.field[chunk_id].get("content_with_weight", "")
),
"document_id": sres.field[chunk_id]["doc_id"],
"docnm_kwd": sres.field[chunk_id]["docnm_kwd"],
"important_keywords": sres.field[chunk_id].get("important_kwd", []),
"tag_kwd": sres.field[chunk_id].get("tag_kwd", []),
"questions": sres.field[chunk_id].get("question_kwd", []),
"dataset_id": sres.field[chunk_id].get("kb_id", sres.field[chunk_id].get("dataset_id")),
"image_id": sres.field[chunk_id].get("img_id", ""),
"available": bool(int(sres.field[chunk_id].get("available_int", "1"))),
"positions": sres.field[chunk_id].get("position_int", []),
}
res["chunks"].append(d)
_ = Chunk(**d)
return get_result(data=res)
@manager.route("/datasets/<dataset_id>/documents/<document_id>/chunks/<chunk_id>", methods=["GET"]) # noqa: F821
@login_required
@add_tenant_id_to_kwargs
async def get_chunk(tenant_id, dataset_id, document_id, chunk_id):
if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id):
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
doc = DocumentService.query(id=document_id, kb_id=dataset_id)
if not doc:
return get_error_data_result(message=f"You don't own the document {document_id}.")
try:
chunk = settings.docStoreConn.get(chunk_id, search.index_name(tenant_id), [dataset_id])
if chunk is None or str(chunk.get("doc_id", chunk.get("document_id"))) != str(document_id):
return get_result(data=False, message="Chunk not found!", code=RetCode.DATA_ERROR)
return get_result(data=_strip_chunk_runtime_fields(chunk))
except Exception as e:
if str(e).find("NotFoundError") >= 0:
return get_result(data=False, message="Chunk not found!", code=RetCode.DATA_ERROR)
return server_error_response(e)
@manager.route("/datasets/<dataset_id>/documents/<document_id>/chunks", methods=["POST"]) # noqa: F821
@login_required
@add_tenant_id_to_kwargs
async def add_chunk(tenant_id, dataset_id, document_id):
if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id):
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
doc = DocumentService.query(id=document_id, kb_id=dataset_id)
if not doc:
return get_error_data_result(message=f"You don't own the document {document_id}.")
doc = doc[0]
req = await get_request_json()
if is_content_empty(req.get("content")):
return get_error_data_result(message="`content` is required")
if "important_keywords" in req and not isinstance(req["important_keywords"], list):
return get_error_data_result("`important_keywords` is required to be a list")
if "questions" in req and not isinstance(req["questions"], list):
return get_error_data_result("`questions` is required to be a list")
chunk_id = xxhash.xxh64((req["content"] + document_id).encode("utf-8")).hexdigest()
d = {
"id": chunk_id,
"content_ltks": rag_tokenizer.tokenize(req["content"]),
"content_with_weight": req["content"],
}
d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
d["important_kwd"] = req.get("important_keywords", [])
d["important_tks"] = rag_tokenizer.tokenize(" ".join(req.get("important_keywords", [])))
d["question_kwd"] = [str(q).strip() for q in req.get("questions", []) if str(q).strip()]
d["question_tks"] = rag_tokenizer.tokenize("\n".join(req.get("questions", [])))
d["create_time"] = str(datetime.datetime.now()).replace("T", " ")[:19]
d["create_timestamp_flt"] = datetime.datetime.now().timestamp()
d["kb_id"] = dataset_id
d["docnm_kwd"] = doc.name
d["doc_id"] = document_id
if "tag_kwd" in req:
if not isinstance(req["tag_kwd"], list):
return get_error_data_result("`tag_kwd` is required to be a list")
if not all(isinstance(t, str) for t in req["tag_kwd"]):
return get_error_data_result("`tag_kwd` must be a list of strings")
d["tag_kwd"] = req["tag_kwd"]
if "tag_feas" in req:
try:
d["tag_feas"] = validate_tag_features(req["tag_feas"])
except ValueError as exc:
return get_error_data_result(f"`tag_feas` {exc}")
image_base64 = req.get("image_base64")
if image_base64:
d["img_id"] = f"{dataset_id}-{chunk_id}"
d["doc_type_kwd"] = "image"
tenant_embd_id = DocumentService.get_tenant_embd_id(document_id)
if tenant_embd_id:
model_config = get_model_config_by_id(tenant_embd_id)
else:
embd_id = DocumentService.get_embd_id(document_id)
model_config = get_model_config_by_type_and_name(tenant_id, LLMType.EMBEDDING.value, embd_id)
embd_mdl = TenantLLMService.model_instance(model_config)
v, c = embd_mdl.encode([doc.name, req["content"] if not d["question_kwd"] else "\n".join(d["question_kwd"])])
v = 0.1 * v[0] + 0.9 * v[1]
d[f"q_{len(v)}_vec"] = v.tolist()
settings.docStoreConn.insert([d], search.index_name(tenant_id), dataset_id)
if image_base64:
store_chunk_image(dataset_id, chunk_id, base64.b64decode(image_base64))
DocumentService.increment_chunk_num(doc.id, doc.kb_id, c, 1, 0)
key_mapping = {
"id": "id",
"content_with_weight": "content",
"doc_id": "document_id",
"important_kwd": "important_keywords",
"tag_kwd": "tag_kwd",
"question_kwd": "questions",
"kb_id": "dataset_id",
"create_timestamp_flt": "create_timestamp",
"create_time": "create_time",
"document_keyword": "document",
"img_id": "image_id",
}
renamed_chunk = {new_key: d[key] for key, new_key in key_mapping.items() if key in d}
_ = Chunk(**renamed_chunk)
return get_result(data={"chunk": renamed_chunk})
@manager.route("/datasets/<dataset_id>/documents/<document_id>/chunks", methods=["DELETE"]) # noqa: F821
@login_required
@add_tenant_id_to_kwargs
async def rm_chunk(tenant_id, dataset_id, document_id):
if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id):
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
docs = DocumentService.query(id=document_id, kb_id=dataset_id)
if not docs:
return get_error_data_result(message=f"You don't own the document {document_id}.")
req = await get_request_json()
if not req:
return get_result()
chunk_ids = req.get("chunk_ids")
if not chunk_ids:
if req.get("delete_all") is True:
doc = docs[0]
DocumentService.delete_chunk_images(doc, tenant_id)
chunk_number = settings.docStoreConn.delete({"doc_id": document_id}, search.index_name(tenant_id), dataset_id)
if chunk_number != 0:
DocumentService.decrement_chunk_num(document_id, dataset_id, 1, chunk_number, 0)
return get_result(message=f"deleted {chunk_number} chunks")
return get_result()
unique_chunk_ids, duplicate_messages = check_duplicate_ids(chunk_ids, "chunk")
chunk_number = settings.docStoreConn.delete(
{"doc_id": document_id, "id": unique_chunk_ids},
search.index_name(tenant_id),
dataset_id,
)
if chunk_number != 0:
DocumentService.decrement_chunk_num(document_id, dataset_id, 1, chunk_number, 0)
if chunk_number != len(unique_chunk_ids):
if len(unique_chunk_ids) == 0:
return get_result(message=f"deleted {chunk_number} chunks")
return get_error_data_result(message=f"rm_chunk deleted chunks {chunk_number}, expect {len(unique_chunk_ids)}")
if duplicate_messages:
return get_result(
message=f"Partially deleted {chunk_number} chunks with {len(duplicate_messages)} errors",
data={"success_count": chunk_number, "errors": duplicate_messages},
)
return get_result(message=f"deleted {chunk_number} chunks")
@manager.route("/datasets/<dataset_id>/documents/<document_id>/chunks/<chunk_id>", methods=["PATCH"]) # noqa: F821
@login_required
@add_tenant_id_to_kwargs
async def update_chunk(tenant_id, dataset_id, document_id, chunk_id):
if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id):
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
doc = DocumentService.query(id=document_id, kb_id=dataset_id)
if not doc:
return get_error_data_result(message=f"You don't own the document {document_id}.")
doc = doc[0]
chunk = settings.docStoreConn.get(chunk_id, search.index_name(tenant_id), [dataset_id])
if chunk is None or str(chunk.get("doc_id", chunk.get("document_id"))) != str(document_id):
return get_error_data_result(f"Can't find this chunk {chunk_id}")
req = await get_request_json()
content = req.get("content")
if content is not None:
if is_content_empty(content):
return get_error_data_result(message="`content` is required")
else:
content = chunk.get("content_with_weight", "")
d = {"id": chunk_id, "content_with_weight": content}
d["content_ltks"] = rag_tokenizer.tokenize(d["content_with_weight"])
d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
if "important_keywords" in req:
if not isinstance(req["important_keywords"], list):
return get_error_data_result("`important_keywords` should be a list")
d["important_kwd"] = req.get("important_keywords", [])
d["important_tks"] = rag_tokenizer.tokenize(" ".join(req["important_keywords"]))
if "questions" in req:
if not isinstance(req["questions"], list):
return get_error_data_result("`questions` should be a list")
d["question_kwd"] = [str(q).strip() for q in req.get("questions", []) if str(q).strip()]
d["question_tks"] = rag_tokenizer.tokenize("\n".join(req["questions"]))
if "available" in req:
d["available_int"] = int(req["available"])
if "positions" in req:
if not isinstance(req["positions"], list):
return get_error_data_result("`positions` should be a list")
d["position_int"] = req["positions"]
if "tag_kwd" in req:
if not isinstance(req["tag_kwd"], list):
return get_error_data_result("`tag_kwd` should be a list")
if not all(isinstance(t, str) for t in req["tag_kwd"]):
return get_error_data_result("`tag_kwd` must be a list of strings")
d["tag_kwd"] = req["tag_kwd"]
if "tag_feas" in req:
try:
d["tag_feas"] = validate_tag_features(req["tag_feas"])
except ValueError as exc:
return get_error_data_result(f"`tag_feas` {exc}")
image_base64 = req.get("image_base64")
if image_base64:
d["img_id"] = f"{dataset_id}-{chunk_id}"
d["doc_type_kwd"] = "image"
tenant_embd_id = DocumentService.get_tenant_embd_id(document_id)
if tenant_embd_id:
model_config = get_model_config_by_id(tenant_embd_id)
else:
embd_id = DocumentService.get_embd_id(document_id)
model_config = get_model_config_by_type_and_name(tenant_id, LLMType.EMBEDDING.value, embd_id)
embd_mdl = TenantLLMService.model_instance(model_config)
if doc.parser_id == ParserType.QA:
arr = [t for t in re.split(r"[\n\t]", d["content_with_weight"]) if len(t) > 1]
if len(arr) != 2:
return get_error_data_result(message="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, _ = embd_mdl.encode(
[
doc.name,
d["content_with_weight"] if not d.get("question_kwd") else "\n".join(d["question_kwd"]),
]
)
v = 0.1 * v[0] + 0.9 * v[1] if doc.parser_id != ParserType.QA else v[1]
d[f"q_{len(v)}_vec"] = v.tolist()
settings.docStoreConn.update({"id": chunk_id}, d, search.index_name(tenant_id), dataset_id)
if image_base64:
store_chunk_image(dataset_id, chunk_id, base64.b64decode(image_base64))
return get_result()
@manager.route("/datasets/<dataset_id>/documents/<document_id>/chunks", methods=["PATCH"]) # noqa: F821
@login_required
@add_tenant_id_to_kwargs
async def switch_chunks(tenant_id, dataset_id, document_id):
if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id):
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
req = await get_request_json()
if not req.get("chunk_ids"):
return get_error_data_result(message="`chunk_ids` is required.")
if "available_int" not in req and "available" not in req:
return get_error_data_result(message="`available_int` or `available` is required.")
available_int = int(req["available_int"]) if "available_int" in req else (1 if req.get("available") else 0)
try:
def _switch_sync():
e, doc = DocumentService.get_by_id(document_id)
if not e:
return get_error_data_result(message="Document not found!")
if not doc or str(doc.kb_id) != str(dataset_id):
return get_error_data_result(message="Document not found!")
for cid in req["chunk_ids"]:
if not settings.docStoreConn.update(
{"id": cid},
{"available_int": available_int},
search.index_name(tenant_id),
doc.kb_id,
):
return get_error_data_result(message="Index updating failure")
return get_result(data=True)
return await thread_pool_exec(_switch_sync)
except Exception as e:
return server_error_response(e)

670
api/apps/sdk/doc.py
View File

@@ -13,12 +13,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
import datetime
import re
from io import BytesIO
import xxhash
from pydantic import BaseModel, Field, validator
from quart import request, send_file
from api.db.db_models import APIToken, Document, Task
@@ -31,42 +27,16 @@ from api.db.services.llm_service import LLMBundle
from api.db.services.task_service import TaskService, cancel_all_task_of, queue_tasks
from api.db.services.tenant_llm_service import TenantLLMService
from api.utils.api_utils import check_duplicate_ids, construct_json_result, get_error_data_result, get_request_json, get_result, server_error_response, token_required
from api.utils.image_utils import store_chunk_image
from common import settings
from common.constants import LLMType, ParserType, RetCode, TaskStatus
from common.constants import LLMType, RetCode, TaskStatus
from common.metadata_utils import convert_conditions, meta_filter
from common.misc_utils import thread_pool_exec
from common.string_utils import is_content_empty, remove_redundant_spaces
from common.tag_feature_utils import validate_tag_features
from rag.app.qa import beAdoc, rmPrefix
from rag.app.tag import label_question
from rag.nlp import rag_tokenizer, search
from rag.nlp import search
from rag.prompts.generator import cross_languages, keyword_extraction
MAXIMUM_OF_UPLOADING_FILES = 256
class Chunk(BaseModel):
id: str = ""
content: str = ""
document_id: str = ""
docnm_kwd: str = ""
important_keywords: list = Field(default_factory=list)
tag_kwd: list = Field(default_factory=list)
questions: list = Field(default_factory=list)
question_tks: str = ""
image_id: str = ""
available: bool = True
positions: list[list[int]] = Field(default_factory=list)
@validator("positions")
def validate_positions(cls, value):
for sublist in value:
if len(sublist) != 5:
raise ValueError("Each sublist in positions must have a length of 5")
return value
@manager.route("/datasets/<dataset_id>/documents/<document_id>", methods=["GET"]) # noqa: F821
@token_required
async def download(tenant_id, dataset_id, document_id):
@@ -329,642 +299,6 @@ async def stop_parsing(tenant_id, dataset_id):
return get_result()
@manager.route("/datasets/<dataset_id>/documents/<document_id>/chunks", methods=["GET"]) # noqa: F821
@token_required
async def list_chunks(tenant_id, dataset_id, document_id):
"""
List chunks of a document.
---
tags:
- Chunks
security:
- ApiKeyAuth: []
parameters:
- in: path
name: dataset_id
type: string
required: true
description: ID of the dataset.
- in: path
name: document_id
type: string
required: true
description: ID of the document.
- in: query
name: page
type: integer
required: false
default: 1
description: Page number.
- in: query
name: page_size
type: integer
required: false
default: 30
description: Number of items per page.
- in: query
name: id
type: string
required: false
default: ""
description: Chunk id.
- in: header
name: Authorization
type: string
required: true
description: Bearer token for authentication.
responses:
200:
description: List of chunks.
schema:
type: object
properties:
total:
type: integer
description: Total number of chunks.
chunks:
type: array
items:
type: object
properties:
id:
type: string
description: Chunk ID.
content:
type: string
description: Chunk content.
document_id:
type: string
description: ID of the document.
important_keywords:
type: array
items:
type: string
description: Important keywords.
tag_kwd:
type: array
items:
type: string
description: Tag keywords.
image_id:
type: string
description: Image ID associated with the chunk.
doc:
type: object
description: Document details.
"""
if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id):
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
doc = DocumentService.query(id=document_id, kb_id=dataset_id)
if not doc:
return get_error_data_result(message=f"You don't own the document {document_id}.")
doc = doc[0]
req = request.args
doc_id = document_id
page = int(req.get("page", 1))
size = int(req.get("page_size", 30))
question = req.get("keywords", "")
query = {
"doc_ids": [doc_id],
"page": page,
"size": size,
"question": question,
"sort": True,
}
if "available" in req:
query["available_int"] = 1 if req["available"] == "true" else 0
key_mapping = {
"chunk_num": "chunk_count",
"kb_id": "dataset_id",
"token_num": "token_count",
"parser_id": "chunk_method",
}
run_mapping = {
"0": "UNSTART",
"1": "RUNNING",
"2": "CANCEL",
"3": "DONE",
"4": "FAIL",
}
doc = doc.to_dict()
renamed_doc = {}
for key, value in doc.items():
new_key = key_mapping.get(key, key)
renamed_doc[new_key] = value
if key == "run":
renamed_doc["run"] = run_mapping.get(str(value))
res = {"total": 0, "chunks": [], "doc": renamed_doc}
if req.get("id"):
chunk = settings.docStoreConn.get(req.get("id"), search.index_name(tenant_id), [dataset_id])
if not chunk:
return get_result(message=f"Chunk not found: {dataset_id}/{req.get('id')}", code=RetCode.NOT_FOUND)
k = []
for n in chunk.keys():
if re.search(r"(_vec$|_sm_|_tks|_ltks)", n):
k.append(n)
for n in k:
del chunk[n]
if not chunk:
return get_error_data_result(f"Chunk `{req.get('id')}` not found.")
res["total"] = 1
final_chunk = {
"id": chunk.get("id", chunk.get("chunk_id")),
"content": chunk["content_with_weight"],
"document_id": chunk.get("doc_id", chunk.get("document_id")),
"docnm_kwd": chunk["docnm_kwd"],
"important_keywords": chunk.get("important_kwd", []),
"questions": chunk.get("question_kwd", []),
"dataset_id": chunk.get("kb_id", chunk.get("dataset_id")),
"image_id": chunk.get("img_id", ""),
"available": bool(chunk.get("available_int", 1)),
"positions": chunk.get("position_int", []),
"tag_kwd": chunk.get("tag_kwd", []),
"tag_feas": chunk.get("tag_feas", {}),
}
res["chunks"].append(final_chunk)
_ = Chunk(**final_chunk)
elif settings.docStoreConn.index_exist(search.index_name(tenant_id), dataset_id):
sres = await settings.retriever.search(query, search.index_name(tenant_id), [dataset_id], emb_mdl=None, highlight=True)
res["total"] = sres.total
for id in sres.ids:
d = {
"id": id,
"content": (remove_redundant_spaces(sres.highlight[id]) if question and id in sres.highlight else sres.field[id].get("content_with_weight", "")),
"document_id": sres.field[id]["doc_id"],
"docnm_kwd": sres.field[id]["docnm_kwd"],
"important_keywords": sres.field[id].get("important_kwd", []),
"tag_kwd": sres.field[id].get("tag_kwd", []),
"questions": sres.field[id].get("question_kwd", []),
"dataset_id": sres.field[id].get("kb_id", sres.field[id].get("dataset_id")),
"image_id": sres.field[id].get("img_id", ""),
"available": bool(int(sres.field[id].get("available_int", "1"))),
"positions": sres.field[id].get("position_int", []),
}
res["chunks"].append(d)
_ = Chunk(**d) # validate the chunk
return get_result(data=res)
@manager.route( # noqa: F821
"/datasets/<dataset_id>/documents/<document_id>/chunks", methods=["POST"]
)
@token_required
async def add_chunk(tenant_id, dataset_id, document_id):
"""
Add a chunk to a document.
---
tags:
- Chunks
security:
- ApiKeyAuth: []
parameters:
- in: path
name: dataset_id
type: string
required: true
description: ID of the dataset.
- in: path
name: document_id
type: string
required: true
description: ID of the document.
- in: body
name: body
description: Chunk data.
required: true
schema:
type: object
properties:
content:
type: string
required: true
description: Content of the chunk.
important_keywords:
type: array
items:
type: string
description: Important keywords.
image_base64:
type: string
description: Base64-encoded image to associate with the chunk.
- in: header
name: Authorization
type: string
required: true
description: Bearer token for authentication.
responses:
200:
description: Chunk added successfully.
schema:
type: object
properties:
chunk:
type: object
properties:
id:
type: string
description: Chunk ID.
content:
type: string
description: Chunk content.
document_id:
type: string
description: ID of the document.
important_keywords:
type: array
items:
type: string
description: Important keywords.
"""
if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id):
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
doc = DocumentService.query(id=document_id, kb_id=dataset_id)
if not doc:
return get_error_data_result(message=f"You don't own the document {document_id}.")
doc = doc[0]
req = await get_request_json()
if is_content_empty(req.get("content")):
return get_error_data_result(message="`content` is required")
if "important_keywords" in req:
if not isinstance(req["important_keywords"], list):
return get_error_data_result("`important_keywords` is required to be a list")
if "questions" in req:
if not isinstance(req["questions"], list):
return get_error_data_result("`questions` is required to be a list")
chunk_id = xxhash.xxh64((req["content"] + document_id).encode("utf-8")).hexdigest()
d = {
"id": chunk_id,
"content_ltks": rag_tokenizer.tokenize(req["content"]),
"content_with_weight": req["content"],
}
d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
d["important_kwd"] = req.get("important_keywords", [])
d["important_tks"] = rag_tokenizer.tokenize(" ".join(req.get("important_keywords", [])))
d["question_kwd"] = [str(q).strip() for q in req.get("questions", []) if str(q).strip()]
d["question_tks"] = rag_tokenizer.tokenize("\n".join(req.get("questions", [])))
d["create_time"] = str(datetime.datetime.now()).replace("T", " ")[:19]
d["create_timestamp_flt"] = datetime.datetime.now().timestamp()
d["kb_id"] = dataset_id
d["docnm_kwd"] = doc.name
d["doc_id"] = document_id
if "tag_kwd" in req:
if not isinstance(req["tag_kwd"], list):
return get_error_data_result("`tag_kwd` is required to be a list")
if not all(isinstance(t, str) for t in req["tag_kwd"]):
return get_error_data_result("`tag_kwd` must be a list of strings")
d["tag_kwd"] = req["tag_kwd"]
if "tag_feas" in req:
try:
d["tag_feas"] = validate_tag_features(req["tag_feas"])
except ValueError as exc:
return get_error_data_result(f"`tag_feas` {exc}")
import base64
image_base64 = req.get("image_base64", None)
if image_base64:
d["img_id"] = "{}-{}".format(dataset_id, chunk_id)
d["doc_type_kwd"] = "image"
tenant_embd_id = DocumentService.get_tenant_embd_id(document_id)
if tenant_embd_id:
model_config = get_model_config_by_id(tenant_embd_id)
else:
embd_id = DocumentService.get_embd_id(document_id)
model_config = get_model_config_by_type_and_name(tenant_id, LLMType.EMBEDDING.value, embd_id)
embd_mdl = TenantLLMService.model_instance(model_config)
v, c = embd_mdl.encode([doc.name, req["content"] if not d["question_kwd"] else "\n".join(d["question_kwd"])])
v = 0.1 * v[0] + 0.9 * v[1]
d["q_%d_vec" % len(v)] = v.tolist()
settings.docStoreConn.insert([d], search.index_name(tenant_id), dataset_id)
if image_base64:
store_chunk_image(dataset_id, chunk_id, base64.b64decode(image_base64))
DocumentService.increment_chunk_num(doc.id, doc.kb_id, c, 1, 0)
# rename keys
key_mapping = {
"id": "id",
"content_with_weight": "content",
"doc_id": "document_id",
"important_kwd": "important_keywords",
"tag_kwd": "tag_kwd",
"question_kwd": "questions",
"kb_id": "dataset_id",
"create_timestamp_flt": "create_timestamp",
"create_time": "create_time",
"document_keyword": "document",
"img_id": "image_id",
}
renamed_chunk = {}
for key, value in d.items():
if key in key_mapping:
new_key = key_mapping.get(key, key)
renamed_chunk[new_key] = value
_ = Chunk(**renamed_chunk) # validate the chunk
return get_result(data={"chunk": renamed_chunk})
# return get_result(data={"chunk_id": chunk_id})
@manager.route( # noqa: F821
"datasets/<dataset_id>/documents/<document_id>/chunks", methods=["DELETE"]
)
@token_required
async def rm_chunk(tenant_id, dataset_id, document_id):
"""
Remove chunks from a document.
---
tags:
- Chunks
security:
- ApiKeyAuth: []
parameters:
- in: path
name: dataset_id
type: string
required: true
description: ID of the dataset.
- in: path
name: document_id
type: string
required: true
description: ID of the document.
- in: body
name: body
description: Chunk removal parameters.
required: true
schema:
type: object
properties:
chunk_ids:
type: array
items:
type: string
description: |
List of chunk IDs to remove.
If omitted, `null`, or an empty array is provided, no chunks will be deleted.
- in: header
name: Authorization
type: string
required: true
description: Bearer token for authentication.
responses:
200:
description: Chunks removed successfully.
schema:
type: object
"""
if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id):
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
docs = DocumentService.get_by_ids([document_id])
if not docs:
raise LookupError(f"Can't find the document with ID {document_id}!")
req = await get_request_json()
if not req:
return get_result()
chunk_ids = req.get("chunk_ids")
if not chunk_ids:
if req.get("delete_all") is True:
doc = docs[0]
# Clean up storage assets while index rows still exist for discovery
DocumentService.delete_chunk_images(doc, tenant_id)
condition = {"doc_id": document_id}
chunk_number = settings.docStoreConn.delete(condition, search.index_name(tenant_id), dataset_id)
if chunk_number != 0:
DocumentService.decrement_chunk_num(document_id, dataset_id, 1, chunk_number, 0)
return get_result(message=f"deleted {chunk_number} chunks")
else:
return get_result()
condition = {"doc_id": document_id}
unique_chunk_ids, duplicate_messages = check_duplicate_ids(chunk_ids, "chunk")
condition["id"] = unique_chunk_ids
chunk_number = settings.docStoreConn.delete(condition, search.index_name(tenant_id), dataset_id)
if chunk_number != 0:
DocumentService.decrement_chunk_num(document_id, dataset_id, 1, chunk_number, 0)
if chunk_number != len(unique_chunk_ids):
if len(unique_chunk_ids) == 0:
return get_result(message=f"deleted {chunk_number} chunks")
return get_error_data_result(message=f"rm_chunk deleted chunks {chunk_number}, expect {len(unique_chunk_ids)}")
if duplicate_messages:
return get_result(
message=f"Partially deleted {chunk_number} chunks with {len(duplicate_messages)} errors",
data={"success_count": chunk_number, "errors": duplicate_messages},
)
return get_result(message=f"deleted {chunk_number} chunks")
@manager.route( # noqa: F821
"/datasets/<dataset_id>/documents/<document_id>/chunks/<chunk_id>", methods=["PUT"]
)
@token_required
async def update_chunk(tenant_id, dataset_id, document_id, chunk_id):
"""
Update a chunk within a document.
---
tags:
- Chunks
security:
- ApiKeyAuth: []
parameters:
- in: path
name: dataset_id
type: string
required: true
description: ID of the dataset.
- in: path
name: document_id
type: string
required: true
description: ID of the document.
- in: path
name: chunk_id
type: string
required: true
description: ID of the chunk to update.
- in: body
name: body
description: Chunk update parameters.
required: true
schema:
type: object
properties:
content:
type: string
description: Updated content of the chunk.
important_keywords:
type: array
items:
type: string
description: Updated important keywords.
tag_kwd:
type: array
items:
type: string
description: Updated tag keywords.
available:
type: boolean
description: Availability status of the chunk.
- in: header
name: Authorization
type: string
required: true
description: Bearer token for authentication.
responses:
200:
description: Chunk updated successfully.
schema:
type: object
"""
chunk = settings.docStoreConn.get(chunk_id, search.index_name(tenant_id), [dataset_id])
if chunk is None:
return get_error_data_result(f"Can't find this chunk {chunk_id}")
if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id):
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
doc = DocumentService.query(id=document_id, kb_id=dataset_id)
if not doc:
return get_error_data_result(message=f"You don't own the document {document_id}.")
doc = doc[0]
req = await get_request_json()
content = req.get("content")
if content is not None:
if is_content_empty(content):
return get_error_data_result(message="`content` is required")
else:
content = chunk.get("content_with_weight", "")
d = {"id": chunk_id, "content_with_weight": content}
d["content_ltks"] = rag_tokenizer.tokenize(d["content_with_weight"])
d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
if "important_keywords" in req:
if not isinstance(req["important_keywords"], list):
return get_error_data_result("`important_keywords` should be a list")
d["important_kwd"] = req.get("important_keywords", [])
d["important_tks"] = rag_tokenizer.tokenize(" ".join(req["important_keywords"]))
if "questions" in req:
if not isinstance(req["questions"], list):
return get_error_data_result("`questions` should be a list")
d["question_kwd"] = [str(q).strip() for q in req.get("questions", []) if str(q).strip()]
d["question_tks"] = rag_tokenizer.tokenize("\n".join(req["questions"]))
if "available" in req:
d["available_int"] = int(req["available"])
if "positions" in req:
if not isinstance(req["positions"], list):
return get_error_data_result("`positions` should be a list")
d["position_int"] = req["positions"]
if "tag_kwd" in req:
if not isinstance(req["tag_kwd"], list):
return get_error_data_result("`tag_kwd` should be a list")
if not all(isinstance(t, str) for t in req["tag_kwd"]):
return get_error_data_result("`tag_kwd` must be a list of strings")
d["tag_kwd"] = req["tag_kwd"]
if "tag_feas" in req:
try:
d["tag_feas"] = validate_tag_features(req["tag_feas"])
except ValueError as exc:
return get_error_data_result(f"`tag_feas` {exc}")
tenant_embd_id = DocumentService.get_tenant_embd_id(document_id)
if tenant_embd_id:
model_config = get_model_config_by_id(tenant_embd_id)
else:
embd_id = DocumentService.get_embd_id(document_id)
model_config = get_model_config_by_type_and_name(tenant_id, LLMType.EMBEDDING.value, embd_id)
embd_mdl = TenantLLMService.model_instance(model_config)
if doc.parser_id == ParserType.QA:
arr = [t for t in re.split(r"[\n\t]", d["content_with_weight"]) if len(t) > 1]
if len(arr) != 2:
return get_error_data_result(message="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, d["content_with_weight"] if not d.get("question_kwd") else "\n".join(d["question_kwd"])])
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()
settings.docStoreConn.update({"id": chunk_id}, d, search.index_name(tenant_id), dataset_id)
return get_result()
@manager.route( # noqa: F821
"/datasets/<dataset_id>/documents/<document_id>/chunks/switch", methods=["POST"]
)
@token_required
async def switch_chunks(tenant_id, dataset_id, document_id):
"""
Switch availability of specified chunks (same as chunk_app switch).
---
tags:
- Chunks
security:
- ApiKeyAuth: []
parameters:
- in: path
name: dataset_id
type: string
required: true
description: ID of the dataset.
- in: path
name: document_id
type: string
required: true
description: ID of the document.
- in: body
name: body
required: true
schema:
type: object
properties:
chunk_ids:
type: array
items:
type: string
description: List of chunk IDs to switch.
available_int:
type: integer
description: 1 for available, 0 for unavailable.
available:
type: boolean
description: Availability status (alternative to available_int).
- in: header
name: Authorization
type: string
required: true
description: Bearer token for authentication.
responses:
200:
description: Chunks availability switched successfully.
"""
if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id):
return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
req = await get_request_json()
if not req.get("chunk_ids"):
return get_error_data_result(message="`chunk_ids` is required.")
if "available_int" not in req and "available" not in req:
return get_error_data_result(message="`available_int` or `available` is required.")
available_int = int(req["available_int"]) if "available_int" in req else (1 if req.get("available") else 0)
try:
def _switch_sync():
e, doc = DocumentService.get_by_id(document_id)
if not e:
return get_error_data_result(message="Document not found!")
if not doc or str(doc.kb_id) != str(dataset_id):
return get_error_data_result(message="Document not found!")
for cid in req["chunk_ids"]:
if not settings.docStoreConn.update(
{"id": cid},
{"available_int": available_int},
search.index_name(tenant_id),
doc.kb_id,
):
return get_error_data_result(message="Index updating failure")
return get_result(data=True)
return await thread_pool_exec(_switch_sync)
except Exception as e:
return server_error_response(e)
@manager.route("/retrieval", methods=["POST"]) # noqa: F821
@token_required
async def retrieval_test(tenant_id):

728
docs/references/http_api_reference.md
View File

File diff suppressed because it is too large Load Diff

12
docs/references/python_api_reference.md
View File

@@ -855,7 +855,7 @@ print("Async bulk parsing cancelled.")
### Add chunk
```python
Document.add_chunk(content:str, important_keywords:list[str] = [], image_base64:str = None, *, tag_kwd:list[str] = []) -> Chunk
Document.add_chunk(content:str, important_keywords:list[str] = [], questions:list[str] = [], image_base64:str = None, *, tag_kwd:list[str] = []) -> Chunk
```
Adds a chunk to the current document.
@@ -870,6 +870,10 @@ The text content of the chunk.
The key terms or phrases to tag with the chunk.
##### questions: `list[str]`
Optional questions to use when embedding the chunk.
##### image_base64: `string`
A base64-encoded image to associate with the chunk. If the chunk already has an image, the new image will be vertically concatenated below the existing one.
@@ -889,6 +893,7 @@ A `Chunk` object contains the following attributes:
- `content`: `string` The text content of the chunk.
- `important_keywords`: `list[str]` A list of key terms or phrases tagged with the chunk.
- `tag_kwd`: `list[str]` A list of tag keywords associated with the chunk.
- `questions`: `list[str]` A list of questions associated with the chunk.
- `image_id`: `string` The image ID associated with the chunk (empty string if no image).
- `create_time`: `string` The time when the chunk was created (added to the document).
- `create_timestamp`: `float` The timestamp representing the creation time of the chunk, expressed in seconds since January 1, 1970.
@@ -1023,16 +1028,19 @@ Updates content or configurations for the current chunk.
#### Parameters
##### update_message: `dict[str, str|list[str]|int]` *Required*
##### update_message: `dict[str, str|list[str]|bool]` *Required*
A dictionary representing the attributes to update, with the following keys:
- `"content"`: `string` The text content of the chunk.
- `"important_keywords"`: `list[str]` A list of key terms or phrases to tag with the chunk.
- `"questions"`: `list[str]` A list of questions associated with the chunk.
- `"tag_kwd"`: `list[str]` A list of tag keywords to associate with the chunk.
- `"positions"`: `list` Updated source positions for the chunk.
- `"available"`: `bool` The chunk's availability status in the dataset. Value options:
- `False`: Unavailable
- `True`: Available (default)
- `"image_base64"`: `string` Base64-encoded image content to associate with the chunk.
#### Returns

4
sdk/python/ragflow_sdk/modules/chunk.py
View File

@@ -54,11 +54,11 @@ class Chunk(Base):
def update(self, update_message: dict):
res = self.put(f"/datasets/{self.dataset_id}/documents/{self.document_id}/chunks/{self.id}", update_message)
res = self.patch(f"/datasets/{self.dataset_id}/documents/{self.document_id}/chunks/{self.id}", update_message)
res = res.json()
if res.get("code") != 0:
raise ChunkUpdateError(
code=res.get("code"),
message=res.get("message"),
details=res.get("details")
)
)

8
test/testcases/test_http_api/common.py
View File

@@ -173,9 +173,15 @@ def list_chunks(auth, dataset_id, document_id, params=None):
return res.json()
def get_chunk(auth, dataset_id, document_id, chunk_id):
url = f"{HOST_ADDRESS}{CHUNK_API_URL}/{chunk_id}".format(dataset_id=dataset_id, document_id=document_id)
res = requests.get(url=url, headers=HEADERS, auth=auth)
return res.json()
def update_chunk(auth, dataset_id, document_id, chunk_id, payload=None):
url = f"{HOST_ADDRESS}{CHUNK_API_URL}/{chunk_id}".format(dataset_id=dataset_id, document_id=document_id)
res = requests.put(url=url, headers=HEADERS, auth=auth, json=payload)
res = requests.patch(url=url, headers=HEADERS, auth=auth, json=payload)
return res.json()

24
test/testcases/test_http_api/test_chunk_management_within_dataset/conftest.py
View File

@@ -18,17 +18,20 @@
from time import sleep
import pytest
from common import batch_add_chunks, delete_all_chunks, list_documents, parse_documents
from utils import wait_for
from common import add_chunk, batch_add_chunks, delete_all_chunks
@wait_for(30, 1, "Document parsing timeout")
def condition(_auth, _dataset_id):
res = list_documents(_auth, _dataset_id)
for doc in res["data"]["docs"]:
if doc["run"] != "DONE":
return False
return True
def _add_baseline_chunk(auth, dataset_id, document_id):
add_chunk(auth, dataset_id, document_id, {"content": "ragflow test upload"})
@pytest.fixture(scope="class")
def add_chunks(HttpApiAuth, add_document):
dataset_id, document_id = add_document
_add_baseline_chunk(HttpApiAuth, dataset_id, document_id)
chunk_ids = batch_add_chunks(HttpApiAuth, dataset_id, document_id, 4)
sleep(1) # issues/6487
return dataset_id, document_id, chunk_ids
@pytest.fixture(scope="function")
@@ -39,8 +42,7 @@ def add_chunks_func(request, HttpApiAuth, add_document):
request.addfinalizer(cleanup)
dataset_id, document_id = add_document
parse_documents(HttpApiAuth, dataset_id, {"document_ids": [document_id]})
condition(HttpApiAuth, dataset_id)
_add_baseline_chunk(HttpApiAuth, dataset_id, document_id)
chunk_ids = batch_add_chunks(HttpApiAuth, dataset_id, document_id, 4)
# issues/6487
sleep(1)

8
test/testcases/test_http_api/test_chunk_management_within_dataset/test_add_chunk.py
View File

@@ -39,12 +39,8 @@ class TestAuthorization:
@pytest.mark.parametrize(
"invalid_auth, expected_code, expected_message",
[
(None, 0, "`Authorization` can't be empty"),
(
RAGFlowHttpApiAuth(INVALID_API_TOKEN),
109,
"Authentication error: API key is invalid!",
),
(None, 401, "<Unauthorized '401: Unauthorized'>"),
(RAGFlowHttpApiAuth(INVALID_API_TOKEN), 401, "<Unauthorized '401: Unauthorized'>"),
],
)
def test_invalid_auth(self, invalid_auth, expected_code, expected_message):

10
test/testcases/test_http_api/test_chunk_management_within_dataset/test_delete_chunks.py
View File

@@ -26,12 +26,8 @@ class TestAuthorization:
@pytest.mark.parametrize(
"invalid_auth, expected_code, expected_message",
[
(None, 0, "`Authorization` can't be empty"),
(
RAGFlowHttpApiAuth(INVALID_API_TOKEN),
109,
"Authentication error: API key is invalid!",
),
(None, 401, "<Unauthorized '401: Unauthorized'>"),
(RAGFlowHttpApiAuth(INVALID_API_TOKEN), 401, "<Unauthorized '401: Unauthorized'>"),
],
)
def test_invalid_auth(self, invalid_auth, expected_code, expected_message):
@@ -58,7 +54,7 @@ class TestChunksDeletion:
@pytest.mark.parametrize(
"document_id, expected_code, expected_message",
[
(INVALID_ID_32, 100, f"""LookupError("Can't find the document with ID {INVALID_ID_32}!")"""),
(INVALID_ID_32, 102, f"You don't own the document {INVALID_ID_32}."),
],
)
def test_invalid_document_id(self, HttpApiAuth, add_chunks_func, document_id, expected_code, expected_message):

19
test/testcases/test_http_api/test_chunk_management_within_dataset/test_list_chunks.py
View File

@@ -17,7 +17,7 @@ import os
from concurrent.futures import ThreadPoolExecutor, as_completed
import pytest
from common import batch_add_chunks, list_chunks
from common import batch_add_chunks, get_chunk, list_chunks
from configs import INVALID_API_TOKEN, INVALID_ID_32
from libs.auth import RAGFlowHttpApiAuth
@@ -27,12 +27,8 @@ class TestAuthorization:
@pytest.mark.parametrize(
"invalid_auth, expected_code, expected_message",
[
(None, 0, "`Authorization` can't be empty"),
(
RAGFlowHttpApiAuth(INVALID_API_TOKEN),
109,
"Authentication error: API key is invalid!",
),
(None, 401, "<Unauthorized '401: Unauthorized'>"),
(RAGFlowHttpApiAuth(INVALID_API_TOKEN), 401, "<Unauthorized '401: Unauthorized'>"),
],
)
def test_invalid_auth(self, invalid_auth, expected_code, expected_message):
@@ -139,6 +135,15 @@ class TestChunksList:
else:
assert res["message"] == expected_message
@pytest.mark.p1
@pytest.mark.skipif(os.getenv("DOC_ENGINE") == "infinity", reason="issues/6499")
def test_get_chunk(self, HttpApiAuth, add_chunks):
dataset_id, document_id, chunk_ids = add_chunks
res = get_chunk(HttpApiAuth, dataset_id, document_id, chunk_ids[0])
assert res["code"] == 0
assert res["data"]["id"] == chunk_ids[0]
assert res["data"]["doc_id"] == document_id
@pytest.mark.p3
def test_invalid_params(self, HttpApiAuth, add_chunks):
dataset_id, document_id, _ = add_chunks

8
test/testcases/test_http_api/test_chunk_management_within_dataset/test_update_chunk.py
View File

@@ -28,12 +28,8 @@ class TestAuthorization:
@pytest.mark.parametrize(
"invalid_auth, expected_code, expected_message",
[
(None, 0, "`Authorization` can't be empty"),
(
RAGFlowHttpApiAuth(INVALID_API_TOKEN),
109,
"Authentication error: API key is invalid!",
),
(None, 401, "<Unauthorized '401: Unauthorized'>"),
(RAGFlowHttpApiAuth(INVALID_API_TOKEN), 401, "<Unauthorized '401: Unauthorized'>"),
],
)
def test_invalid_auth(self, invalid_auth, expected_code, expected_message):

104
test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py
View File

@@ -14,6 +14,7 @@
# limitations under the License.
#
import asyncio
import inspect
import importlib.util
import sys
from pathlib import Path
@@ -309,6 +310,19 @@ def _load_doc_module(monkeypatch):
return module
def _load_restful_chunk_module(monkeypatch):
repo_root = Path(__file__).resolve().parents[4]
helper_path = repo_root / "test" / "testcases" / "test_web_api" / "test_chunk_app" / "test_chunk_routes_unit.py"
spec = importlib.util.spec_from_file_location("test_restful_chunk_route_helpers", helper_path)
helper = importlib.util.module_from_spec(spec)
spec.loader.exec_module(helper)
return helper._load_chunk_api_module(monkeypatch)
def _route_core(func):
return inspect.unwrap(func)
def _patch_send_file(monkeypatch, module):
async def _fake_send_file(file_obj, **kwargs):
return {"file": file_obj, "filename": kwargs.get("attachment_filename")}
@@ -336,7 +350,7 @@ def _patch_docstore(monkeypatch, module, **kwargs):
@pytest.mark.p2
class TestDocRoutesUnit:
def test_chunk_positions_validation_error(self, monkeypatch):
module = _load_doc_module(monkeypatch)
module = _load_restful_chunk_module(monkeypatch)
with pytest.raises(ValueError) as exc_info:
module.Chunk(positions=[[1, 2, 3, 4]])
assert "length of 5" in str(exc_info.value)
@@ -484,25 +498,44 @@ class TestDocRoutesUnit:
assert res["code"] == 0
def test_list_chunks_branches(self, monkeypatch):
module = _load_doc_module(monkeypatch)
module = _load_restful_chunk_module(monkeypatch)
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False)
res = _run(module.list_chunks.__wrapped__("tenant-1", "ds-1", "doc-1"))
res = _run(_route_core(module.list_chunks)("tenant-1", "ds-1", "doc-1"))
assert "don't own the dataset" in res["message"]
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True)
monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [])
res = _run(module.list_chunks.__wrapped__("tenant-1", "ds-1", "doc-1"))
res = _run(_route_core(module.list_chunks)("tenant-1", "ds-1", "doc-1"))
assert "don't own the document" in res["message"]
monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc()])
monkeypatch.setattr(module, "request", SimpleNamespace(args=_DummyArgs({})))
_patch_docstore(monkeypatch, module, index_exist=lambda *_args, **_kwargs: False)
res = _run(_route_core(module.list_chunks)("tenant-1", "ds-1", "doc-1"))
assert res["code"] == 0
assert res["data"]["total"] == 0
assert res["data"]["chunks"] == []
monkeypatch.setattr(module, "request", SimpleNamespace(args=_DummyArgs({"id": "chunk-1"})))
_patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: None)
res = _run(module.list_chunks.__wrapped__("tenant-1", "ds-1", "doc-1"))
res = _run(_route_core(module.list_chunks)("tenant-1", "ds-1", "doc-1"))
assert res["code"] == module.RetCode.DATA_ERROR
assert "Chunk not found" in res["message"]
_patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: {"id_vec": [1], "content_with_weight_vec": [2]})
res = _run(module.list_chunks.__wrapped__("tenant-1", "ds-1", "doc-1"))
assert "Chunk `chunk-1` not found." in res["message"]
_patch_docstore(
monkeypatch,
module,
get=lambda *_args, **_kwargs: {
"chunk_id": "chunk-1",
"content_with_weight": "x",
"doc_id": "other-doc",
"docnm_kwd": "doc",
"position_int": [[1, 2, 3, 4, 5]],
},
)
res = _run(_route_core(module.list_chunks)("tenant-1", "ds-1", "doc-1"))
assert res["code"] == module.RetCode.DATA_ERROR
assert "Chunk not found" in res["message"]
_patch_docstore(
monkeypatch,
@@ -515,29 +548,29 @@ class TestDocRoutesUnit:
"position_int": [[1, 2, 3, 4, 5]],
},
)
res = _run(module.list_chunks.__wrapped__("tenant-1", "ds-1", "doc-1"))
res = _run(_route_core(module.list_chunks)("tenant-1", "ds-1", "doc-1"))
assert res["code"] == 0
assert res["data"]["total"] == 1
assert res["data"]["chunks"][0]["id"] == "chunk-1"
def test_add_chunk_access_guard(self, monkeypatch):
module = _load_doc_module(monkeypatch)
module = _load_restful_chunk_module(monkeypatch)
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False)
res = _run(module.add_chunk.__wrapped__("tenant-1", "ds-1", "doc-1"))
res = _run(_route_core(module.add_chunk)("tenant-1", "ds-1", "doc-1"))
assert "don't own the dataset" in res["message"]
def test_rm_chunk_branches(self, monkeypatch):
module = _load_doc_module(monkeypatch)
module = _load_restful_chunk_module(monkeypatch)
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False)
res = _run(module.rm_chunk.__wrapped__("tenant-1", "ds-1", "doc-1"))
res = _run(_route_core(module.rm_chunk)("tenant-1", "ds-1", "doc-1"))
assert "don't own the dataset" in res["message"]
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True)
monkeypatch.setattr(module.DocumentService, "get_by_ids", lambda _ids: [])
with pytest.raises(LookupError):
_run(module.rm_chunk.__wrapped__("tenant-1", "ds-1", "doc-1"))
monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [])
res = _run(_route_core(module.rm_chunk)("tenant-1", "ds-1", "doc-1"))
assert "don't own the document" in res["message"]
monkeypatch.setattr(module.DocumentService, "get_by_ids", lambda _ids: [_DummyDoc()])
monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc()])
monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({}))
_patch_docstore(
monkeypatch,
@@ -545,32 +578,37 @@ class TestDocRoutesUnit:
delete=lambda *_args, **_kwargs: (_ for _ in ()).throw(AssertionError("delete must not run for empty chunk ids")),
)
monkeypatch.setattr(module.DocumentService, "decrement_chunk_num", lambda *_args, **_kwargs: None)
res = _run(module.rm_chunk.__wrapped__("tenant-1", "ds-1", "doc-1"))
res = _run(_route_core(module.rm_chunk)("tenant-1", "ds-1", "doc-1"))
assert res["code"] == 0
monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"chunk_ids": ["c1", "c1"]}))
monkeypatch.setattr(module, "check_duplicate_ids", lambda _ids, _kind: (["c1"], ["Duplicate chunk ids: c1"]))
_patch_docstore(monkeypatch, module, delete=lambda *_args, **_kwargs: 1)
res = _run(module.rm_chunk.__wrapped__("tenant-1", "ds-1", "doc-1"))
res = _run(_route_core(module.rm_chunk)("tenant-1", "ds-1", "doc-1"))
assert res["code"] == 0
assert res["data"]["errors"] == ["Duplicate chunk ids: c1"]
def test_update_chunk_branches(self, monkeypatch):
module = _load_doc_module(monkeypatch)
_patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: None)
res = _run(module.update_chunk.__wrapped__("tenant-1", "ds-1", "doc-1", "chunk-1"))
assert "Can't find this chunk" in res["message"]
_patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: {"content_with_weight": "q\na"})
module = _load_restful_chunk_module(monkeypatch)
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False)
res = _run(module.update_chunk.__wrapped__("tenant-1", "ds-1", "doc-1", "chunk-1"))
_patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: (_ for _ in ()).throw(AssertionError("chunk lookup must not run before access check")))
res = _run(_route_core(module.update_chunk)("tenant-1", "ds-1", "doc-1", "chunk-1"))
assert "don't own the dataset" in res["message"]
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True)
monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [])
res = _run(module.update_chunk.__wrapped__("tenant-1", "ds-1", "doc-1", "chunk-1"))
res = _run(_route_core(module.update_chunk)("tenant-1", "ds-1", "doc-1", "chunk-1"))
assert "don't own the document" in res["message"]
monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc()])
_patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: None)
res = _run(_route_core(module.update_chunk)("tenant-1", "ds-1", "doc-1", "chunk-1"))
assert "Can't find this chunk" in res["message"]
_patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: {"doc_id": "other-doc", "content_with_weight": "q\na"})
res = _run(_route_core(module.update_chunk)("tenant-1", "ds-1", "doc-1", "chunk-1"))
assert "Can't find this chunk" in res["message"]
doc = _DummyDoc(parser_id="naive")
monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [doc])
monkeypatch.setattr(module.rag_tokenizer, "tokenize", lambda text: text or "")
@@ -584,25 +622,25 @@ class TestDocRoutesUnit:
return [np.array([0.2, 0.8]), np.array([0.3, 0.7])], 1
monkeypatch.setattr(module.TenantLLMService, "model_instance", lambda *_args, **_kwargs: _EmbedModel())
_patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: {"doc_id": "doc-1", "content_with_weight": "x"}, update=lambda *_args, **_kwargs: None)
monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"positions": "bad"}))
res = _run(module.update_chunk.__wrapped__("tenant-1", "ds-1", "doc-1", "chunk-1"))
res = _run(_route_core(module.update_chunk)("tenant-1", "ds-1", "doc-1", "chunk-1"))
assert "`positions` should be a list" in res["message"]
_patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: {"content_with_weight": "x"}, update=lambda *_args, **_kwargs: None)
monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"positions": [[1, 2, 3, 4, 5]]}))
res = _run(module.update_chunk.__wrapped__("tenant-1", "ds-1", "doc-1", "chunk-1"))
res = _run(_route_core(module.update_chunk)("tenant-1", "ds-1", "doc-1", "chunk-1"))
assert res["code"] == 0
qa_doc = _DummyDoc(parser_id=module.ParserType.QA)
monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [qa_doc])
monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"content": "no-separator"}))
res = _run(module.update_chunk.__wrapped__("tenant-1", "ds-1", "doc-1", "chunk-1"))
res = _run(_route_core(module.update_chunk)("tenant-1", "ds-1", "doc-1", "chunk-1"))
assert "Q&A must be separated" in res["message"]
monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"content": "Q?\nA!"}))
_patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: {"content_with_weight": "Q?\nA!"}, update=lambda *_args, **_kwargs: None)
_patch_docstore(monkeypatch, module, get=lambda *_args, **_kwargs: {"doc_id": "doc-1", "content_with_weight": "Q?\nA!"}, update=lambda *_args, **_kwargs: None)
monkeypatch.setattr(module, "beAdoc", lambda d, *_args, **_kwargs: d)
res = _run(module.update_chunk.__wrapped__("tenant-1", "ds-1", "doc-1", "chunk-1"))
res = _run(_route_core(module.update_chunk)("tenant-1", "ds-1", "doc-1", "chunk-1"))
assert res["code"] == 0
def test_retrieval_validation_matrix(self, monkeypatch):

14
test/testcases/test_web_api/conftest.py
View File

@@ -157,17 +157,17 @@ def add_document(request, WebApiAuth, add_dataset, ragflow_tmp_dir):
@pytest.fixture(scope="class")
def add_chunks(request, WebApiAuth, add_document):
def cleanup():
res = list_chunks(WebApiAuth, {"doc_id": document_id})
res = list_chunks(WebApiAuth, dataset_id, document_id)
if res["code"] == 0:
chunk_ids = [chunk["chunk_id"] for chunk in res["data"]["chunks"]]
delete_chunks(WebApiAuth, {"doc_id": document_id, "chunk_ids": chunk_ids})
chunk_ids = [chunk["id"] for chunk in res["data"]["chunks"]]
delete_chunks(WebApiAuth, dataset_id, document_id, {"chunk_ids": chunk_ids})
request.addfinalizer(cleanup)
kb_id, document_id = add_document
dataset_id, document_id = add_document
parse_documents(WebApiAuth, {"doc_ids": [document_id], "run": "1"})
condition(WebApiAuth, kb_id)
chunk_ids = batch_add_chunks(WebApiAuth, document_id, 4)
condition(WebApiAuth, dataset_id)
chunk_ids = batch_add_chunks(WebApiAuth, dataset_id, document_id, 4)
# issues/6487
sleep(1)
return kb_id, document_id, chunk_ids
return dataset_id, document_id, chunk_ids

14
test/testcases/test_web_api/test_chunk_app/conftest.py
View File

@@ -34,16 +34,16 @@ def condition(_auth, _kb_id):
@pytest.fixture(scope="function")
def add_chunks_func(request, WebApiAuth, add_document):
def cleanup():
res = list_chunks(WebApiAuth, {"doc_id": document_id})
chunk_ids = [chunk["chunk_id"] for chunk in res["data"]["chunks"]]
delete_chunks(WebApiAuth, {"doc_id": document_id, "chunk_ids": chunk_ids})
res = list_chunks(WebApiAuth, dataset_id, document_id)
chunk_ids = [chunk["id"] for chunk in res["data"]["chunks"]]
delete_chunks(WebApiAuth, dataset_id, document_id, {"chunk_ids": chunk_ids})
request.addfinalizer(cleanup)
kb_id, document_id = add_document
dataset_id, document_id = add_document
parse_documents(WebApiAuth, {"doc_ids": [document_id], "run": "1"})
condition(WebApiAuth, kb_id)
chunk_ids = batch_add_chunks(WebApiAuth, document_id, 4)
condition(WebApiAuth, dataset_id)
chunk_ids = batch_add_chunks(WebApiAuth, dataset_id, document_id, 4)
# issues/6487
sleep(1)
return kb_id, document_id, chunk_ids
return dataset_id, document_id, chunk_ids

455
test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py
View File

@@ -15,7 +15,7 @@
#
import asyncio
import base64
import inspect
import importlib.util
import json
import sys
@@ -73,6 +73,7 @@ class _DummyRetCode:
DATA_ERROR = 102
EXCEPTION_ERROR = 100
OPERATING_ERROR = 103
NOT_FOUND = 404
class _DummyParserType:
@@ -81,7 +82,7 @@ class _DummyParserType:
class _DummyRetriever:
async def search(self, query, _index_name, _kb_ids, highlight=None):
async def search(self, query, _index_name, _kb_ids, *args, highlight=None, **kwargs):
class _SRes:
total = 1
ids = ["chunk-1"]
@@ -138,6 +139,9 @@ class _DummyDocStore:
def insert(self, docs, *_args, **_kwargs):
self.inserted.extend(docs)
def index_exist(self, *_args, **_kwargs):
return True
class _DummyStorage:
def __init__(self):
@@ -179,6 +183,10 @@ def _run(coro):
return asyncio.run(coro)
def _route_core(func):
return inspect.unwrap(func)
def _load_chunk_module(monkeypatch):
repo_root = Path(__file__).resolve().parents[4]
@@ -279,15 +287,33 @@ def _load_chunk_module(monkeypatch):
api_utils_mod = ModuleType("api.utils.api_utils")
api_utils_mod.get_json_result = lambda data=None, message="", code=0: {"code": code, "message": message, "data": data}
api_utils_mod.get_data_error_result = lambda message="": {"code": _DummyRetCode.DATA_ERROR, "message": message, "data": False}
api_utils_mod.get_result = lambda data=None, message="", code=0: {"code": code, "message": message, "data": data}
api_utils_mod.get_error_data_result = lambda message="": {"code": _DummyRetCode.DATA_ERROR, "message": message, "data": False}
api_utils_mod.server_error_response = lambda exc: {"code": _DummyRetCode.EXCEPTION_ERROR, "message": repr(exc), "data": False}
api_utils_mod.validate_request = lambda *_args, **_kwargs: (lambda fn: fn)
api_utils_mod.add_tenant_id_to_kwargs = lambda func: func
api_utils_mod.check_duplicate_ids = lambda ids, _kind: (list(dict.fromkeys(ids)), [] if len(ids) == len(set(ids)) else [f"Duplicate {_kind} ids"])
api_utils_mod.get_request_json = lambda: _AwaitableValue({})
monkeypatch.setitem(sys.modules, "api.utils.api_utils", api_utils_mod)
image_utils_mod = ModuleType("api.utils.image_utils")
image_utils_mod.store_chunk_image = lambda *_args, **_kwargs: None
monkeypatch.setitem(sys.modules, "api.utils.image_utils", image_utils_mod)
services_pkg = ModuleType("api.db.services")
services_pkg.__path__ = []
monkeypatch.setitem(sys.modules, "api.db.services", services_pkg)
joint_services_pkg = ModuleType("api.db.joint_services")
joint_services_pkg.__path__ = []
monkeypatch.setitem(sys.modules, "api.db.joint_services", joint_services_pkg)
tenant_model_service_mod = ModuleType("api.db.joint_services.tenant_model_service")
tenant_model_service_mod.get_model_config_by_id = lambda *_args, **_kwargs: {"llm_name": "embed", "model_type": "embedding"}
tenant_model_service_mod.get_model_config_by_type_and_name = lambda *_args, **_kwargs: {"llm_name": "embed", "model_type": "embedding"}
tenant_model_service_mod.get_tenant_default_model_by_type = lambda *_args, **_kwargs: {"llm_name": "chat", "model_type": "chat"}
monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_service_mod)
document_service_mod = ModuleType("api.db.services.document_service")
class _DocumentService:
@@ -302,6 +328,18 @@ def _load_chunk_module(monkeypatch):
def get_by_id(doc_id):
return True, _DummyDoc(doc_id=doc_id, parser_id=_DummyParserType.NAIVE)
@staticmethod
def query(**kwargs):
return [_DummyDoc(doc_id=kwargs.get("id", "doc-1"), kb_id=kwargs.get("kb_id", "kb-1"))]
@staticmethod
def get_by_ids(ids):
return [_DummyDoc(doc_id=ids[0] if ids else "doc-1")]
@staticmethod
def delete_chunk_images(*_args, **_kwargs):
return None
@staticmethod
def get_embd_id(_doc_id):
return "embed-1"
@@ -334,6 +372,10 @@ def _load_chunk_module(monkeypatch):
def get_kb_ids(_tenant_id):
return ["kb-1"]
@staticmethod
def accessible(**_kwargs):
return True
@staticmethod
def get_by_id(_kb_id):
return True, SimpleNamespace(pagerank=0.6, tenant_embd_id=2, tenant_llm_id=1)
@@ -415,6 +457,10 @@ def _load_chunk_module(monkeypatch):
def increase_usage_by_id(model_id, used_tokens):
return True
@staticmethod
def model_instance(_model_config):
return _DummyLLMBundle()
class _TenantService:
@staticmethod
def get_by_id(tenant_id):
@@ -455,6 +501,19 @@ def _load_chunk_module(monkeypatch):
return module
def _load_chunk_api_module(monkeypatch):
_load_chunk_module(monkeypatch)
repo_root = Path(__file__).resolve().parents[4]
module_name = "test_chunk_api_routes_unit_module"
module_path = repo_root / "api" / "apps" / "restful_apis" / "chunk_api.py"
spec = importlib.util.spec_from_file_location(module_name, module_path)
module = importlib.util.module_from_spec(spec)
module.manager = _DummyManager()
monkeypatch.setitem(sys.modules, module_name, module)
spec.loader.exec_module(module)
return module
def _set_request_json(monkeypatch, module, payload):
monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue(payload))
@@ -465,347 +524,133 @@ def set_tenant_info():
@pytest.mark.p2
def test_list_chunk_exception_branches_unit(monkeypatch):
module = _load_chunk_module(monkeypatch)
def test_restful_chunk_list_get_and_delete_unit(monkeypatch):
module = _load_chunk_api_module(monkeypatch)
module.request = SimpleNamespace(args={"keywords": "chunk", "available": "true"}, headers={})
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "keywords": "chunk", "available_int": 0})
res = _run(module.list_chunk())
res = _run(_route_core(module.list_chunks)("tenant-1", "kb-1", "doc-1"))
assert res["code"] == 0, res
assert res["data"]["total"] == 1, res
assert res["data"]["chunks"][0]["available_int"] == 1, res
assert res["data"]["chunks"][0]["id"] == "chunk-1", res
assert res["data"]["chunks"][0]["available"] is True, res
monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "")
_set_request_json(monkeypatch, module, {"doc_id": "doc-1"})
res = _run(module.list_chunk())
assert res["code"] == module.RetCode.DATA_ERROR, res
assert res["message"] == "Tenant not found!", res
monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "tenant-1")
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (False, None))
_set_request_json(monkeypatch, module, {"doc_id": "doc-1"})
res = _run(module.list_chunk())
assert res["message"] == "Document not found!", res
async def _raise_not_found(*_args, **_kwargs):
raise Exception("x not_found y")
monkeypatch.setattr(module.settings.retriever, "search", _raise_not_found)
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, _DummyDoc()))
_set_request_json(monkeypatch, module, {"doc_id": "doc-1"})
res = _run(module.list_chunk())
assert res["code"] == module.RetCode.DATA_ERROR, res
assert res["message"] == "No chunk found!", res
async def _raise_generic(*_args, **_kwargs):
raise RuntimeError("boom")
monkeypatch.setattr(module.settings.retriever, "search", _raise_generic)
_set_request_json(monkeypatch, module, {"doc_id": "doc-1"})
res = _run(module.list_chunk())
assert res["code"] == module.RetCode.EXCEPTION_ERROR, res
assert "boom" in res["message"], res
@pytest.mark.p2
def test_get_chunk_sanitize_and_exception_matrix_unit(monkeypatch):
module = _load_chunk_module(monkeypatch)
module.request = SimpleNamespace(args={"chunk_id": "chunk-1"}, headers={})
res = module.get()
res = _run(_route_core(module.get_chunk)("tenant-1", "kb-1", "doc-1", "chunk-1"))
assert res["code"] == 0, res
assert "q_2_vec" not in res["data"], res
assert "content_tks" not in res["data"], res
assert "content_ltks" not in res["data"], res
assert "content_sm_ltks" not in res["data"], res
monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [])
res = module.get()
assert res["message"] == "Tenant not found!", res
monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [_DummyTenant("tenant-1")])
module.settings.docStoreConn.chunk = None
res = module.get()
assert res["code"] == module.RetCode.EXCEPTION_ERROR, res
assert "Chunk not found" in res["message"], res
def _raise_not_found(*_args, **_kwargs):
raise Exception("NotFoundError: chunk-1")
monkeypatch.setattr(module.settings.docStoreConn, "get", _raise_not_found)
res = module.get()
assert res["code"] == module.RetCode.DATA_ERROR, res
assert res["message"] == "Chunk not found!", res
def _raise_generic(*_args, **_kwargs):
raise RuntimeError("get boom")
monkeypatch.setattr(module.settings.docStoreConn, "get", _raise_generic)
res = module.get()
assert res["code"] == module.RetCode.EXCEPTION_ERROR, res
assert "get boom" in res["message"], res
monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"chunk_ids": ["chunk-1"]}))
res = _run(_route_core(module.rm_chunk)("tenant-1", "kb-1", "doc-1"))
assert res["code"] == 0, res
assert module.settings.docStoreConn.deleted_inputs[-1]["doc_id"] == "doc-1"
@pytest.mark.p2
def test_set_chunk_bytes_qa_image_and_guard_matrix_unit(monkeypatch):
module = _load_chunk_module(monkeypatch)
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_id": "chunk-1", "content_with_weight": 1})
with pytest.raises(TypeError, match="expected string or bytes-like object"):
_run(module.set())
_set_request_json(
monkeypatch,
module,
{"doc_id": "doc-1", "chunk_id": "chunk-1", "content_with_weight": "abc", "important_kwd": "bad"},
)
res = _run(module.set())
assert res["message"] == "`important_kwd` should be a list", res
_set_request_json(
monkeypatch,
module,
{"doc_id": "doc-1", "chunk_id": "chunk-1", "content_with_weight": "abc", "question_kwd": "bad"},
)
res = _run(module.set())
assert res["message"] == "`question_kwd` should be a list", res
monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "")
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_id": "chunk-1", "content_with_weight": "abc"})
res = _run(module.set())
assert res["message"] == "Tenant not found!", res
monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "tenant-1")
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (False, None))
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_id": "chunk-1", "content_with_weight": "abc"})
res = _run(module.set())
assert res["message"] == "Document not found!", res
def test_restful_chunk_add_update_and_switch_unit(monkeypatch):
module = _load_chunk_api_module(monkeypatch)
module.request = SimpleNamespace(args={}, headers={})
monkeypatch.setattr(
module.DocumentService,
"get_by_id",
lambda _doc_id: (True, _DummyDoc(doc_id="doc-1", parser_id=module.ParserType.NAIVE)),
)
_set_request_json(
monkeypatch,
module,
{"doc_id": "doc-1", "chunk_id": "chunk-1", "content_with_weight": "abc", "tag_feas": [0.1]},
"get_request_json",
lambda: _AwaitableValue(
{
"content": "chunk",
"important_keywords": ["i1"],
"questions": ["q1"],
"tag_kwd": ["tag"],
"tag_feas": {"tag": 0.2},
}
),
)
res = _run(module.set())
assert "`tag_feas` must be an object mapping string tags to finite numeric scores" in res["message"], res
_set_request_json(
monkeypatch,
module,
{
"doc_id": "doc-1",
"chunk_id": "chunk-1",
"content_with_weight": b"bytes-content",
"important_kwd": ["important"],
"question_kwd": ["question"],
"tag_kwd": ["tag"],
"tag_feas": {"tag": 0.1},
"available_int": 0,
},
)
res = _run(module.set())
res = _run(_route_core(module.add_chunk)("tenant-1", "kb-1", "doc-1"))
assert res["code"] == 0, res
assert module.settings.docStoreConn.updated[-1][1]["content_with_weight"] == "bytes-content"
assert res["data"]["chunk"]["content"] == "chunk", res
assert module.settings.docStoreConn.inserted, "insert should be called"
assert module.DocumentService.increment_calls, "increment_chunk_num should be called"
monkeypatch.setattr(
module.DocumentService,
"get_by_id",
lambda _doc_id: (True, _DummyDoc(doc_id="doc-1", parser_id=module.ParserType.QA)),
)
_set_request_json(
monkeypatch,
module,
{
"doc_id": "doc-1",
"chunk_id": "chunk-2",
"content_with_weight": "Q:Question\nA:Answer",
"image_base64": base64.b64encode(b"image").decode("utf-8"),
"img_id": "bucket-name",
},
"get_request_json",
lambda: _AwaitableValue(
{
"content": "updated chunk",
"important_keywords": ["i2"],
"questions": ["q2"],
"tag_kwd": ["tag2"],
"positions": [[1, 2, 3, 4, 5]],
"available": False,
}
),
)
res = _run(module.set())
res = _run(_route_core(module.update_chunk)("tenant-1", "kb-1", "doc-1", "chunk-1"))
assert res["code"] == 0, res
assert module.settings.STORAGE_IMPL.put_calls, "image storage branch should be called"
updated = module.settings.docStoreConn.updated[-1][1]
assert updated["content_with_weight"] == "updated chunk"
assert updated["available_int"] == 0
assert updated["position_int"] == [[1, 2, 3, 4, 5]]
async def _raise_thread_pool(_func):
raise RuntimeError("set tp boom")
monkeypatch.setattr(module, "thread_pool_exec", _raise_thread_pool)
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_id": "chunk-1", "content_with_weight": "abc"})
res = _run(module.set())
assert res["code"] == module.RetCode.EXCEPTION_ERROR, res
assert "set tp boom" in res["message"], res
@pytest.mark.p2
def test_switch_chunk_success_failure_and_exception_unit(monkeypatch):
module = _load_chunk_module(monkeypatch)
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (False, None))
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1"], "available_int": 1})
res = _run(module.switch())
assert res["message"] == "Document not found!", res
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, _DummyDoc()))
monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "tenant-1")
monkeypatch.setattr(module.settings.docStoreConn, "update", lambda *_args, **_kwargs: False)
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1", "c2"], "available_int": 0})
res = _run(module.switch())
assert res["message"] == "Index updating failure", res
monkeypatch.setattr(module.settings.docStoreConn, "update", lambda *_args, **_kwargs: True)
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1", "c2"], "available_int": 1})
res = _run(module.switch())
monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"chunk_ids": ["chunk-1"], "available": True}))
res = _run(_route_core(module.switch_chunks)("tenant-1", "kb-1", "doc-1"))
assert res["code"] == 0, res
assert res["data"] is True, res
async def _raise_thread_pool(_func):
raise RuntimeError("switch tp boom")
monkeypatch.setattr(module, "thread_pool_exec", _raise_thread_pool)
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1"], "available_int": 1})
res = _run(module.switch())
assert res["code"] == module.RetCode.EXCEPTION_ERROR, res
assert "switch tp boom" in res["message"], res
@pytest.mark.p2
def test_rm_chunk_delete_exception_partial_compensation_and_cleanup_unit(monkeypatch):
module = _load_chunk_module(monkeypatch)
def test_restful_chunk_guard_branches_unit(monkeypatch):
module = _load_chunk_api_module(monkeypatch)
module.request = SimpleNamespace(args={}, headers={})
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (False, None))
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1"]})
res = _run(module.rm())
assert res["message"] == "Document not found!", res
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False)
res = _run(_route_core(module.list_chunks)("tenant-1", "kb-1", "doc-1"))
assert res["message"] == "You don't own the dataset kb-1.", res
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": []})
monkeypatch.setattr(
module.DocumentService,
"get_by_id",
lambda _doc_id: (_ for _ in ()).throw(AssertionError("get_by_id must not run for empty delete payload")),
)
monkeypatch.setattr(
module.settings.docStoreConn,
"delete",
lambda *_args, **_kwargs: (_ for _ in ()).throw(AssertionError("delete must not run for empty delete payload")),
)
res = _run(module.rm())
assert res["code"] == 0, res
monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True)
monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [])
res = _run(_route_core(module.list_chunks)("tenant-1", "kb-1", "doc-1"))
assert res["message"] == "You don't own the document doc-1.", res
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, _DummyDoc()))
monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc()])
module.request = SimpleNamespace(args={"id": "chunk-1"}, headers={})
module.settings.docStoreConn.chunk = None
res = _run(_route_core(module.list_chunks)("tenant-1", "kb-1", "doc-1"))
assert res["code"] == module.RetCode.DATA_ERROR, res
assert "Chunk not found" in res["message"], res
def _raise_delete(*_args, **_kwargs):
raise RuntimeError("delete boom")
module.settings.docStoreConn.chunk = {
"id": "chunk-1",
"doc_id": "other-doc",
"content_with_weight": "chunk",
"docnm_kwd": "Doc",
}
res = _run(_route_core(module.list_chunks)("tenant-1", "kb-1", "doc-1"))
assert res["code"] == module.RetCode.DATA_ERROR, res
assert "Chunk not found" in res["message"], res
monkeypatch.setattr(module.settings.docStoreConn, "delete", _raise_delete)
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1"]})
res = _run(module.rm())
assert res["message"] == "Chunk deleting failure", res
module.settings.docStoreConn.chunk = None
module.request = SimpleNamespace(args={}, headers={})
res = _run(_route_core(module.get_chunk)("tenant-1", "kb-1", "doc-1", "chunk-1"))
assert res["code"] == module.RetCode.DATA_ERROR, res
assert "Chunk not found" in res["message"], res
def _delete(condition, *_args, **_kwargs):
module.settings.docStoreConn.deleted_inputs.append(condition)
if not module.settings.docStoreConn.to_delete:
return 0
return module.settings.docStoreConn.to_delete.pop(0)
monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"content": ""}))
res = _run(_route_core(module.add_chunk)("tenant-1", "kb-1", "doc-1"))
assert res["message"] == "`content` is required", res
module.settings.docStoreConn.to_delete = [0]
monkeypatch.setattr(module.settings.docStoreConn, "delete", _delete)
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1"]})
res = _run(module.rm())
assert res["message"] == "Index updating failure", res
module.settings.docStoreConn.chunk = {"id": "chunk-1", "doc_id": "doc-1", "content_with_weight": "chunk"}
monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"important_keywords": "bad"}))
res = _run(_route_core(module.update_chunk)("tenant-1", "kb-1", "doc-1", "chunk-1"))
assert res["message"] == "`important_keywords` should be a list", res
module.settings.docStoreConn.to_delete = [1, 2]
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1", "c2", "c3"]})
res = _run(module.rm())
assert res["code"] == 0, res
assert module.DocumentService.decrement_calls, "decrement_chunk_num should be called"
assert len(module.settings.STORAGE_IMPL.rm_calls) >= 1
monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"chunk_ids": []}))
res = _run(_route_core(module.switch_chunks)("tenant-1", "kb-1", "doc-1"))
assert res["message"] == "`chunk_ids` is required.", res
module.settings.docStoreConn.to_delete = [1]
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": "c1"})
res = _run(module.rm())
assert res["code"] == 0, res
async def _raise_thread_pool(_func):
raise RuntimeError("rm tp boom")
monkeypatch.setattr(module, "thread_pool_exec", _raise_thread_pool)
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "chunk_ids": ["c1"]})
res = _run(module.rm())
assert res["code"] == module.RetCode.EXCEPTION_ERROR, res
assert "rm tp boom" in res["message"], res
@pytest.mark.p2
def test_create_chunk_guards_pagerank_and_success_unit(monkeypatch):
module = _load_chunk_module(monkeypatch)
module.request = SimpleNamespace(headers={"X-Request-ID": "req-1"}, args={})
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "content_with_weight": "chunk", "important_kwd": "bad"})
res = _run(module.create())
assert res["message"] == "`important_kwd` is required to be a list", res
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "content_with_weight": "chunk", "question_kwd": "bad"})
res = _run(module.create())
assert res["message"] == "`question_kwd` is required to be a list", res
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (False, None))
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "content_with_weight": "chunk"})
res = _run(module.create())
assert res["message"] == "Document not found!", res
monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _doc_id: (True, _DummyDoc(doc_id="doc-1")))
monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "")
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "content_with_weight": "chunk"})
res = _run(module.create())
assert res["message"] == "Tenant not found!", res
monkeypatch.setattr(module.DocumentService, "get_tenant_id", lambda _doc_id: "tenant-1")
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (False, None))
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "content_with_weight": "chunk"})
res = _run(module.create())
assert res["message"] == "Knowledgebase not found!", res
monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _kb_id: (True, SimpleNamespace(pagerank=0.8)))
_set_request_json(
monkeypatch,
module,
{"doc_id": "doc-1", "content_with_weight": "chunk", "tag_feas": [0.2]},
)
res = _run(module.create())
assert "`tag_feas` must be an object mapping string tags to finite numeric scores" in res["message"], res
_set_request_json(
monkeypatch,
module,
{
"doc_id": "doc-1",
"content_with_weight": "chunk",
"important_kwd": ["i1"],
"question_kwd": ["q1"],
"tag_feas": {"tag": 0.2},
},
)
res = _run(module.create())
assert res["code"] == 0, res
assert res["data"]["chunk_id"], res
assert module.settings.docStoreConn.inserted, "insert should be called"
inserted = module.settings.docStoreConn.inserted[-1]
assert "pagerank_flt" in inserted
assert module.DocumentService.increment_calls, "increment_chunk_num should be called"
async def _raise_thread_pool(_func):
raise RuntimeError("create tp boom")
monkeypatch.setattr(module, "thread_pool_exec", _raise_thread_pool)
_set_request_json(monkeypatch, module, {"doc_id": "doc-1", "content_with_weight": "chunk"})
res = _run(module.create())
assert res["code"] == module.RetCode.EXCEPTION_ERROR, res
assert "create tp boom" in res["message"], res
monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"chunk_ids": ["chunk-1"]}))
res = _run(_route_core(module.switch_chunks)("tenant-1", "kb-1", "doc-1"))
assert res["message"] == "`available_int` or `available` is required.", res
@pytest.mark.p2

230
test/testcases/test_web_api/test_chunk_app/test_create_chunk.py
View File

@@ -16,24 +16,28 @@
from concurrent.futures import ThreadPoolExecutor, as_completed
import pytest
from test_common import add_chunk, delete_document, get_chunk, list_chunks
from configs import INVALID_API_TOKEN
from libs.auth import RAGFlowWebApiAuth
from test_common import add_chunk, delete_document, get_chunk, list_chunks
def validate_chunk_details(auth, kb_id, doc_id, payload, res):
chunk_id = res["data"]["chunk_id"]
res = get_chunk(auth, {"chunk_id": chunk_id})
assert res["code"] == 0, res
chunk = res["data"]
assert chunk["doc_id"] == doc_id
assert chunk["kb_id"] == kb_id
assert chunk["content_with_weight"] == payload["content_with_weight"]
if "important_kwd" in payload:
assert chunk["important_kwd"] == payload["important_kwd"]
if "question_kwd" in payload:
expected = [str(q).strip() for q in payload.get("question_kwd", [])]
assert chunk["question_kwd"] == expected
def validate_chunk_details(auth, dataset_id, document_id, payload, res):
chunk = res["data"]["chunk"]
assert chunk["dataset_id"] == dataset_id
assert chunk["document_id"] == document_id
assert chunk["content"] == payload["content"]
if "important_keywords" in payload:
assert chunk["important_keywords"] == payload["important_keywords"]
if "questions" in payload:
expected = [str(q).strip() for q in payload.get("questions", []) if str(q).strip()]
assert chunk["questions"] == expected
if "tag_kwd" in payload:
assert chunk["tag_kwd"] == payload["tag_kwd"]
fetched = get_chunk(auth, dataset_id, document_id, chunk["id"])
assert fetched["code"] == 0, fetched
assert fetched["data"]["id"] == chunk["id"]
assert fetched["data"]["doc_id"] == document_id
@pytest.mark.p2
@@ -46,7 +50,7 @@ class TestAuthorization:
],
)
def test_invalid_auth(self, invalid_auth, expected_code, expected_message):
res = add_chunk(invalid_auth)
res = add_chunk(invalid_auth, "dataset_id", "document_id", {"content": "chunk test"})
assert res["code"] == expected_code, res
assert res["message"] == expected_message, res
@@ -56,33 +60,22 @@ class TestAddChunk:
@pytest.mark.parametrize(
"payload, expected_code, expected_message",
[
({"content_with_weight": None}, 100, """TypeError("unsupported operand type(s) for +: 'NoneType' and 'str'")"""),
({"content_with_weight": ""}, 100, """Exception('Error: 413 - {"error":"Input validation error: `inputs` cannot be empty","error_type":"Validation"}')"""),
pytest.param(
{"content_with_weight": 1},
100,
"""TypeError("unsupported operand type(s) for +: 'int' and 'str'")""",
marks=pytest.mark.skip,
),
({"content_with_weight": "a"}, 0, ""),
({"content_with_weight": " "}, 0, ""),
({"content_with_weight": "\n!?。；！？\"'"}, 0, ""),
({"content": None}, 102, "`content` is required"),
({"content": ""}, 102, "`content` is required"),
({"content": "a"}, 0, ""),
({"content": " "}, 102, "`content` is required"),
({"content": "\n!?。；！？\"'"}, 0, ""),
],
)
def test_content(self, WebApiAuth, add_document, payload, expected_code, expected_message):
kb_id, doc_id = add_document
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
if res["code"] == 0:
chunks_count = res["data"]["doc"]["chunk_num"]
else:
chunks_count = 0
res = add_chunk(WebApiAuth, {**payload, "doc_id": doc_id})
dataset_id, document_id = add_document
chunks_count = list_chunks(WebApiAuth, dataset_id, document_id)["data"]["doc"]["chunk_count"]
res = add_chunk(WebApiAuth, dataset_id, document_id, payload)
assert res["code"] == expected_code, res
if expected_code == 0:
validate_chunk_details(WebApiAuth, kb_id, doc_id, payload, res)
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
assert res["code"] == 0, res
assert res["data"]["doc"]["chunk_num"] == chunks_count + 1, res
validate_chunk_details(WebApiAuth, dataset_id, document_id, payload, res)
res = list_chunks(WebApiAuth, dataset_id, document_id)
assert res["data"]["doc"]["chunk_count"] == chunks_count + 1, res
else:
assert res["message"] == expected_message, res
@@ -90,32 +83,20 @@ class TestAddChunk:
@pytest.mark.parametrize(
"payload, expected_code, expected_message",
[
({"content_with_weight": "chunk test", "important_kwd": ["a", "b", "c"]}, 0, ""),
({"content_with_weight": "chunk test", "important_kwd": [""]}, 0, ""),
(
{"content_with_weight": "chunk test", "important_kwd": [1]},
100,
"TypeError('sequence item 0: expected str instance, int found')",
),
({"content_with_weight": "chunk test", "important_kwd": ["a", "a"]}, 0, ""),
({"content_with_weight": "chunk test", "important_kwd": "abc"}, 102, "`important_kwd` is required to be a list"),
({"content_with_weight": "chunk test", "important_kwd": 123}, 102, "`important_kwd` is required to be a list"),
({"content": "chunk test", "important_keywords": ["a", "b", "c"]}, 0, ""),
({"content": "chunk test", "important_keywords": [""]}, 0, ""),
({"content": "chunk test", "important_keywords": [1]}, 100, "TypeError('sequence item 0: expected str instance, int found')"),
({"content": "chunk test", "important_keywords": ["a", "a"]}, 0, ""),
({"content": "chunk test", "important_keywords": "abc"}, 102, "`important_keywords` is required to be a list"),
({"content": "chunk test", "important_keywords": 123}, 102, "`important_keywords` is required to be a list"),
],
)
def test_important_keywords(self, WebApiAuth, add_document, payload, expected_code, expected_message):
kb_id, doc_id = add_document
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
if res["code"] == 0:
chunks_count = res["data"]["doc"]["chunk_num"]
else:
chunks_count = 0
res = add_chunk(WebApiAuth, {**payload, "doc_id": doc_id})
dataset_id, document_id = add_document
res = add_chunk(WebApiAuth, dataset_id, document_id, payload)
assert res["code"] == expected_code, res
if expected_code == 0:
validate_chunk_details(WebApiAuth, kb_id, doc_id, payload, res)
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
assert res["code"] == 0, res
assert res["data"]["doc"]["chunk_num"] == chunks_count + 1, res
validate_chunk_details(WebApiAuth, dataset_id, document_id, payload, res)
else:
assert res["message"] == expected_message, res
@@ -123,130 +104,95 @@ class TestAddChunk:
@pytest.mark.parametrize(
"payload, expected_code, expected_message",
[
({"content_with_weight": "chunk test", "question_kwd": ["a", "b", "c"]}, 0, ""),
({"content_with_weight": "chunk test", "question_kwd": [""]}, 100, """Exception('Error: 413 - {"error":"Input validation error: `inputs` cannot be empty","error_type":"Validation"}')"""),
({"content_with_weight": "chunk test", "question_kwd": [1]}, 100, "TypeError('sequence item 0: expected str instance, int found')"),
({"content_with_weight": "chunk test", "question_kwd": ["a", "a"]}, 0, ""),
({"content_with_weight": "chunk test", "question_kwd": "abc"}, 102, "`question_kwd` is required to be a list"),
({"content_with_weight": "chunk test", "question_kwd": 123}, 102, "`question_kwd` is required to be a list"),
({"content": "chunk test", "questions": ["a", "b", "c"]}, 0, ""),
({"content": "chunk test", "questions": [""]}, 0, ""),
({"content": "chunk test", "questions": [1]}, 100, "TypeError('sequence item 0: expected str instance, int found')"),
({"content": "chunk test", "questions": ["a", "a"]}, 0, ""),
({"content": "chunk test", "questions": "abc"}, 102, "`questions` is required to be a list"),
({"content": "chunk test", "questions": 123}, 102, "`questions` is required to be a list"),
],
)
def test_questions(self, WebApiAuth, add_document, payload, expected_code, expected_message):
kb_id, doc_id = add_document
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
if res["code"] == 0:
chunks_count = res["data"]["doc"]["chunk_num"]
else:
chunks_count = 0
res = add_chunk(WebApiAuth, {**payload, "doc_id": doc_id})
dataset_id, document_id = add_document
res = add_chunk(WebApiAuth, dataset_id, document_id, payload)
assert res["code"] == expected_code, res
if expected_code == 0:
validate_chunk_details(WebApiAuth, kb_id, doc_id, payload, res)
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
assert res["code"] == 0, res
assert res["data"]["doc"]["chunk_num"] == chunks_count + 1, res
validate_chunk_details(WebApiAuth, dataset_id, document_id, payload, res)
else:
assert res["message"] == expected_message, res
@pytest.mark.p2
def test_get_chunk_not_found(self, WebApiAuth):
res = get_chunk(WebApiAuth, {"chunk_id": "missing_chunk_id"})
assert res["code"] != 0, res
assert "Chunk not found" in res["message"], res
def test_add_chunk_with_tag_fields(self, WebApiAuth, add_document):
dataset_id, document_id = add_document
payload = {
"content": "chunk with tags",
"tag_kwd": ["tag1", "tag2"],
"important_keywords": ["tag"],
"questions": ["question"],
}
res = add_chunk(WebApiAuth, dataset_id, document_id, payload)
assert res["code"] == 0, res
validate_chunk_details(WebApiAuth, dataset_id, document_id, payload, res)
@pytest.mark.p2
def test_create_chunk_with_tag_fields(self, WebApiAuth, add_document):
_, doc_id = add_document
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
if res["code"] == 0:
chunks_count = res["data"]["doc"]["chunk_num"]
else:
chunks_count = 0
payload = {
"doc_id": doc_id,
"content_with_weight": "chunk with tags",
"tag_feas": {"tag1": 0.1, "tag2": 0.2},
"important_kwd": ["tag"],
"question_kwd": ["question"],
}
res = add_chunk(WebApiAuth, payload)
assert res["code"] == 0, res
assert res["data"]["chunk_id"], res
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
assert res["code"] == 0, res
assert res["data"]["doc"]["chunk_num"] == chunks_count + 1, res
def test_get_chunk_not_found(self, WebApiAuth, add_document):
dataset_id, document_id = add_document
res = get_chunk(WebApiAuth, dataset_id, document_id, "missing_chunk_id")
assert res["code"] == 102, res
assert "Chunk not found" in res["message"], res
@pytest.mark.p3
@pytest.mark.parametrize(
"doc_id, expected_code, expected_message",
"document_id, expected_code, expected_message",
[
("", 102, "Document not found!"),
("invalid_document_id", 102, "Document not found!"),
("invalid_document_id", 102, "You don't own the document invalid_document_id."),
],
)
def test_invalid_document_id(self, WebApiAuth, add_document, doc_id, expected_code, expected_message):
_, _ = add_document
res = add_chunk(WebApiAuth, {"doc_id": doc_id, "content_with_weight": "chunk test"})
def test_invalid_document_id(self, WebApiAuth, add_document, document_id, expected_code, expected_message):
dataset_id, _ = add_document
res = add_chunk(WebApiAuth, dataset_id, document_id, {"content": "chunk test"})
assert res["code"] == expected_code, res
assert res["message"] == expected_message, res
@pytest.mark.p3
def test_repeated_add_chunk(self, WebApiAuth, add_document):
payload = {"content_with_weight": "chunk test"}
kb_id, doc_id = add_document
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
if res["code"] != 0:
assert False, res
chunks_count = res["data"]["doc"]["chunk_num"]
payload = {"content": "chunk test"}
dataset_id, document_id = add_document
chunks_count = list_chunks(WebApiAuth, dataset_id, document_id)["data"]["doc"]["chunk_count"]
res = add_chunk(WebApiAuth, {**payload, "doc_id": doc_id})
res = add_chunk(WebApiAuth, dataset_id, document_id, payload)
assert res["code"] == 0, res
validate_chunk_details(WebApiAuth, kb_id, doc_id, payload, res)
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
if res["code"] != 0:
assert False, res
assert res["data"]["doc"]["chunk_num"] == chunks_count + 1, res
validate_chunk_details(WebApiAuth, dataset_id, document_id, payload, res)
res = add_chunk(WebApiAuth, {**payload, "doc_id": doc_id})
res = add_chunk(WebApiAuth, dataset_id, document_id, payload)
assert res["code"] == 0, res
validate_chunk_details(WebApiAuth, kb_id, doc_id, payload, res)
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
if res["code"] != 0:
assert False, res
assert res["data"]["doc"]["chunk_num"] == chunks_count + 2, res
validate_chunk_details(WebApiAuth, dataset_id, document_id, payload, res)
res = list_chunks(WebApiAuth, dataset_id, document_id)
assert res["data"]["doc"]["chunk_count"] == chunks_count + 2, res
@pytest.mark.p2
def test_add_chunk_to_deleted_document(self, WebApiAuth, add_document):
kb_id, doc_id = add_document
delete_document(WebApiAuth, kb_id, {"ids": [doc_id]})
res = add_chunk(WebApiAuth, {"doc_id": doc_id, "content_with_weight": "chunk test"})
dataset_id, document_id = add_document
delete_document(WebApiAuth, dataset_id, {"ids": [document_id]})
res = add_chunk(WebApiAuth, dataset_id, document_id, {"content": "chunk test"})
assert res["code"] == 102, res
assert res["message"] == "Document not found!", res
assert res["message"] == f"You don't own the document {document_id}.", res
@pytest.mark.skip(reason="issues/6411")
@pytest.mark.p3
def test_concurrent_add_chunk(self, WebApiAuth, add_document):
count = 50
_, doc_id = add_document
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
if res["code"] == 0:
chunks_count = res["data"]["doc"]["chunk_num"]
else:
chunks_count = 0
dataset_id, document_id = add_document
chunks_count = list_chunks(WebApiAuth, dataset_id, document_id)["data"]["doc"]["chunk_count"]
with ThreadPoolExecutor(max_workers=5) as executor:
futures = [
executor.submit(
add_chunk,
WebApiAuth,
{"doc_id": doc_id, "content_with_weight": f"chunk test {i}"},
)
executor.submit(add_chunk, WebApiAuth, dataset_id, document_id, {"content": f"chunk test {i}"})
for i in range(count)
]
responses = list(as_completed(futures))
assert len(responses) == count, responses
assert all(future.result()["code"] == 0 for future in futures)
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
assert res["code"] == 0, res
assert res["data"]["doc"]["chunk_num"] == chunks_count + count
res = list_chunks(WebApiAuth, dataset_id, document_id)
assert res["data"]["doc"]["chunk_count"] == chunks_count + count

96
test/testcases/test_web_api/test_chunk_app/test_list_chunks.py
View File

@@ -17,9 +17,9 @@ import os
from concurrent.futures import ThreadPoolExecutor, as_completed
import pytest
from test_common import batch_add_chunks, list_chunks, update_chunk
from configs import INVALID_API_TOKEN
from libs.auth import RAGFlowWebApiAuth
from test_common import batch_add_chunks, list_chunks, update_chunk
@pytest.mark.p2
@@ -32,7 +32,7 @@ class TestAuthorization:
],
)
def test_invalid_auth(self, invalid_auth, expected_code, expected_message):
res = list_chunks(invalid_auth, {"doc_id": "document_id"})
res = list_chunks(invalid_auth, "dataset_id", "document_id")
assert res["code"] == expected_code, res
assert res["message"] == expected_message, res
@@ -42,21 +42,18 @@ class TestChunksList:
@pytest.mark.parametrize(
"params, expected_code, expected_page_size, expected_message",
[
pytest.param({"page": None, "size": 2}, 100, 0, """TypeError("int() argument must be a string, a bytes-like object or a real number, not 'NoneType'")""", marks=pytest.mark.skip),
pytest.param({"page": 0, "size": 2}, 100, 0, "ValueError('Search does not support negative slicing.')", marks=pytest.mark.skip),
({"page": 2, "size": 2}, 0, 2, ""),
({"page": 3, "size": 2}, 0, 1, ""),
({"page": "3", "size": 2}, 0, 1, ""),
pytest.param({"page": -1, "size": 2}, 100, 0, "ValueError('Search does not support negative slicing.')", marks=pytest.mark.skip),
pytest.param({"page": "a", "size": 2}, 100, 0, """ValueError("invalid literal for int() with base 10: \'a\'")""", marks=pytest.mark.skip),
({"page": None, "page_size": 2}, 0, 2, ""),
pytest.param({"page": 0, "page_size": 2}, 100, 0, "ValueError('Search does not support negative slicing.')", marks=pytest.mark.skip),
({"page": 2, "page_size": 2}, 0, 2, ""),
({"page": 3, "page_size": 2}, 0, 1, ""),
({"page": "3", "page_size": 2}, 0, 1, ""),
pytest.param({"page": -1, "page_size": 2}, 100, 0, "ValueError('Search does not support negative slicing.')", marks=pytest.mark.skip),
pytest.param({"page": "a", "page_size": 2}, 100, 0, """ValueError("invalid literal for int() with base 10: 'a'")""", marks=pytest.mark.skip),
],
)
def test_page(self, WebApiAuth, add_chunks, params, expected_code, expected_page_size, expected_message):
_, doc_id, _ = add_chunks
payload = {"doc_id": doc_id}
if params:
payload.update(params)
res = list_chunks(WebApiAuth, payload)
dataset_id, document_id, _ = add_chunks
res = list_chunks(WebApiAuth, dataset_id, document_id, params=params)
assert res["code"] == expected_code, res
if expected_code == 0:
assert len(res["data"]["chunks"]) == expected_page_size, res
@@ -67,21 +64,18 @@ class TestChunksList:
@pytest.mark.parametrize(
"params, expected_code, expected_page_size, expected_message",
[
({"size": None}, 100, 0, """TypeError("int() argument must be a string, a bytes-like object or a real number, not 'NoneType'")"""),
pytest.param({"size": 0}, 0, 5, ""),
({"size": 1}, 0, 1, ""),
({"size": 6}, 0, 5, ""),
({"size": "1"}, 0, 1, ""),
pytest.param({"size": -1}, 0, 5, "", marks=pytest.mark.skip),
pytest.param({"size": "a"}, 100, 0, """ValueError("invalid literal for int() with base 10: \'a\'")""", marks=pytest.mark.skip),
({"page_size": None}, 0, 5, ""),
pytest.param({"page_size": 0}, 0, 5, ""),
({"page_size": 1}, 0, 1, ""),
({"page_size": 6}, 0, 5, ""),
({"page_size": "1"}, 0, 1, ""),
pytest.param({"page_size": -1}, 0, 5, "", marks=pytest.mark.skip),
pytest.param({"page_size": "a"}, 100, 0, """ValueError("invalid literal for int() with base 10: 'a'")""", marks=pytest.mark.skip),
],
)
def test_page_size(self, WebApiAuth, add_chunks, params, expected_code, expected_page_size, expected_message):
_, doc_id, _ = add_chunks
payload = {"doc_id": doc_id}
if params:
payload.update(params)
res = list_chunks(WebApiAuth, payload)
dataset_id, document_id, _ = add_chunks
res = list_chunks(WebApiAuth, dataset_id, document_id, params=params)
assert res["code"] == expected_code, res
if expected_code == 0:
assert len(res["data"]["chunks"]) == expected_page_size, res
@@ -89,29 +83,22 @@ class TestChunksList:
assert res["message"] == expected_message, res
@pytest.mark.p2
def test_available_int_filter(self, WebApiAuth, add_chunks):
_, doc_id, chunk_ids = add_chunks
def test_available_filter(self, WebApiAuth, add_chunks):
dataset_id, document_id, chunk_ids = add_chunks
chunk_id = chunk_ids[0]
res = update_chunk(
WebApiAuth,
{"doc_id": doc_id, "chunk_id": chunk_id, "content_with_weight": "unchanged content", "available_int": 0},
)
res = update_chunk(WebApiAuth, dataset_id, document_id, chunk_id, {"content": "unchanged content", "available": False})
assert res["code"] == 0, res
from time import sleep
sleep(1)
res = list_chunks(WebApiAuth, {"doc_id": doc_id, "available_int": 0})
res = list_chunks(WebApiAuth, dataset_id, document_id, params={"available": "false"})
assert res["code"] == 0, res
assert len(res["data"]["chunks"]) >= 1, res
assert all(chunk["available_int"] == 0 for chunk in res["data"]["chunks"]), res
assert all(chunk["available"] is False for chunk in res["data"]["chunks"]), res
# Restore the class-scoped fixture state for subsequent keyword cases.
res = update_chunk(
WebApiAuth,
{"doc_id": doc_id, "chunk_id": chunk_id, "content_with_weight": "chunk test 0", "available_int": 1},
)
res = update_chunk(WebApiAuth, dataset_id, document_id, chunk_id, {"content": "chunk test 0", "available": True})
assert res["code"] == 0, res
sleep(1)
@@ -123,49 +110,44 @@ class TestChunksList:
({"keywords": ""}, 5),
({"keywords": "1"}, 1),
pytest.param({"keywords": "chunk"}, 4, marks=pytest.mark.skipif(os.getenv("DOC_ENGINE") == "infinity", reason="issues/6509")),
({"keywords": "content"}, 1),
({"keywords": "unknown"}, 0),
],
)
def test_keywords(self, WebApiAuth, add_chunks, params, expected_page_size):
_, doc_id, _ = add_chunks
payload = {"doc_id": doc_id}
if params:
payload.update(params)
res = list_chunks(WebApiAuth, payload)
dataset_id, document_id, _ = add_chunks
res = list_chunks(WebApiAuth, dataset_id, document_id, params=params)
assert res["code"] == 0, res
assert len(res["data"]["chunks"]) == expected_page_size, res
@pytest.mark.p3
def test_invalid_params(self, WebApiAuth, add_chunks):
_, doc_id, _ = add_chunks
payload = {"doc_id": doc_id, "a": "b"}
res = list_chunks(WebApiAuth, payload)
dataset_id, document_id, _ = add_chunks
res = list_chunks(WebApiAuth, dataset_id, document_id, params={"a": "b"})
assert res["code"] == 0, res
assert len(res["data"]["chunks"]) == 5, res
@pytest.mark.p3
def test_concurrent_list(self, WebApiAuth, add_chunks):
_, doc_id, _ = add_chunks
dataset_id, document_id, _ = add_chunks
count = 100
with ThreadPoolExecutor(max_workers=5) as executor:
futures = [executor.submit(list_chunks, WebApiAuth, {"doc_id": doc_id}) for i in range(count)]
futures = [executor.submit(list_chunks, WebApiAuth, dataset_id, document_id) for _ in range(count)]
responses = list(as_completed(futures))
assert len(responses) == count, responses
assert all(len(future.result()["data"]["chunks"]) == 5 for future in futures)
@pytest.mark.p1
def test_default(self, WebApiAuth, add_document):
_, doc_id = add_document
dataset_id, document_id = add_document
res = list_chunks(WebApiAuth, dataset_id, document_id)
chunks_count = res["data"]["doc"]["chunk_count"]
batch_add_chunks(WebApiAuth, dataset_id, document_id, 31)
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
chunks_count = res["data"]["doc"]["chunk_num"]
batch_add_chunks(WebApiAuth, doc_id, 31)
# issues/6487
from time import sleep
sleep(3)
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
res = list_chunks(WebApiAuth, dataset_id, document_id)
assert res["code"] == 0
assert len(res["data"]["chunks"]) == 30
assert res["data"]["doc"]["chunk_num"] == chunks_count + 31
assert res["data"]["doc"]["chunk_count"] == chunks_count + 31

104
test/testcases/test_web_api/test_chunk_app/test_rm_chunks.py
View File

@@ -16,9 +16,9 @@
from concurrent.futures import ThreadPoolExecutor, as_completed
import pytest
from test_common import batch_add_chunks, delete_chunks, list_chunks
from configs import INVALID_API_TOKEN
from libs.auth import RAGFlowWebApiAuth
from test_common import batch_add_chunks, delete_chunks, list_chunks
@pytest.mark.p2
@@ -31,7 +31,7 @@ class TestAuthorization:
],
)
def test_invalid_auth(self, invalid_auth, expected_code, expected_message):
res = delete_chunks(invalid_auth, {"doc_id": "document_id", "chunk_ids": ["1"]})
res = delete_chunks(invalid_auth, "dataset_id", "document_id", {"chunk_ids": ["1"]})
assert res["code"] == expected_code
assert res["message"] == expected_message
@@ -39,17 +39,16 @@ class TestAuthorization:
class TestChunksDeletion:
@pytest.mark.p3
@pytest.mark.parametrize(
"doc_id, expected_code, expected_message",
"document_id, expected_code, expected_message",
[
("", 102, "Document not found!"),
("invalid_document_id", 102, "Document not found!"),
("invalid_document_id", 100, "Can't find the document with ID invalid_document_id!"),
],
)
def test_invalid_document_id(self, WebApiAuth, add_chunks_func, doc_id, expected_code, expected_message):
_, _, chunk_ids = add_chunks_func
res = delete_chunks(WebApiAuth, {"doc_id": doc_id, "chunk_ids": chunk_ids})
def test_invalid_document_id(self, WebApiAuth, add_chunks_func, document_id, expected_code, expected_message):
dataset_id, _, chunk_ids = add_chunks_func
res = delete_chunks(WebApiAuth, dataset_id, document_id, {"chunk_ids": chunk_ids})
assert res["code"] == expected_code, res
assert res["message"] == expected_message, res
assert expected_message in res["message"], res
@pytest.mark.parametrize(
"payload",
@@ -60,61 +59,41 @@ class TestChunksDeletion:
],
)
def test_delete_partial_invalid_id(self, WebApiAuth, add_chunks_func, payload):
_, doc_id, chunk_ids = add_chunks_func
if callable(payload):
payload = payload(chunk_ids)
payload["doc_id"] = doc_id
res = delete_chunks(WebApiAuth, payload)
assert res["code"] == 0, res
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
assert res["code"] == 0, res
assert len(res["data"]["chunks"]) == 0, res
assert res["data"]["total"] == 0, res
dataset_id, document_id, chunk_ids = add_chunks_func
payload = payload(chunk_ids)
res = delete_chunks(WebApiAuth, dataset_id, document_id, payload)
assert res["code"] == 102, res
assert "rm_chunk deleted chunks" in res["message"], res
@pytest.mark.p3
def test_repeated_deletion(self, WebApiAuth, add_chunks_func):
_, doc_id, chunk_ids = add_chunks_func
payload = {"chunk_ids": chunk_ids, "doc_id": doc_id}
res = delete_chunks(WebApiAuth, payload)
dataset_id, document_id, chunk_ids = add_chunks_func
payload = {"chunk_ids": chunk_ids}
res = delete_chunks(WebApiAuth, dataset_id, document_id, payload)
assert res["code"] == 0, res
res = delete_chunks(WebApiAuth, payload)
res = delete_chunks(WebApiAuth, dataset_id, document_id, payload)
assert res["code"] == 102, res
assert res["message"] == "Index updating failure", res
assert res["message"] == f"rm_chunk deleted chunks 0, expect {len(chunk_ids)}", res
@pytest.mark.p3
def test_duplicate_deletion(self, WebApiAuth, add_chunks_func):
_, doc_id, chunk_ids = add_chunks_func
payload = {"chunk_ids": chunk_ids * 2, "doc_id": doc_id}
res = delete_chunks(WebApiAuth, payload)
dataset_id, document_id, chunk_ids = add_chunks_func
res = delete_chunks(WebApiAuth, dataset_id, document_id, {"chunk_ids": chunk_ids * 2})
assert res["code"] == 0, res
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
res = list_chunks(WebApiAuth, dataset_id, document_id)
assert res["code"] == 0, res
assert len(res["data"]["chunks"]) == 0, res
assert res["data"]["total"] == 0, res
@pytest.mark.p2
def test_delete_scalar_chunk_id_payload(self, WebApiAuth, add_chunks_func):
_, doc_id, chunk_ids = add_chunks_func
payload = {"chunk_ids": chunk_ids[0], "doc_id": doc_id}
res = delete_chunks(WebApiAuth, payload)
assert res["code"] == 0, res
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
assert res["code"] == 0, res
assert len(res["data"]["chunks"]) == 3, res
assert res["data"]["total"] == 3, res
@pytest.mark.p2
def test_delete_duplicate_ids_dedup_behavior(self, WebApiAuth, add_chunks_func):
_, doc_id, chunk_ids = add_chunks_func
payload = {"chunk_ids": [chunk_ids[0], chunk_ids[0]], "doc_id": doc_id}
res = delete_chunks(WebApiAuth, payload)
dataset_id, document_id, chunk_ids = add_chunks_func
res = delete_chunks(WebApiAuth, dataset_id, document_id, {"chunk_ids": [chunk_ids[0], chunk_ids[0]]})
assert res["code"] == 0, res
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
res = list_chunks(WebApiAuth, dataset_id, document_id)
assert res["code"] == 0, res
assert len(res["data"]["chunks"]) == 3, res
assert res["data"]["total"] == 3, res
@@ -122,16 +101,12 @@ class TestChunksDeletion:
@pytest.mark.p3
def test_concurrent_deletion(self, WebApiAuth, add_document):
count = 100
_, doc_id = add_document
chunk_ids = batch_add_chunks(WebApiAuth, doc_id, count)
dataset_id, document_id = add_document
chunk_ids = batch_add_chunks(WebApiAuth, dataset_id, document_id, count)
with ThreadPoolExecutor(max_workers=5) as executor:
futures = [
executor.submit(
delete_chunks,
WebApiAuth,
{"doc_id": doc_id, "chunk_ids": chunk_ids[i : i + 1]},
)
executor.submit(delete_chunks, WebApiAuth, dataset_id, document_id, {"chunk_ids": chunk_ids[i : i + 1]})
for i in range(count)
]
responses = list(as_completed(futures))
@@ -141,45 +116,40 @@ class TestChunksDeletion:
@pytest.mark.p3
def test_delete_1k(self, WebApiAuth, add_document):
chunks_num = 1_000
_, doc_id = add_document
chunk_ids = batch_add_chunks(WebApiAuth, doc_id, chunks_num)
dataset_id, document_id = add_document
chunk_ids = batch_add_chunks(WebApiAuth, dataset_id, document_id, chunks_num)
from time import sleep
sleep(1)
res = delete_chunks(WebApiAuth, {"doc_id": doc_id, "chunk_ids": chunk_ids})
res = delete_chunks(WebApiAuth, dataset_id, document_id, {"chunk_ids": chunk_ids})
assert res["code"] == 0
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
if res["code"] != 0:
assert False, res
res = list_chunks(WebApiAuth, dataset_id, document_id)
assert res["code"] == 0, res
assert len(res["data"]["chunks"]) == 0, res
assert res["data"]["total"] == 0, res
@pytest.mark.parametrize(
"payload, expected_code, expected_message, remaining",
[
pytest.param(None, 100, """TypeError("argument of type \'NoneType\' is not iterable")""", 5, marks=pytest.mark.skip),
pytest.param({"chunk_ids": ["invalid_id"]}, 102, "Index updating failure", 4, marks=pytest.mark.p3),
pytest.param("not json", 100, """UnboundLocalError("local variable \'duplicate_messages\' referenced before assignment")""", 5, marks=pytest.mark.skip(reason="pull/6376")),
pytest.param({"chunk_ids": ["invalid_id"]}, 102, "rm_chunk deleted chunks 0, expect 1", 4, marks=pytest.mark.p3),
pytest.param(lambda r: {"chunk_ids": r[:1]}, 0, "", 3, marks=pytest.mark.p3),
pytest.param(lambda r: {"chunk_ids": r}, 0, "", 0, marks=pytest.mark.p1),
pytest.param({"chunk_ids": []}, 0, "", 4, marks=pytest.mark.p3),
],
)
def test_basic_scenarios(self, WebApiAuth, add_chunks_func, payload, expected_code, expected_message, remaining):
_, doc_id, chunk_ids = add_chunks_func
dataset_id, document_id, chunk_ids = add_chunks_func
if callable(payload):
payload = payload(chunk_ids)
payload["doc_id"] = doc_id
res = delete_chunks(WebApiAuth, payload)
res = delete_chunks(WebApiAuth, dataset_id, document_id, payload)
assert res["code"] == expected_code, res
if res["code"] != 0:
assert res["message"] == expected_message, res
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
if res["code"] != 0:
assert False, res
res = list_chunks(WebApiAuth, dataset_id, document_id)
assert res["code"] == 0, res
assert len(res["data"]["chunks"]) == remaining, res
assert res["data"]["total"] == remaining, res

181
test/testcases/test_web_api/test_chunk_app/test_update_chunk.py
View File

@@ -13,16 +13,15 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
import base64
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from random import randint
from time import sleep
import pytest
from test_common import delete_document, list_chunks, update_chunk
from configs import INVALID_API_TOKEN
from libs.auth import RAGFlowWebApiAuth
from test_common import delete_document, list_chunks, update_chunk
@pytest.mark.p2
@@ -35,178 +34,144 @@ class TestAuthorization:
],
)
def test_invalid_auth(self, invalid_auth, expected_code, expected_message):
res = update_chunk(invalid_auth, {"doc_id": "doc_id", "chunk_id": "chunk_id", "content_with_weight": "test"})
res = update_chunk(invalid_auth, "dataset_id", "document_id", "chunk_id", {"content": "test"})
assert res["code"] == expected_code, res
assert res["message"] == expected_message, res
def _find_chunk(auth, dataset_id, document_id, chunk_id):
res = list_chunks(auth, dataset_id, document_id, params={"id": chunk_id})
assert res["code"] == 0, res
return res["data"]["chunks"][0]
class TestUpdateChunk:
@pytest.mark.p1
@pytest.mark.parametrize(
"payload, expected_code, expected_message",
[
({"content_with_weight": None}, 100, "TypeError('expected string or bytes-like object')"),
({"content_with_weight": ""}, 102, "`content_with_weight` is required"),
({"content_with_weight": 1}, 100, "TypeError('expected string or bytes-like object')"),
({"content_with_weight": "update chunk"}, 0, ""),
({"content_with_weight": " "}, 102, "`content_with_weight` is required"),
({"content_with_weight": "\n!?。；！？\"'"}, 0, ""),
({"content": None}, 0, ""),
({"content": ""}, 102, "`content` is required"),
pytest.param({"content": 1}, 100, "TypeError('expected string or bytes-like object')", marks=pytest.mark.skip),
({"content": "update chunk"}, 0, ""),
({"content": " "}, 102, "`content` is required"),
({"content": "\n!?。；！？\"'"}, 0, ""),
],
)
def test_content(self, WebApiAuth, add_chunks, payload, expected_code, expected_message):
_, doc_id, chunk_ids = add_chunks
dataset_id, document_id, chunk_ids = add_chunks
chunk_id = chunk_ids[0]
update_payload = {"doc_id": doc_id, "chunk_id": chunk_id}
if payload:
update_payload.update(payload)
res = update_chunk(WebApiAuth, update_payload)
res = update_chunk(WebApiAuth, dataset_id, document_id, chunk_id, payload)
assert res["code"] == expected_code, res
if expected_code != 0:
assert res["message"] == expected_message, res
else:
sleep(1)
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
for chunk in res["data"]["chunks"]:
if chunk["chunk_id"] == chunk_id:
assert chunk["content_with_weight"] == payload["content_with_weight"]
chunk = _find_chunk(WebApiAuth, dataset_id, document_id, chunk_id)
if payload["content"] is not None:
assert chunk["content"] == payload["content"]
@pytest.mark.p2
@pytest.mark.parametrize(
"payload, expected_code, expected_message",
[
({"important_kwd": ["a", "b", "c"]}, 0, ""),
({"important_kwd": [""]}, 0, ""),
({"important_kwd": [1]}, 100, "TypeError('sequence item 0: expected str instance, int found')"),
({"important_kwd": ["a", "a"]}, 0, ""),
({"important_kwd": "abc"}, 102, "`important_kwd` should be a list"),
({"important_kwd": 123}, 102, "`important_kwd` should be a list"),
({"important_keywords": ["a", "b", "c"]}, 0, ""),
({"important_keywords": [""]}, 0, ""),
({"important_keywords": [1]}, 100, "TypeError('sequence item 0: expected str instance, int found')"),
({"important_keywords": ["a", "a"]}, 0, ""),
({"important_keywords": "abc"}, 102, "`important_keywords` should be a list"),
({"important_keywords": 123}, 102, "`important_keywords` should be a list"),
],
)
def test_important_keywords(self, WebApiAuth, add_chunks, payload, expected_code, expected_message):
_, doc_id, chunk_ids = add_chunks
dataset_id, document_id, chunk_ids = add_chunks
chunk_id = chunk_ids[0]
update_payload = {"doc_id": doc_id, "chunk_id": chunk_id, "content_with_weight": "unchanged content"} # Add content_with_weight as it's required
if payload:
update_payload.update(payload)
res = update_chunk(WebApiAuth, update_payload)
res = update_chunk(WebApiAuth, dataset_id, document_id, chunk_id, payload)
assert res["code"] == expected_code, res
if expected_code != 0:
assert res["message"] == expected_message, res
else:
sleep(1)
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
for chunk in res["data"]["chunks"]:
if chunk["chunk_id"] == chunk_id:
assert chunk["important_kwd"] == payload["important_kwd"]
chunk = _find_chunk(WebApiAuth, dataset_id, document_id, chunk_id)
assert chunk["important_keywords"] == payload["important_keywords"]
@pytest.mark.p2
@pytest.mark.parametrize(
"payload, expected_code, expected_message",
[
({"question_kwd": ["a", "b", "c"]}, 0, ""),
({"question_kwd": [""]}, 100, """Exception('Error: 413 - {"error":"Input validation error: `inputs` cannot be empty","error_type":"Validation"}')"""),
({"question_kwd": [1]}, 100, "TypeError('sequence item 0: expected str instance, int found')"),
({"question_kwd": ["a", "a"]}, 0, ""),
({"question_kwd": "abc"}, 102, "`question_kwd` should be a list"),
({"question_kwd": 123}, 102, "`question_kwd` should be a list"),
({"questions": ["a", "b", "c"]}, 0, ""),
({"questions": [""]}, 0, ""),
({"questions": [1]}, 100, "TypeError('sequence item 0: expected str instance, int found')"),
({"questions": ["a", "a"]}, 0, ""),
({"questions": "abc"}, 102, "`questions` should be a list"),
({"questions": 123}, 102, "`questions` should be a list"),
],
)
def test_questions(self, WebApiAuth, add_chunks, payload, expected_code, expected_message):
_, doc_id, chunk_ids = add_chunks
dataset_id, document_id, chunk_ids = add_chunks
chunk_id = chunk_ids[0]
update_payload = {"doc_id": doc_id, "chunk_id": chunk_id, "content_with_weight": "unchanged content"} # Add content_with_weight as it's required
if payload:
update_payload.update(payload)
res = update_chunk(WebApiAuth, update_payload)
res = update_chunk(WebApiAuth, dataset_id, document_id, chunk_id, payload)
assert res["code"] == expected_code, res
if expected_code != 0:
assert res["message"] == expected_message, res
else:
sleep(1)
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
for chunk in res["data"]["chunks"]:
if chunk["chunk_id"] == chunk_id:
assert chunk["question_kwd"] == payload["question_kwd"]
chunk = _find_chunk(WebApiAuth, dataset_id, document_id, chunk_id)
assert chunk["questions"] == [str(q).strip() for q in payload["questions"] if str(q).strip()]
@pytest.mark.p2
@pytest.mark.parametrize(
"payload, expected_code, expected_message",
[
({"available_int": 1}, 0, ""),
({"available_int": 0}, 0, ""),
({"available": True}, 0, ""),
({"available": 1}, 0, ""),
({"available": False}, 0, ""),
({"available": 0}, 0, ""),
],
)
def test_available(self, WebApiAuth, add_chunks, payload, expected_code, expected_message):
_, doc_id, chunk_ids = add_chunks
dataset_id, document_id, chunk_ids = add_chunks
chunk_id = chunk_ids[0]
update_payload = {"doc_id": doc_id, "chunk_id": chunk_id, "content_with_weight": "unchanged content"}
if payload:
update_payload.update(payload)
res = update_chunk(WebApiAuth, update_payload)
res = update_chunk(WebApiAuth, dataset_id, document_id, chunk_id, payload)
assert res["code"] == expected_code, res
if expected_code != 0:
assert res["message"] == expected_message, res
else:
sleep(1)
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
for chunk in res["data"]["chunks"]:
if chunk["chunk_id"] == chunk_id:
assert chunk["available_int"] == payload["available_int"]
chunk = _find_chunk(WebApiAuth, dataset_id, document_id, chunk_id)
assert chunk["available"] == bool(payload["available"])
@pytest.mark.p2
def test_update_chunk_qa_multiline_content(self, WebApiAuth, add_chunks):
_, doc_id, chunk_ids = add_chunks
payload = {"doc_id": doc_id, "chunk_id": chunk_ids[0], "content_with_weight": "Question line\nAnswer line"}
res = update_chunk(WebApiAuth, payload)
dataset_id, document_id, chunk_ids = add_chunks
payload = {"content": "Question line\nAnswer line"}
res = update_chunk(WebApiAuth, dataset_id, document_id, chunk_ids[0], payload)
assert res["code"] == 0, res
sleep(1)
res = list_chunks(WebApiAuth, {"doc_id": doc_id})
assert res["code"] == 0, res
chunk = next(chunk for chunk in res["data"]["chunks"] if chunk["chunk_id"] == chunk_ids[0])
assert chunk["content_with_weight"] == payload["content_with_weight"], res
@pytest.mark.p2
def test_update_chunk_with_image_payload(self, WebApiAuth, add_chunks):
_, doc_id, chunk_ids = add_chunks
payload = {
"doc_id": doc_id,
"chunk_id": chunk_ids[0],
"content_with_weight": "content with image",
"image_base64": base64.b64encode(b"img").decode("utf-8"),
"img_id": "bucket-name",
}
res = update_chunk(WebApiAuth, payload)
assert res["code"] == 0, res
chunk = _find_chunk(WebApiAuth, dataset_id, document_id, chunk_ids[0])
assert chunk["content"] == payload["content"], chunk
@pytest.mark.p3
@pytest.mark.parametrize(
"doc_id_param, expected_code, expected_message",
"document_id, expected_code, expected_message",
[
("", 102, "Tenant not found!"),
("invalid_doc_id", 102, "Tenant not found!"),
("invalid_doc_id", 102, "You don't own the document invalid_doc_id."),
],
)
def test_invalid_document_id_for_update(self, WebApiAuth, add_chunks, doc_id_param, expected_code, expected_message):
_, _, chunk_ids = add_chunks
chunk_id = chunk_ids[0]
payload = {"doc_id": doc_id_param, "chunk_id": chunk_id, "content_with_weight": "test content"}
res = update_chunk(WebApiAuth, payload)
def test_invalid_document_id_for_update(self, WebApiAuth, add_chunks, document_id, expected_code, expected_message):
dataset_id, _, chunk_ids = add_chunks
res = update_chunk(WebApiAuth, dataset_id, document_id, chunk_ids[0], {"content": "test content"})
assert res["code"] == expected_code
assert expected_message in res["message"]
@pytest.mark.p3
def test_repeated_update_chunk(self, WebApiAuth, add_chunks):
_, doc_id, chunk_ids = add_chunks
payload1 = {"doc_id": doc_id, "chunk_id": chunk_ids[0], "content_with_weight": "chunk test 1"}
res = update_chunk(WebApiAuth, payload1)
dataset_id, document_id, chunk_ids = add_chunks
res = update_chunk(WebApiAuth, dataset_id, document_id, chunk_ids[0], {"content": "chunk test 1"})
assert res["code"] == 0
payload2 = {"doc_id": doc_id, "chunk_id": chunk_ids[0], "content_with_weight": "chunk test 2"}
res = update_chunk(WebApiAuth, payload2)
res = update_chunk(WebApiAuth, dataset_id, document_id, chunk_ids[0], {"content": "chunk test 2"})
assert res["code"] == 0
@pytest.mark.p3
@@ -215,17 +180,11 @@ class TestUpdateChunk:
[
({"unknown_key": "unknown_value"}, 0, ""),
({}, 0, ""),
pytest.param(None, 100, """TypeError("int() argument must be a string, a bytes-like object or a real number, not 'NoneType'")""", marks=pytest.mark.skip),
],
)
def test_invalid_params(self, WebApiAuth, add_chunks, payload, expected_code, expected_message):
_, doc_id, chunk_ids = add_chunks
chunk_id = chunk_ids[0]
update_payload = {"doc_id": doc_id, "chunk_id": chunk_id, "content_with_weight": "unchanged content"}
if payload is not None:
update_payload.update(payload)
res = update_chunk(WebApiAuth, update_payload)
dataset_id, document_id, chunk_ids = add_chunks
res = update_chunk(WebApiAuth, dataset_id, document_id, chunk_ids[0], payload)
assert res["code"] == expected_code, res
if expected_code != 0:
assert res["message"] == expected_message, res
@@ -234,14 +193,17 @@ class TestUpdateChunk:
@pytest.mark.skipif(os.getenv("DOC_ENGINE") == "infinity", reason="issues/6554")
def test_concurrent_update_chunk(self, WebApiAuth, add_chunks):
count = 50
_, doc_id, chunk_ids = add_chunks
dataset_id, document_id, chunk_ids = add_chunks
with ThreadPoolExecutor(max_workers=5) as executor:
futures = [
executor.submit(
update_chunk,
WebApiAuth,
{"doc_id": doc_id, "chunk_id": chunk_ids[randint(0, 3)], "content_with_weight": f"update chunk test {i}"},
dataset_id,
document_id,
chunk_ids[randint(0, 3)],
{"content": f"update chunk test {i}"},
)
for i in range(count)
]
@@ -251,9 +213,8 @@ class TestUpdateChunk:
@pytest.mark.p3
def test_update_chunk_to_deleted_document(self, WebApiAuth, add_chunks):
kb_id, doc_id, chunk_ids = add_chunks
delete_document(WebApiAuth, kb_id, {"ids": [doc_id]})
payload = {"doc_id": doc_id, "chunk_id": chunk_ids[0], "content_with_weight": "test content"}
res = update_chunk(WebApiAuth, payload)
dataset_id, document_id, chunk_ids = add_chunks
delete_document(WebApiAuth, dataset_id, {"ids": [document_id]})
res = update_chunk(WebApiAuth, dataset_id, document_id, chunk_ids[0], {"content": "test content"})
assert res["code"] == 102, res
assert res["message"] == "Tenant not found!", res
assert res["message"] in [f"You don't own the document {document_id}.", f"Can't find this chunk {chunk_ids[0]}"]

43
test/testcases/test_web_api/test_common.py
View File

@@ -28,7 +28,8 @@ HEADERS = {"Content-Type": "application/json"}
KB_APP_URL = f"/{VERSION}/kb"
DATASETS_URL = f"/api/{VERSION}/datasets"
DOCUMENT_APP_URL = f"/{VERSION}/document"
CHUNK_API_URL = f"/{VERSION}/chunk"
CHUNK_APP_URL = f"/{VERSION}/chunk"
CHUNK_API_URL = f"/api/{VERSION}/datasets/{{dataset_id}}/documents/{{document_id}}/chunks"
# SESSION_WITH_CHAT_ASSISTANT_API_URL = "/api/v1/chats/{chat_id}/sessions"
# SESSION_WITH_AGENT_API_URL = "/api/v1/agents/{agent_id}/sessions"
MEMORY_API_URL = f"/api/{VERSION}/memories"
@@ -441,47 +442,53 @@ def bulk_upload_documents(auth, kb_id, num, tmp_path):
return document_ids
# CHUNK APP
def add_chunk(auth, payload=None, *, headers=HEADERS, data=None):
res = requests.post(url=f"{HOST_ADDRESS}{CHUNK_API_URL}/create", headers=headers, auth=auth, json=payload, data=data)
# CHUNK MANAGEMENT
def add_chunk(auth, dataset_id, document_id, payload=None, *, headers=HEADERS, data=None):
url = f"{HOST_ADDRESS}{CHUNK_API_URL}".format(dataset_id=dataset_id, document_id=document_id)
res = requests.post(url=url, headers=headers, auth=auth, json=payload, data=data)
return res.json()
def list_chunks(auth, payload=None, *, headers=HEADERS):
res = requests.post(url=f"{HOST_ADDRESS}{CHUNK_API_URL}/list", headers=headers, auth=auth, json=payload)
def list_chunks(auth, dataset_id, document_id, params=None, *, headers=HEADERS):
url = f"{HOST_ADDRESS}{CHUNK_API_URL}".format(dataset_id=dataset_id, document_id=document_id)
res = requests.get(url=url, headers=headers, auth=auth, params=params)
return res.json()
def get_chunk(auth, params=None, *, headers=HEADERS):
res = requests.get(url=f"{HOST_ADDRESS}{CHUNK_API_URL}/get", headers=headers, auth=auth, params=params)
def get_chunk(auth, dataset_id, document_id, chunk_id, *, headers=HEADERS):
url = f"{HOST_ADDRESS}{CHUNK_API_URL}/{chunk_id}".format(dataset_id=dataset_id, document_id=document_id)
res = requests.get(url=url, headers=headers, auth=auth)
return res.json()
def update_chunk(auth, payload=None, *, headers=HEADERS):
res = requests.post(url=f"{HOST_ADDRESS}{CHUNK_API_URL}/set", headers=headers, auth=auth, json=payload)
def update_chunk(auth, dataset_id, document_id, chunk_id, payload=None, *, headers=HEADERS):
url = f"{HOST_ADDRESS}{CHUNK_API_URL}/{chunk_id}".format(dataset_id=dataset_id, document_id=document_id)
res = requests.patch(url=url, headers=headers, auth=auth, json=payload)
return res.json()
def switch_chunks(auth, payload=None, *, headers=HEADERS):
res = requests.post(url=f"{HOST_ADDRESS}{CHUNK_API_URL}/switch", headers=headers, auth=auth, json=payload)
def switch_chunks(auth, dataset_id, document_id, payload=None, *, headers=HEADERS):
url = f"{HOST_ADDRESS}{CHUNK_API_URL}".format(dataset_id=dataset_id, document_id=document_id)
res = requests.patch(url=url, headers=headers, auth=auth, json=payload)
return res.json()
def delete_chunks(auth, payload=None, *, headers=HEADERS):
res = requests.post(url=f"{HOST_ADDRESS}{CHUNK_API_URL}/rm", headers=headers, auth=auth, json=payload)
def delete_chunks(auth, dataset_id, document_id, payload=None, *, headers=HEADERS):
url = f"{HOST_ADDRESS}{CHUNK_API_URL}".format(dataset_id=dataset_id, document_id=document_id)
res = requests.delete(url=url, headers=headers, auth=auth, json=payload)
return res.json()
def retrieval_chunks(auth, payload=None, *, headers=HEADERS):
res = requests.post(url=f"{HOST_ADDRESS}{CHUNK_API_URL}/retrieval_test", headers=headers, auth=auth, json=payload)
res = requests.post(url=f"{HOST_ADDRESS}{CHUNK_APP_URL}/retrieval_test", headers=headers, auth=auth, json=payload)
return res.json()
def batch_add_chunks(auth, doc_id, num):
def batch_add_chunks(auth, dataset_id, document_id, num):
chunk_ids = []
for i in range(num):
res = add_chunk(auth, {"doc_id": doc_id, "content_with_weight": f"chunk test {i}"})
chunk_ids.append(res["data"]["chunk_id"])
res = add_chunk(auth, dataset_id, document_id, {"content": f"chunk test {i}"})
chunk_ids.append(res["data"]["chunk"]["id"])
return chunk_ids

7
test/testcases/test_web_api/test_kb_app/test_kb_tags_meta.py
View File

@@ -60,10 +60,11 @@ def _seed_tag(auth, kb_id, document_id, chunk_id):
tag = f"tag_{uuid.uuid4().hex[:8]}"
res = update_chunk(
auth,
kb_id,
document_id,
chunk_id,
{
"doc_id": document_id,
"chunk_id": chunk_id,
"content_with_weight": f"tag seed {tag}",
"content": f"tag seed {tag}",
"tag_kwd": [tag],
},
)

13
test/unit_test/common/test_delete_query_construction.py
View File

@@ -212,18 +212,17 @@ class TestDeleteQueryConstruction:
assert len(kw_terms) == 1
class TestChunkAppDeleteCondition:
class TestChunkApiDeleteCondition:
"""
Tests that verify the chunk_app.py rm endpoint passes the correct
Tests that verify the RESTful chunk DELETE endpoint passes the correct
condition to docStoreConn.delete.
"""
def test_rm_endpoint_includes_doc_id_in_condition(self):
"""
The /chunk/rm endpoint MUST include doc_id in the condition
The /api/v1/datasets/<dataset_id>/documents/<document_id>/chunks endpoint
MUST include doc_id in the condition
passed to settings.docStoreConn.delete.
This is the fix applied to api/apps/chunk_app.py
"""
# Simulate what the rm endpoint should construct
req = {
@@ -248,7 +247,7 @@ class TestChunkAppDeleteCondition:
class TestSDKDocDeleteCondition:
"""
Tests that verify the SDK doc.py rm_chunk endpoint constructs
Tests that verify the RESTful chunk delete endpoint constructs
the correct deletion condition.
"""
@@ -261,7 +260,7 @@ class TestSDKDocDeleteCondition:
document_id = "doc456"
chunk_ids = ["chunk1", "chunk2"]
# The CORRECT condition construction (from sdk/doc.py):
# The CORRECT condition construction (from restful_apis/chunk_api.py):
condition = {"doc_id": document_id}
if chunk_ids:
condition["id"] = chunk_ids

8
web/src/hooks/route-hook.ts
View File

@@ -2,6 +2,7 @@ import {
KnowledgeRouteKey,
KnowledgeSearchParams,
} from '@/constants/knowledge';
import { Routes } from '@/routes';
import { useCallback } from 'react';
import { useLocation, useNavigate, useSearchParams } from 'react-router';
@@ -27,13 +28,16 @@ export const useThirdPathName = () => {
export const useGetKnowledgeSearchParams = () => {
const [currentQueryParameters] = useSearchParams();
const { pathname } = useLocation();
const isDataflowResultPage = pathname === Routes.DataflowResult;
return {
type: currentQueryParameters.get(KnowledgeSearchParams.Type) || '',
documentId:
currentQueryParameters.get(KnowledgeSearchParams.DocumentId) || '',
knowledgeId:
currentQueryParameters.get(KnowledgeSearchParams.KnowledgeId) || '',
knowledgeId: isDataflowResultPage
? currentQueryParameters.get('knowledgeId') || ''
: currentQueryParameters.get(KnowledgeSearchParams.KnowledgeId) || '',
};
};

36
web/src/hooks/use-chunk-request.ts
View File

@@ -40,6 +40,7 @@ export const useSelectChunkList = () => {
export const useDeleteChunk = () => {
const queryClient = useQueryClient();
const { setPaginationParams } = useSetPaginationParams();
const { knowledgeId } = useGetKnowledgeSearchParams();
const {
data,
isPending: loading,
@@ -47,7 +48,10 @@ export const useDeleteChunk = () => {
} = useMutation({
mutationKey: ['deleteChunk'],
mutationFn: async (params: { chunkIds: string[]; doc_id: string }) => {
const { data } = await kbService.rmChunk(params);
const { data } = await kbService.rmChunk({
...params,
kb_id: knowledgeId,
});
if (data.code === 0) {
setPaginationParams(1);
queryClient.invalidateQueries({ queryKey: ['fetchChunkList'] });
@@ -62,6 +66,7 @@ export const useDeleteChunk = () => {
export const useCreateChunk = () => {
const { t } = useTranslation();
const queryClient = useQueryClient();
const { knowledgeId } = useGetKnowledgeSearchParams();
const {
data,
@@ -74,7 +79,10 @@ export const useCreateChunk = () => {
if (payload.chunk_id) {
service = kbService.setChunk;
}
const { data } = await service(payload);
const { data } = await service({
...payload,
kb_id: payload.kb_id || knowledgeId,
});
if (data.code === 0) {
message.success(t('message.created'));
setTimeout(() => {
@@ -88,14 +96,20 @@ export const useCreateChunk = () => {
return { data, loading, createChunk: mutateAsync };
};
export const useFetchChunk = (chunkId?: string): ResponseType<any> => {
export const useFetchChunk = (
chunkId?: string,
documentId?: string,
): ResponseType<any> => {
const { knowledgeId } = useGetKnowledgeSearchParams();
const { data } = useQuery({
queryKey: ['fetchChunk'],
enabled: !!chunkId,
queryKey: ['fetchChunk', knowledgeId, documentId, chunkId],
enabled: !!chunkId && !!documentId && !!knowledgeId,
initialData: {},
gcTime: 0,
queryFn: async () => {
const data = await kbService.getChunk({
kb_id: knowledgeId,
doc_id: documentId,
chunk_id: chunkId,
});
@@ -115,7 +129,7 @@ export const useFetchNextChunkList = (
}> &
IChunkListResult => {
const { pagination, setPagination } = useGetPaginationWithRouter();
const { documentId } = useGetKnowledgeSearchParams();
const { documentId, knowledgeId } = useGetKnowledgeSearchParams();
const { searchString, handleInputChange } = useHandleSearchChange();
const [available, setAvailable] = useState<number | undefined>();
const debouncedSearchString = useDebounce(searchString, { wait: 500 });
@@ -127,6 +141,7 @@ export const useFetchNextChunkList = (
} = useQuery({
queryKey: [
'fetchChunkList',
knowledgeId,
documentId,
pagination.current,
pagination.pageSize,
@@ -136,9 +151,10 @@ export const useFetchNextChunkList = (
placeholderData: (previousData: any) =>
previousData ?? { data: [], total: 0, documentInfo: {} }, // https://github.com/TanStack/query/issues/8183
gcTime: 0,
enabled,
enabled: enabled && !!knowledgeId && !!documentId,
queryFn: async () => {
const { data } = await kbService.chunkList({
kb_id: knowledgeId,
doc_id: documentId,
page: pagination.current,
size: pagination.pageSize,
@@ -195,6 +211,7 @@ export const useFetchNextChunkList = (
export const useSwitchChunk = () => {
const { t } = useTranslation();
const { knowledgeId } = useGetKnowledgeSearchParams();
const {
data,
isPending: loading,
@@ -206,7 +223,10 @@ export const useSwitchChunk = () => {
available_int?: number;
doc_id: string;
}) => {
const { data } = await kbService.switchChunk(params);
const { data } = await kbService.switchChunk({
...params,
kb_id: knowledgeId,
});
if (data.code === 0) {
message.success(t('message.modified'));
}

10
web/src/pages/chunk/parsed-result/add-knowledge/components/knowledge-chunk/components/chunk-creating-modal/index.tsx
View File

@@ -26,7 +26,6 @@ import type { ChunkDocType } from '@/interfaces/database/knowledge';
import React, { useCallback, useEffect, useState } from 'react';
import { FieldValues, FormProvider, useForm } from 'react-hook-form';
import { useTranslation } from 'react-i18next';
import { useDeleteChunkByIds } from '../../hooks';
import {
transformTagFeaturesArrayToObject,
transformTagFeaturesObjectToArray,
@@ -75,8 +74,7 @@ const ChunkCreatingModal: React.FC<IModalProps<any> & kFProps> = ({
},
});
const [checked, setChecked] = useState(false);
const { removeChunk } = useDeleteChunkByIds();
const { data } = useFetchChunk(chunkId);
const { data } = useFetchChunk(chunkId, doc_id);
const { t } = useTranslation();
const isEditMode = !!chunkId;
@@ -99,12 +97,6 @@ const ChunkCreatingModal: React.FC<IModalProps<any> & kFProps> = ({
const handleOk = form.handleSubmit(onSubmit);
const handleRemove = useCallback(() => {
if (chunkId) {
return removeChunk([chunkId], doc_id);
}
}, [chunkId, doc_id, removeChunk]);
const handleCheck = useCallback(() => {
setChecked(!checked);
}, [checked]);

2
web/src/pages/dataflow-result/components/chunk-creating-modal/index.tsx
View File

@@ -57,7 +57,7 @@ const ChunkCreatingModal: React.FC<IModalProps<any> & kFProps> = ({
});
const [checked, setChecked] = useState(false);
const { removeChunk } = useDeleteChunkByIds();
const { data } = useFetchChunk(chunkId);
const { data } = useFetchChunk(chunkId, doc_id);
const { t } = useTranslation();
const isTagParser = parserId === 'tag';

9
web/src/pages/dataset/dataset-overview/overview-table.tsx
View File

@@ -51,7 +51,7 @@ import { DocumentLog, FileLogsTableProps, IFileLogItem } from './interface';
export const getFileLogsTableColumns = (
t: TFunction<'translation', string>,
showLog: (row: Row<IFileLogItem & DocumentLog>, active: LogTabs) => void,
kowledgeId: string,
knowledgeId: string,
navigateToDataflowResult: (
props: NavigateToDataflowResultProps,
) => () => void,
@@ -210,7 +210,8 @@ export const getFileLogsTableColumns = (
size="icon-sm"
onClick={navigateToDataflowResult({
id: row.original.id,
[PipelineResultSearchParams.KnowledgeId]: kowledgeId,
[PipelineResultSearchParams.KnowledgeId]:
row.original.kb_id || knowledgeId,
[PipelineResultSearchParams.DocumentId]:
row.original.document_id,
[PipelineResultSearchParams.IsReadOnly]: 'false',
@@ -358,7 +359,7 @@ const FileLogsTable: FC<FileLogsTableProps> = ({
const [isModalVisible, setIsModalVisible] = useState(false);
const { navigateToDataflowResult } = useNavigatePage();
const [logInfo, setLogInfo] = useState<IFileLogItem>();
const kowledgeId = useParams().id;
const knowledgeId = useParams().id;
const showLog = (row: Row<IFileLogItem & DocumentLog>) => {
const logDetail = {
taskId: row.original?.dsl?.task_id,
@@ -382,7 +383,7 @@ const FileLogsTable: FC<FileLogsTableProps> = ({
? getFileLogsTableColumns(
t,
showLog,
kowledgeId || '',
knowledgeId || '',
navigateToDataflowResult,
dataSourceInfo,
)

165
web/src/services/knowledge-service.ts
View File

@@ -21,12 +21,6 @@ const {
documentCreate,
documentChangeParser,
documentThumbnails,
chunkList,
createChunk,
setChunk,
getChunk,
switchChunk,
rmChunk,
retrievalTest,
documentRun,
documentUpload,
@@ -97,31 +91,6 @@ const methods = {
url: setMeta,
method: 'post',
},
// chunk管理
chunkList: {
url: chunkList,
method: 'post',
},
createChunk: {
url: createChunk,
method: 'post',
},
setChunk: {
url: setChunk,
method: 'post',
},
getChunk: {
url: getChunk,
method: 'get',
},
switchChunk: {
url: switchChunk,
method: 'post',
},
rmChunk: {
url: rmChunk,
method: 'post',
},
retrievalTest: {
url: retrievalTest,
method: 'post',
@@ -178,7 +147,139 @@ const methods = {
},
};
const kbService = registerServer<keyof typeof methods>(methods, request);
const baseKbService = registerServer<keyof typeof methods>(methods, request);
const getDatasetId = (params: Record<string, any>) =>
params.dataset_id || params.kb_id || params.knowledge_id;
const getDocumentId = (params: Record<string, any>) =>
params.document_id || params.doc_id;
const mapChunkToLegacy = (chunk: Record<string, any>) => ({
...chunk,
chunk_id: chunk.chunk_id || chunk.id,
content_with_weight: chunk.content_with_weight || chunk.content,
doc_id: chunk.doc_id || chunk.document_id,
doc_name: chunk.doc_name || chunk.docnm_kwd,
image_id: chunk.image_id || chunk.img_id,
important_kwd: chunk.important_kwd || chunk.important_keywords || [],
question_kwd: chunk.question_kwd || chunk.questions || [],
available_int: chunk.available_int ?? (chunk.available === false ? 0 : 1),
positions: chunk.positions || chunk.position_int || [],
});
const mapDocumentToLegacy = (doc: Record<string, any>) => ({
...doc,
chunk_num: doc.chunk_num ?? doc.chunk_count,
kb_id: doc.kb_id || doc.dataset_id,
});
const mapChunkPayloadToRest = (payload: Record<string, any>) => ({
content: payload.content ?? payload.content_with_weight,
important_keywords: payload.important_keywords ?? payload.important_kwd,
questions: payload.questions ?? payload.question_kwd,
tag_kwd: payload.tag_kwd,
tag_feas: payload.tag_feas,
positions: payload.positions,
available:
payload.available ??
(payload.available_int === undefined
? undefined
: payload.available_int === 1),
image_base64: payload.image_base64,
});
const getAvailableParam = (available?: number) => {
if (available === undefined) {
return undefined;
}
return available === 1 ? 'true' : 'false';
};
const chunkService = {
chunkList: async (params: Record<string, any>) => {
const datasetId = getDatasetId(params);
const documentId = getDocumentId(params);
const response = await request.get(api.chunkList(datasetId, documentId), {
params: {
page: params.page,
page_size: params.page_size || params.size,
keywords: params.keywords,
available: getAvailableParam(params.available_int),
},
});
if (response.data?.code === 0) {
response.data.data = {
...response.data.data,
chunks: (response.data.data?.chunks || []).map(mapChunkToLegacy),
doc: mapDocumentToLegacy(response.data.data?.doc || {}),
};
}
return response;
},
createChunk: async (payload: Record<string, any>) => {
const datasetId = getDatasetId(payload);
const documentId = getDocumentId(payload);
const response = await request.post(api.chunkList(datasetId, documentId), {
data: mapChunkPayloadToRest(payload),
});
if (response.data?.code === 0 && response.data.data?.chunk) {
response.data.data.chunk = mapChunkToLegacy(response.data.data.chunk);
}
return response;
},
setChunk: (payload: Record<string, any>) => {
const datasetId = getDatasetId(payload);
const documentId = getDocumentId(payload);
const chunkId = payload.chunk_id || payload.id;
return request.patch(api.chunkDetail(datasetId, documentId, chunkId), {
data: mapChunkPayloadToRest(payload),
});
},
getChunk: async (params: Record<string, any>) => {
const datasetId = getDatasetId(params);
const documentId = getDocumentId(params);
const chunkId = params.chunk_id || params.id;
const response = await request.get(
api.chunkDetail(datasetId, documentId, chunkId),
);
if (response.data?.code === 0) {
response.data.data = mapChunkToLegacy(response.data.data || {});
}
return response;
},
switchChunk: (params: Record<string, any>) => {
const datasetId = getDatasetId(params);
const documentId = getDocumentId(params);
return request.patch(api.chunkList(datasetId, documentId), {
data: {
chunk_ids: params.chunk_ids || params.chunkIds,
available_int: params.available_int,
},
});
},
rmChunk: (params: Record<string, any>) => {
const datasetId = getDatasetId(params);
const documentId = getDocumentId(params);
return request.delete(api.chunkList(datasetId, documentId), {
data: {
chunk_ids: params.chunk_ids || params.chunkIds,
delete_all: params.delete_all,
},
});
},
};
const kbService = {
...baseKbService,
...chunkService,
};
export const listTag = (knowledgeId: string) =>
request.get(api.listTag(knowledgeId));

10
web/src/utils/api.ts
View File

@@ -99,12 +99,10 @@ export default {
renameTag: (knowledgeId: string) => `${webAPI}/kb/${knowledgeId}/rename_tag`,
// chunk
chunkList: `${webAPI}/chunk/list`,
createChunk: `${webAPI}/chunk/create`,
setChunk: `${webAPI}/chunk/set`,
getChunk: `${webAPI}/chunk/get`,
switchChunk: `${webAPI}/chunk/switch`,
rmChunk: `${webAPI}/chunk/rm`,
chunkList: (datasetId: string, documentId: string) =>
`${restAPIv1}/datasets/${datasetId}/documents/${documentId}/chunks`,
chunkDetail: (datasetId: string, documentId: string, chunkId: string) =>
`${restAPIv1}/datasets/${datasetId}/documents/${documentId}/chunks/${chunkId}`,
retrievalTest: `${webAPI}/chunk/retrieval_test`,
knowledgeGraph: `${webAPI}/chunk/knowledge_graph`,