zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-20 14:41:05 +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

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],
},
)