fix(api): decrement knowledgebase counters on SDK re-parse / stop-parse (#17236)

This commit is contained in:
deadtrickster
2026-08-05 04:28:46 +02:00
committed by GitHub
parent 61d2747f6a
commit f063cfdb19
6 changed files with 446 additions and 81 deletions

View File

@@ -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.<name> (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