From f063cfdb192e9a65ced2896a068f0854dd21581b Mon Sep 17 00:00:00 2001 From: deadtrickster Date: Wed, 5 Aug 2026 04:28:46 +0200 Subject: [PATCH] fix(api): decrement knowledgebase counters on SDK re-parse / stop-parse (#17236) --- api/apps/restful_apis/chunk_api.py | 25 +- api/apps/restful_apis/document_api.py | 12 +- api/apps/services/document_api_service.py | 25 +- api/db/services/document_counter_service.py | 51 +++ .../test_doc_sdk_routes_unit.py | 290 ++++++++++++++---- .../test_chunk_app/test_chunk_routes_unit.py | 124 ++++++++ 6 files changed, 446 insertions(+), 81 deletions(-) create mode 100644 api/db/services/document_counter_service.py diff --git a/api/apps/restful_apis/chunk_api.py b/api/apps/restful_apis/chunk_api.py index 8ce7319fbc..2d254b8d96 100644 --- a/api/apps/restful_apis/chunk_api.py +++ b/api/apps/restful_apis/chunk_api.py @@ -33,6 +33,7 @@ from api.db.joint_services.tenant_model_service import ( ) from api.db.db_models import Document, Task from api.db.services.doc_metadata_service import DocMetadataService +from api.db.services.document_counter_service import release_reparse_counters from api.db.services.document_service import DocumentService from api.db.services.file2document_service import File2DocumentService from api.db.services.knowledgebase_service import KnowledgebaseService @@ -179,6 +180,22 @@ def _enrich_chunks_with_document_metadata(chunks: list[dict], metadata_fields=No enrich_chunks_with_document_metadata(chunks, metadata_fields) +def _release_doc_counters(doc): + """Roll back the document's and knowledgebase's chunk/token/duration counters + so a re-parse starts from zero. Callers that delete a document's chunks must + do this, otherwise the removed counts stay in the knowledgebase total. The + release re-reads the row under a lock (see release_reparse_counters) so it is + safe against a worker still parsing the document. Returns an error result if + the document is gone, else None. + """ + try: + release_reparse_counters(doc.id) + except LookupError: + logging.exception("Failed to release counters for document %s in knowledgebase %s", doc.id, doc.kb_id) + return get_error_data_result(message=f"Document {doc.id} not found") + return None + + @manager.route("/datasets//chunks", methods=["POST"]) # noqa: F821 @login_required @add_tenant_id_to_kwargs @@ -211,7 +228,7 @@ async def parse(tenant_id, dataset_id): continue if not doc: return get_error_data_result(message=f"you don't own the document {id}") - info = {"run": "1", "progress": 0, "progress_msg": "", "chunk_num": 0, "token_num": 0} + info = {"run": "1", "progress": 0, "progress_msg": ""} if ( DocumentService.filter_update( [ @@ -223,6 +240,8 @@ async def parse(tenant_id, dataset_id): == 0 ): return get_error_data_result("Can't parse document that is currently being processed") + if err := _release_doc_counters(doc[0]): + return err index_name = search.index_name(dataset_tenant_id) if settings.docStoreConn.index_exist(index_name, doc[0].kb_id): settings.docStoreConn.delete({"doc_id": id}, index_name, doc[0].kb_id) @@ -283,8 +302,10 @@ async def stop_parsing(tenant_id, dataset_id): data={"error_code": DOC_STOP_PARSING_INVALID_STATE_ERROR_CODE}, ) cancel_all_task_of(id) - info = {"run": "2", "progress": 0, "chunk_num": 0} + info = {"run": "2", "progress": 0} DocumentService.update_by_id(id, info) + if err := _release_doc_counters(doc[0]): + return err index_name = search.index_name(dataset_tenant_id) if settings.docStoreConn.index_exist(index_name, doc[0].kb_id): settings.docStoreConn.delete({"doc_id": doc[0].id}, index_name, doc[0].kb_id) diff --git a/api/apps/restful_apis/document_api.py b/api/apps/restful_apis/document_api.py index e2931c911e..3e92e4ecd8 100644 --- a/api/apps/restful_apis/document_api.py +++ b/api/apps/restful_apis/document_api.py @@ -40,6 +40,7 @@ from api.db import VALID_FILE_TYPES, FileType from api.db.db_models import API4Conversation, DB from api.db.services import duplicate_name from api.db.services.doc_metadata_service import DocMetadataService +from api.db.services.document_counter_service import release_reparse_counters from api.db.db_models import Task from api.db.services.document_service import DocumentService from api.db.services.file2document_service import File2DocumentService @@ -1734,13 +1735,22 @@ async def stop_parse_documents(tenant_id, dataset_id): continue cancel_all_task_of(doc_id) + # Release the document's partial chunk/token counts from the + # knowledgebase aggregate under the row lock (see + # release_reparse_counters). This is the sole counter adjustment, + # so the status update below must not touch chunk_num. + try: + release_reparse_counters(doc_id) + except LookupError: + logging.exception("Failed to release counters for document %s during stop-parse", doc_id) + errors.append(f"Document not found: {doc_id}") + continue cancel_doc_msg = f"\n{datetime.now().strftime('%H:%M:%S')} Task stopped by user." DocumentService.update_by_id( doc_id, { "run": str(TaskStatus.CANCEL.value), "progress": 0, - "chunk_num": 0, "progress_msg": (doc.progress_msg or "") + cancel_doc_msg, }, ) diff --git a/api/apps/services/document_api_service.py b/api/apps/services/document_api_service.py index d6ca3e5623..e8f7ab78a9 100644 --- a/api/apps/services/document_api_service.py +++ b/api/apps/services/document_api_service.py @@ -15,6 +15,7 @@ # import logging +from api.db.services.document_counter_service import release_reparse_counters from api.db.services.document_service import DocumentService from api.db.services.file2document_service import File2DocumentService from api.db.services.file_service import FileService @@ -128,22 +129,14 @@ def reset_document_for_reparse(doc, tenant_id, parser_id=None, pipeline_id=None) if not e: return get_error_data_result(message="document not found") - # Update document statistics before deleting all document rows. Pipeline - # compilation rows may exist even when token_num is zero, so the doc-store - # cleanup must not be gated by the document counters. - if doc.token_num > 0: - try: - e = DocumentService.increment_chunk_num( - doc.id, - doc.kb_id, - doc.token_num * -1, - doc.chunk_num * -1, - doc.process_duration * -1, - ) - except LookupError: - return get_error_data_result(message="document not found") - if not e: - return get_error_data_result(message="document not found") + # Release the document's chunk/token/duration counters from the knowledgebase + # aggregate under a row lock before clearing the chunks. release_reparse_counters + # guards the zero case internally, so the doc-store cleanup below still runs for + # pipeline compilation rows that exist even when token_num is zero. + try: + release_reparse_counters(doc.id) + except LookupError: + return get_error_data_result(message="Document not found!") settings.docStoreConn.delete({"doc_id": doc.id}, search.index_name(tenant_id), doc.kb_id) # Delete chunk images diff --git a/api/db/services/document_counter_service.py b/api/db/services/document_counter_service.py new file mode 100644 index 0000000000..d079b5c54d --- /dev/null +++ b/api/db/services/document_counter_service.py @@ -0,0 +1,51 @@ +# +# Copyright 2024 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import logging + +from api.db.db_models import DB, Document +from api.db.services.document_service import DocumentService + + +def release_reparse_counters(doc_id): + """Roll back a document's chunk, token, and duration counters and the owning + knowledgebase's chunk/token totals so a re-parse starts from zero. + + The counters are re-read under a ``FOR UPDATE`` row lock in the same + transaction as the decrement, so the release subtracts the row's committed + value at release time rather than a request-time snapshot a concurrent worker + may have already moved past. ``increment_chunk_num`` updates both ledgers + together. This does not serialize a worker that writes its final counts in a + separate transaction after the release commits; fully closing the + stop-parse-during-parse race needs a worker-side cancel check. + + Raises ``LookupError`` if the document row no longer exists so callers can + surface a not-found result. + """ + with DB.atomic(): + fresh = Document.select().where(Document.id == doc_id).for_update().first() + if fresh is None: + raise LookupError(doc_id) + if not (fresh.token_num or fresh.chunk_num or fresh.process_duration): + logging.debug("release_reparse_counters: nothing to release for document %s", doc_id) + return + DocumentService.increment_chunk_num(fresh.id, fresh.kb_id, -fresh.token_num, -fresh.chunk_num, -fresh.process_duration) + logging.debug( + "release_reparse_counters: released document %s (token=%s chunk=%s duration=%s)", + doc_id, + fresh.token_num, + fresh.chunk_num, + fresh.process_duration, + ) diff --git a/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py b/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py index 552288692c..4430a50793 100644 --- a/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py +++ b/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py @@ -14,6 +14,8 @@ # limitations under the License. # import asyncio +import contextlib +import functools import inspect import importlib.util import sys @@ -85,6 +87,7 @@ class _DummyDoc: doc_type=FileType.OTHER, status=True, run=0, + progress_msg="", ): self.id = doc_id self.kb_id = kb_id @@ -97,6 +100,7 @@ class _DummyDoc: self.type = doc_type self.status = status self.run = run + self.progress_msg = progress_msg def to_dict(self): return { @@ -137,7 +141,19 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"): monkeypatch.setitem(sys.modules, "common", common_pkg) apps_mod = ModuleType("api.apps") - apps_mod.login_required = lambda func: func + + def _login_required(func=None, **_kwargs): + # Real login_required is used both bare (chunk_api) and as a factory with + # auth_types=[...] (document_api); the mock must pass through both forms. + if func is None: + return lambda inner: inner + return func + + apps_mod.login_required = _login_required + apps_mod.current_user = SimpleNamespace(id="tenant-1") + apps_mod.AUTH_JWT = None + apps_mod.AUTH_API = None + apps_mod.AUTH_BETA = None monkeypatch.setitem(sys.modules, "api.apps", apps_mod) common_settings_mod = ModuleType("common.settings") @@ -152,6 +168,7 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"): return func(*args, **kwargs) common_misc_utils_mod.thread_pool_exec = _thread_pool_exec + common_misc_utils_mod.get_uuid = lambda: "uuid-1" monkeypatch.setitem(sys.modules, "common.misc_utils", common_misc_utils_mod) common_string_utils_mod = ModuleType("common.string_utils") @@ -180,21 +197,58 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"): def is_null(self, value=True): return _FakeExpr() + class _StubFreshDoc: + id = "doc-1" + kb_id = "kb-1" + token_num = 2 + chunk_num = 1 + process_duration = 0.0 + + class _StubDocQuery: + # Stands in for Document.select().where(...).for_update().first(). + def where(self, *_args, **_kwargs): + return self + + def for_update(self): + return self + + def first(self): + return _StubDocumentModel.fresh_doc + class _StubDocumentModel: id = _FakeField() run = _FakeField() + # The row re-read under lock by _release_doc_counters; tests that assert on + # the decrement override this to mirror the document's current counters. + fresh_doc = _StubFreshDoc() + + @classmethod + def select(cls, *_args, **_kwargs): + return _StubDocQuery() class _StubTaskModel: doc_id = _FakeField() + class _AnyFieldMeta(type): + def __getattr__(cls, _name): + return _FakeField() + + class _StubModel(metaclass=_AnyFieldMeta): + pass + db_models_mod = ModuleType("api.db.db_models") db_models_mod.APIToken = SimpleNamespace(query=lambda **_kwargs: []) db_models_mod.Document = _StubDocumentModel db_models_mod.Task = _StubTaskModel + db_models_mod.DB = SimpleNamespace(atomic=lambda: contextlib.nullcontext(), connection_context=lambda: lambda fn: fn) + # Transitively-loaded real services import assorted model classes (File, + # Knowledgebase, UserTenant, ...); hand them a permissive stub on demand. + db_models_mod.__getattr__ = lambda _name: _StubModel monkeypatch.setitem(sys.modules, "api.db.db_models", db_models_mod) services_pkg = ModuleType("api.db.services") services_pkg.__path__ = [str(repo_root / "api" / "db" / "services")] + services_pkg.duplicate_name = lambda _query, name="", **_kwargs: name monkeypatch.setitem(sys.modules, "api.db.services", services_pkg) doc_metadata_service_mod = ModuleType("api.db.services.doc_metadata_service") @@ -207,10 +261,12 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"): document_service_mod = ModuleType("api.db.services.document_service") document_service_mod.DocumentService = SimpleNamespace( query=lambda **_kwargs: [], + accessible=lambda *_args, **_kwargs: True, filter_update=lambda *_args, **_kwargs: 0, get_by_id=lambda *_args, **_kwargs: (False, None), update_by_id=lambda *_args, **_kwargs: True, decrement_chunk_num=lambda *_args, **_kwargs: None, + increment_chunk_num=lambda *_args, **_kwargs: True, get_embd_id=lambda *_args, **_kwargs: "", get_tenant_embd_id=lambda *_args, **_kwargs: None, ) @@ -233,13 +289,38 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"): monkeypatch.setitem(sys.modules, "api.db.services.knowledgebase_service", knowledgebase_service_mod) task_service_mod = ModuleType("api.db.services.task_service") - task_service_mod.TaskService = SimpleNamespace(filter_delete=lambda *_args, **_kwargs: None) + task_service_mod.TaskService = SimpleNamespace(filter_delete=lambda *_args, **_kwargs: None, query=lambda **_kwargs: []) task_service_mod.cancel_all_task_of = lambda *_args, **_kwargs: None task_service_mod.queue_tasks = lambda *_args, **_kwargs: None monkeypatch.setitem(sys.modules, "api.db.services.task_service", task_service_mod) + file_service_mod = ModuleType("api.db.services.file_service") + file_service_mod.FileService = SimpleNamespace() + monkeypatch.setitem(sys.modules, "api.db.services.file_service", file_service_mod) + + canvas_service_mod = ModuleType("api.db.services.canvas_service") + canvas_service_mod.UserCanvasService = SimpleNamespace() + monkeypatch.setitem(sys.modules, "api.db.services.canvas_service", canvas_service_mod) + + # document_api imports check_kb_team_permission; stub it so the real module + # and its user_service/common_service (peewee) import chain stay out. + check_team_permission_mod = ModuleType("api.common.check_team_permission") + check_team_permission_mod.check_kb_team_permission = lambda *_args, **_kwargs: True + check_team_permission_mod.check_file_team_permission = lambda *_args, **_kwargs: True + monkeypatch.setitem(sys.modules, "api.common.check_team_permission", check_team_permission_mod) + api_utils_mod = ModuleType("api.utils.api_utils") - api_utils_mod.add_tenant_id_to_kwargs = lambda func: func + + def _add_tenant_id_to_kwargs(func): + # Mirror the real decorator's functools.wraps so tests can reach the raw + # handler through ``route.__wrapped__`` and pass ``tenant_id`` explicitly. + @functools.wraps(func) + def wrapper(*args, **kwargs): + return func(*args, **kwargs) + + return wrapper + + api_utils_mod.add_tenant_id_to_kwargs = _add_tenant_id_to_kwargs api_utils_mod.check_duplicate_ids = lambda ids, _kind="item": (ids, []) api_utils_mod.construct_json_result = lambda code=0, message="success", data=None: {"code": code, "message": message, "data": data} api_utils_mod.get_error_data_result = lambda message="Sorry! Data missing!", code=102: {"code": code, "message": message} @@ -248,6 +329,9 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"): key: value for key, value in {"code": code, "message": message, "data": data, "total": total}.items() if value is not None } api_utils_mod.server_error_response = lambda e: {"code": 500, "message": str(e)} + api_utils_mod.get_data_error_result = lambda message="Sorry! Data missing!", code=102: {"code": code, "message": message} + api_utils_mod.get_error_argument_result = lambda message="": {"code": 101, "message": message} + api_utils_mod.get_json_result = lambda code=0, message="success", data=None, **_kwargs: {"code": code, "message": message, "data": data} monkeypatch.setitem(sys.modules, "api.utils.api_utils", api_utils_mod) image_utils_mod = ModuleType("api.utils.image_utils") @@ -278,6 +362,7 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"): common_metadata_utils_mod = ModuleType("common.metadata_utils") common_metadata_utils_mod.convert_conditions = lambda conditions: conditions common_metadata_utils_mod.meta_filter = lambda *_args, **_kwargs: [] + common_metadata_utils_mod.turn2jsonschema = lambda *_args, **_kwargs: {} monkeypatch.setitem(sys.modules, "common.metadata_utils", common_metadata_utils_mod) rag_app_tag_mod = ModuleType("rag.app.tag") @@ -470,10 +555,19 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"): # Return mock tenant with default model configurations return _MockModelConfig2(tenant_id, "chat-model").to_dict() + def _split_model_name(model_name): + parts = model_name.rsplit("@", 2) + if len(parts) == 3: + return parts[0], parts[1], parts[2] + if len(parts) == 2: + return parts[0], "default", parts[1] + return parts[0], "", "" + tenant_model_service_mod.get_model_config_by_id = _get_model_config_by_id tenant_model_service_mod.get_model_config_from_provider_instance = _get_model_config_from_provider_instance tenant_model_service_mod.resolve_model_config = _get_model_config_from_provider_instance tenant_model_service_mod.get_tenant_default_model_by_type = _get_tenant_default_model_by_type + tenant_model_service_mod.split_model_name = _split_model_name monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_service_mod) if module_basename == "document_api": @@ -493,6 +587,29 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"): document_api_service_mod.update_document_status_only = lambda *_args, **_kwargs: None document_api_service_mod.reset_document_for_reparse = lambda *_args, **_kwargs: None monkeypatch.setitem(sys.modules, "api.apps.services.document_api_service", document_api_service_mod) + else: + # chunk_api imports structure_graph_common from api.apps.services; stub + # it in sys.modules so the real module and its transitive imports are + # not loaded. + stub_apps_services = ModuleType("api.apps.services") + monkeypatch.setitem(sys.modules, "api.apps.services", stub_apps_services) + + sgc_mod = ModuleType("api.apps.services.structure_graph_common") + + async def _sgc_keyword_subgraph(*_args, **_kwargs): + return {}, [], [] + + async def _sgc_build_bucket(*_args, **_kwargs): + return [], [] + + sgc_mod.keyword_subgraph = _sgc_keyword_subgraph + sgc_mod.build_bucket = _sgc_build_bucket + monkeypatch.setitem(sys.modules, "api.apps.services.structure_graph_common", sgc_mod) + + # document_counter_service is a real module (release_reparse_counters); evict + # any cached copy so it re-imports against this test's freshly-stubbed + # db_models / document_service rather than a prior test's stubs. + monkeypatch.delitem(sys.modules, "api.db.services.document_counter_service", raising=False) module_path = repo_root / "api" / "apps" / "restful_apis" / f"{module_basename}.py" spec = importlib.util.spec_from_file_location("test_doc_sdk_routes_unit", module_path) @@ -552,80 +669,63 @@ class TestDocRoutesUnit: assert "length of 5" in str(exc_info.value) def test_download_and_download_doc_errors(self, monkeypatch): - module = _load_doc_module(monkeypatch) + module = _load_doc_module(monkeypatch, module_basename="document_api") _patch_send_file(monkeypatch, module) _patch_storage(monkeypatch, module, file_stream=b"") - res = _run(module.download.__wrapped__("tenant-1", "ds-1", "")) - assert res["message"] == "Specify document_id please." - monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: []) - res = _run(module.download.__wrapped__("tenant-1", "ds-1", "doc-1")) - assert "do not own the dataset" in res["message"] - monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [1]) + # download(dataset_id, document_id) + res = _run(module.download("ds-1", "")) + assert res["message"] == "Specify document_id please." + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False) + res = _run(module.download("ds-1", "doc-1")) + assert res["message"] == "Document not found!" + + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True) + monkeypatch.setattr(module.DocumentService, "accessible", lambda *_args, **_kwargs: False) + res = _run(module.download("ds-1", "doc-1")) + assert res["message"] == "Document not found!" + + monkeypatch.setattr(module.DocumentService, "accessible", lambda *_args, **_kwargs: True) monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: []) - res = _run(module.download.__wrapped__("tenant-1", "ds-1", "doc-1")) + res = _run(module.download("ds-1", "doc-1")) assert "not own the document" in res["message"] monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc()]) monkeypatch.setattr(module.File2DocumentService, "get_storage_address", lambda **_kwargs: ("b", "n")) - res = _run(module.download.__wrapped__("tenant-1", "ds-1", "doc-1")) + res = _run(module.download("ds-1", "doc-1")) assert res["message"] == "This file is empty." - monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer"})) - res = _run(module.download_doc("doc-1")) - assert "Authorization is not valid" in res["message"] - - monkeypatch.setattr(module, "request", SimpleNamespace(headers={"Authorization": "Bearer token"})) - monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: []) - res = _run(module.download_doc("doc-1")) - assert "API key is invalid" in res["message"] - - monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1"), SimpleNamespace(tenant_id="tenant-2")]) - res = _run(module.download_doc("doc-1")) - assert "API key configuration is ambiguous" in res["message"] - - monkeypatch.setattr(module.APIToken, "query", lambda **_kwargs: [SimpleNamespace(tenant_id="tenant-1")]) - res = _run(module.download_doc("")) + # download_document(document_id) + res = _run(module.download_document("")) assert res["message"] == "Specify document_id please." + monkeypatch.setattr(module.DocumentService, "accessible", lambda *_args, **_kwargs: False) + res = _run(module.download_document("doc-1")) + assert res["message"] == "Document not found!" + + monkeypatch.setattr(module.DocumentService, "accessible", lambda *_args, **_kwargs: True) monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: []) - res = _run(module.download_doc("doc-1")) + res = _run(module.download_document("doc-1")) assert "not own the document" in res["message"] monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc()]) - kb_query_calls = [] - - def _deny_kb_query(**kwargs): - kb_query_calls.append(kwargs) - return [] - - monkeypatch.setattr(module.KnowledgebaseService, "query", _deny_kb_query) - monkeypatch.setattr( - module.File2DocumentService, - "get_storage_address", - lambda **_kwargs: (_ for _ in ()).throw(AssertionError("storage lookup must not run before tenant authorization")), - ) - res = _run(module.download_doc("doc-1")) - assert res["message"] == "You do not have access to this document." - assert kb_query_calls == [{"id": "kb-1", "tenant_id": "tenant-1"}] - - monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [1]) - monkeypatch.setattr(module.File2DocumentService, "get_storage_address", lambda **_kwargs: ("b", "n")) - _patch_storage(monkeypatch, module, file_stream=b"") - res = _run(module.download_doc("doc-1")) + res = _run(module.download_document("doc-1")) assert res["message"] == "This file is empty." _patch_storage(monkeypatch, module, file_stream=b"abc") - res = _run(module.download_doc("doc-1")) + res = _run(module.download_document("doc-1")) assert res["filename"] == "doc.txt" def test_download_mimetype_from_filename(self, monkeypatch): module = _load_doc_module(monkeypatch, module_basename="document_api") _patch_send_file(monkeypatch, module) _patch_storage(monkeypatch, module, file_stream=b"pdf-bytes") + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True) + monkeypatch.setattr(module.DocumentService, "accessible", lambda *_args, **_kwargs: True) monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc(name="report.pdf", doc_type=FileType.PDF)]) monkeypatch.setattr(module.File2DocumentService, "get_storage_address", lambda **_kwargs: ("b", "n")) - res = _run(module.download.__wrapped__("ds-1", "doc-1")) + res = _run(module.download("ds-1", "doc-1")) assert res["filename"] == "report.pdf" assert res["mimetype"] == "application/pdf" @@ -636,6 +736,7 @@ class TestDocRoutesUnit: assert "don't own the dataset" in res["message"] monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, SimpleNamespace(tenant_id="tenant-1", pipeline_id=None))) monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"document_ids": ["doc-1"]})) monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, [])) toggle_doc = _ToggleBoolDocList(_DummyDoc(progress=0)) @@ -670,6 +771,58 @@ class TestDocRoutesUnit: assert res["code"] == module.RetCode.DATA_ERROR assert "Duplicate document ids" in res["message"] + def test_parse_and_stop_decrement_kb_counters(self, monkeypatch): + # Both routes delete the document's chunks, so the knowledgebase aggregate + # must drop by the document's current counters. Zeroing only the document + # row leaves that amount stranded in the KB total on every re-parse. + module = _load_doc_module(monkeypatch) + # _release_doc_counters re-reads the row under lock; mirror the document's + # current counters so the decrement is driven by that fresh read. + monkeypatch.setattr(module.Document, "fresh_doc", SimpleNamespace(id="doc-1", kb_id="kb-1", token_num=70, chunk_num=7, process_duration=1.5)) + monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, SimpleNamespace(tenant_id="tenant-1", pipeline_id=None))) + monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"document_ids": ["doc-1"]})) + monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, [])) + monkeypatch.setattr(module.DocumentService, "get_by_id", lambda _id: (True, _DummyDoc())) + monkeypatch.setattr(module.File2DocumentService, "get_storage_address", lambda **_kwargs: ("b", "n")) + monkeypatch.setattr(module.TaskService, "filter_delete", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "queue_tasks", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "cancel_all_task_of", lambda *_args, **_kwargs: None) + _patch_docstore(monkeypatch, module, delete=lambda *_args, **_kwargs: None) + + decrements = [] + updates = [] + update_by_id_payloads = [] + + def _capture_filter_update(_conditions, info): + updates.append(info) + return 1 + + monkeypatch.setattr(module.DocumentService, "update_by_id", lambda _id, info: update_by_id_payloads.append(info) or True) + + monkeypatch.setattr(module.DocumentService, "increment_chunk_num", lambda *args: decrements.append(args)) + monkeypatch.setattr(module.DocumentService, "filter_update", _capture_filter_update) + + monkeypatch.setattr(module.DocumentService, "query", lambda **_kwargs: [_DummyDoc(token_num=70, chunk_num=7, process_duration=1.5)]) + assert _run(module.parse.__wrapped__("tenant-1", "ds-1"))["code"] == 0 + assert decrements == [("doc-1", "kb-1", -70, -7, -1.5)] + # The document update must not zero the counters itself; that is exactly + # what strands the difference in the KB aggregate. + assert "chunk_num" not in updates[0] + assert "token_num" not in updates[0] + + decrements.clear() + monkeypatch.setattr( + module.DocumentService, + "query", + lambda **_kwargs: [_DummyDoc(run=module.TaskStatus.RUNNING.value, token_num=70, chunk_num=7, process_duration=1.5)], + ) + assert _run(module.stop_parsing.__wrapped__("tenant-1", "ds-1"))["code"] == 0 + assert decrements == [("doc-1", "kb-1", -70, -7, -1.5)] + # stop_parsing must not zero the counters in its own document update either. + assert "chunk_num" not in update_by_id_payloads[0] + assert "token_num" not in update_by_id_payloads[0] + def test_stop_parsing_branches(self, monkeypatch): module = _load_doc_module(monkeypatch) monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: False) @@ -677,6 +830,7 @@ class TestDocRoutesUnit: assert "don't own the dataset" in res["message"] monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True) + monkeypatch.setattr(module.KnowledgebaseService, "get_by_id", lambda _id: (True, SimpleNamespace(tenant_id="tenant-1", pipeline_id=None))) monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({})) res = _run(module.stop_parsing.__wrapped__("tenant-1", "ds-1")) assert "`document_ids` is required" in res["message"] @@ -733,7 +887,7 @@ class TestDocRoutesUnit: monkeypatch.setattr( module.KnowledgebaseService, "get_by_id", - lambda _id: (True, SimpleNamespace(tenant_id=owner_tenant)), + lambda _id: (True, SimpleNamespace(tenant_id=owner_tenant, pipeline_id=None)), ) monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"document_ids": ["doc-1"]})) monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, [])) @@ -769,7 +923,7 @@ class TestDocRoutesUnit: monkeypatch.setattr( module.KnowledgebaseService, "get_by_id", - lambda _id: (True, SimpleNamespace(tenant_id=owner_tenant)), + lambda _id: (True, SimpleNamespace(tenant_id=owner_tenant, pipeline_id=None)), ) monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"document_ids": ["doc-1"]})) monkeypatch.setattr(module, "check_duplicate_ids", lambda ids, _kind: (ids, [])) @@ -796,6 +950,7 @@ class TestDocRoutesUnit: module = _load_doc_module(monkeypatch, module_basename="document_api") updated = [] deleted = [] + decrements = [] monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: True) monkeypatch.setattr(module, "get_request_json", lambda: _AwaitableValue({"document_ids": ["doc-1"]})) @@ -808,6 +963,14 @@ class TestDocRoutesUnit: ) monkeypatch.setattr(module.TaskService, "query", lambda **_kwargs: [SimpleNamespace(progress=0.5)]) monkeypatch.setattr(module, "cancel_all_task_of", lambda *_args, **_kwargs: None) + # release_reparse_counters re-reads the row under lock and decrements the + # KB by the document's partial counts before chunk_num is zeroed below. + monkeypatch.setattr( + sys.modules["api.db.db_models"].Document, + "fresh_doc", + SimpleNamespace(id="doc-1", kb_id="kb-1", token_num=70, chunk_num=7, process_duration=1.5), + ) + monkeypatch.setattr(module.DocumentService, "increment_chunk_num", lambda *args: decrements.append(args)) monkeypatch.setattr( module.DocumentService, "update_by_id", @@ -824,16 +987,19 @@ class TestDocRoutesUnit: assert res["code"] == 0 assert res["data"]["success_count"] == 1 - assert updated == [ - ( - "doc-1", - { - "run": module.TaskStatus.CANCEL.value, - "progress": 0, - "chunk_num": 0, - }, - ) - ] + assert len(updated) == 1 + updated_doc_id, updated_info = updated[0] + assert updated_doc_id == "doc-1" + assert updated_info["run"] == str(module.TaskStatus.CANCEL.value) + assert updated_info["progress"] == 0 + # release_reparse_counters is the sole counter adjustment; the status + # update must not set chunk_num itself (that would strand KB counts). + assert "chunk_num" not in updated_info + # progress_msg carries a timestamped cancellation marker, so match loosely. + assert "Task stopped by user." in updated_info["progress_msg"] + # The partial chunk/token counts are released from the KB aggregate via the + # row re-read under lock. + assert decrements == [("doc-1", "kb-1", -70, -7, -1.5)] assert deleted == [({"doc_id": "doc-1"}, module.search.index_name("tenant-1"), "kb-1")] deleted.clear() diff --git a/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py b/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py index 3708c76d43..194bbb6259 100644 --- a/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py +++ b/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py @@ -15,6 +15,7 @@ # import asyncio +import contextlib import inspect import importlib.util import sys @@ -223,6 +224,14 @@ def _load_chunk_module(monkeypatch): constants_mod.LLMType = _DummyLLMType constants_mod.ParserType = _DummyParserType constants_mod.PAGERANK_FLD = "pagerank_flt" + constants_mod.TaskStatus = SimpleNamespace( + UNSTART=SimpleNamespace(value="0"), + RUNNING=SimpleNamespace(value="1"), + CANCEL=SimpleNamespace(value="2"), + DONE=SimpleNamespace(value="3"), + FAIL=SimpleNamespace(value="4"), + SCHEDULE=SimpleNamespace(value="5"), + ) monkeypatch.setitem(sys.modules, "common.constants", constants_mod) string_utils_mod = ModuleType("common.string_utils") @@ -232,8 +241,28 @@ def _load_chunk_module(monkeypatch): metadata_utils_mod = ModuleType("common.metadata_utils") metadata_utils_mod.apply_meta_data_filter = lambda *_args, **_kwargs: {} + metadata_utils_mod.convert_conditions = lambda *_args, **_kwargs: {} + metadata_utils_mod.meta_filter = lambda *_args, **_kwargs: {} monkeypatch.setitem(sys.modules, "common.metadata_utils", metadata_utils_mod) + doc_store_base_mod = ModuleType("common.doc_store.doc_store_base") + doc_store_base_mod.OrderByExpr = type("OrderByExpr", (), {}) + monkeypatch.setitem(sys.modules, "common.doc_store", ModuleType("common.doc_store")) + monkeypatch.setitem(sys.modules, "common.doc_store.doc_store_base", doc_store_base_mod) + + tag_feature_utils_mod = ModuleType("common.tag_feature_utils") + tag_feature_utils_mod.validate_tag_features = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "common.tag_feature_utils", tag_feature_utils_mod) + + pagination_utils_mod = ModuleType("api.utils.pagination_utils") + pagination_utils_mod.validate_rest_api_page_size = lambda *_args, **_kwargs: (1, 30) + monkeypatch.setitem(sys.modules, "api.utils.pagination_utils", pagination_utils_mod) + + reference_metadata_utils_mod = ModuleType("api.utils.reference_metadata_utils") + reference_metadata_utils_mod.enrich_chunks_with_document_metadata = lambda chunks, *_args, **_kwargs: chunks + reference_metadata_utils_mod.resolve_reference_metadata_preferences = lambda *_args, **_kwargs: {} + monkeypatch.setitem(sys.modules, "api.utils.reference_metadata_utils", reference_metadata_utils_mod) + misc_utils_mod = ModuleType("common.misc_utils") async def _thread_pool_exec(func): @@ -293,6 +322,7 @@ def _load_chunk_module(monkeypatch): 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({}) + api_utils_mod.construct_json_result = lambda code=0, message="success", data=None: {"code": code, "message": message, "data": data} monkeypatch.setitem(sys.modules, "api.utils.api_utils", api_utils_mod) image_utils_mod = ModuleType("api.utils.image_utils") @@ -312,8 +342,94 @@ def _load_chunk_module(monkeypatch): tenant_model_service_mod.get_model_config_from_provider_instance = lambda *_args, **_kwargs: {"llm_name": "embed", "model_type": "embedding"} tenant_model_service_mod.resolve_model_config = 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"} + tenant_model_service_mod.split_model_name = lambda model_name: (model_name.rsplit("@", 2) + ["", ""])[:3] monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_service_mod) + # chunk_api imports structure_graph_common from api.apps.services; stub it in + # sys.modules so the real module and its transitive imports are not loaded. + stub_apps_services = ModuleType("api.apps.services") + monkeypatch.setitem(sys.modules, "api.apps.services", stub_apps_services) + + sgc_mod = ModuleType("api.apps.services.structure_graph_common") + + async def _sgc_keyword_subgraph(*_args, **_kwargs): + return {}, [], [] + + async def _sgc_build_bucket(*_args, **_kwargs): + return [], [] + + sgc_mod.keyword_subgraph = _sgc_keyword_subgraph + sgc_mod.build_bucket = _sgc_build_bucket + monkeypatch.setitem(sys.modules, "api.apps.services.structure_graph_common", sgc_mod) + + # chunk_api imports DB, Document, Task from db_models; stub them so the real + # module (which pulls in quart_auth) is never loaded. + class _FakeExpr: + def __or__(self, other): + return self + + def __and__(self, other): + return self + + class _FakeField: + def __eq__(self, other): + return _FakeExpr() + + def __ne__(self, other): + return _FakeExpr() + + def is_null(self, value=True): + return _FakeExpr() + + class _StubFreshDoc: + id = "doc-1" + kb_id = "kb-1" + token_num = 2 + chunk_num = 1 + process_duration = 0.0 + + class _StubDocQuery: + def where(self, *_args, **_kwargs): + return self + + def for_update(self): + return self + + def first(self): + return _StubDocumentModel.fresh_doc + + class _StubDocumentModel: + id = _FakeField() + run = _FakeField() + fresh_doc = _StubFreshDoc() + + @classmethod + def select(cls, *_args, **_kwargs): + return _StubDocQuery() + + class _StubTaskModel: + doc_id = _FakeField() + + db_models_mod = ModuleType("api.db.db_models") + db_models_mod.Document = _StubDocumentModel + db_models_mod.Task = _StubTaskModel + db_models_mod.DB = SimpleNamespace(atomic=lambda: contextlib.nullcontext()) + monkeypatch.setitem(sys.modules, "api.db.db_models", db_models_mod) + + file2document_service_mod = ModuleType("api.db.services.file2document_service") + file2document_service_mod.File2DocumentService = SimpleNamespace(get_storage_address=lambda **_kwargs: ("", "")) + monkeypatch.setitem(sys.modules, "api.db.services.file2document_service", file2document_service_mod) + + task_service_mod = ModuleType("api.db.services.task_service") + task_service_mod.TaskService = SimpleNamespace(filter_delete=lambda *_args, **_kwargs: None) + task_service_mod.cancel_all_task_of = lambda *_args, **_kwargs: None + task_service_mod.queue_tasks = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "api.db.services.task_service", task_service_mod) + + document_counter_service_mod = ModuleType("api.db.services.document_counter_service") + document_counter_service_mod.release_reparse_counters = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "api.db.services.document_counter_service", document_counter_service_mod) + document_service_mod = ModuleType("api.db.services.document_service") class _DocumentService: @@ -499,6 +615,14 @@ def _load_chunk_api_module(monkeypatch): module.manager = _DummyManager() monkeypatch.setitem(sys.modules, module_name, module) spec.loader.exec_module(module) + # chunk_api imports these inside the chunk-write helpers; re-expose the shared + # stubs so tests can reach them via module. (setattr / method patching). + module.rag_tokenizer = sys.modules["rag.nlp"].rag_tokenizer + module.beAdoc = sys.modules["rag.app.qa"].beAdoc + module.rmPrefix = sys.modules["rag.app.qa"].rmPrefix + module.label_question = sys.modules["rag.app.tag"].label_question + module.cross_languages = sys.modules["rag.prompts.generator"].cross_languages + module.keyword_extraction = sys.modules["rag.prompts.generator"].keyword_extraction return module