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

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/<dataset_id>/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)

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

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