Refactor: Task Executor (#15154)

### What problem does this PR solve?

1. Break huge function into smaller pieces
2. Add unit test for the smaller pieces function
3. Layer-ed design
a. infra layer - task_context.py, recording_context.py,
write_operation_interceptor.py, ...
    b. service layer - *_service.py
    c. business layer - task_handler.py
4. Default behavior: use "refactor-ed version" - can switch to original
version by change env variable

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
- [x] Performance Improvement

---------

Co-authored-by: Liu An <asiro@qq.com>
Co-authored-by: Zhichang Yu <yuzhichang@gmail.com>

rag/flow/pipeline.py

@@ -49,7 +49,7 @@ class Pipeline(Graph):
             message += "[CANCEL]"
         try:
             bin = REDIS_CONN.get(log_key)
-            obj = json.loads(bin.encode("utf-8"))
+            obj = json.loads(bin.encode("utf-8")) if bin else []
             if obj:
                 if obj[-1]["component_id"] == component_name:
                     obj[-1]["trace"].append(

rag/flow/tokenizer/tokenizer.py

@@ -26,9 +26,9 @@ from common.connection_utils import timeout
 from rag.flow.base import ProcessBase, ProcessParamBase
 from rag.flow.parser.pdf_chunk_metadata import finalize_pdf_chunk
 from rag.flow.tokenizer.schema import TokenizerFromUpstream
+from rag.svr.task_executor_limiter import embed_limiter
 from rag.nlp import rag_tokenizer
 from common import settings
-from rag.svr.task_executor import embed_limiter
 from common.token_utils import truncate
 
 from common.misc_utils import thread_pool_exec

rag/prompts/assign_toc_levels.md

@@ -8,9 +8,11 @@ Task
 - Decide levels yourself to keep a coherent hierarchy. Keep peers at the same depth.
 
 Output
-- Return a valid JSON array only (no extra text).
-- Each element must be {"level": "1|2|3", "title": <original title string>}.
-- title must be the original title string.
+- Return a valid JSON array only (no extra text, no markdown code blocks).
+- Each element MUST be a JSON object with exactly this structure: {"level": "1", "title": "some title"}.
+- title must be the original title string exactly.
+- DO NOT return arrays of arrays like [["1", "title"]] or other formats.
+- The output must be parseable by json.loads() directly.
 
 Examples
 

rag/prompts/generator.py

@@ -887,6 +887,23 @@ async def run_toc_from_text(chunks, chat_mdl, callback=None):
     if not toc_with_levels:
         return []
 
+    # Normalize TOC items to ensure consistent dict format
+    normalized_levels = []
+    for item in toc_with_levels:
+        if isinstance(item, dict):
+            # Already in correct format
+            normalized_levels.append(item)
+        elif isinstance(item, (list, tuple)) and len(item) >= 2:
+            # Convert ["level", "title"] or similar to dict
+            normalized_levels.append({"level": str(item[0]), "title": str(item[1])})
+        else:
+            logging.warning(f"Unexpected TOC item format (type={type(item).__name__}), skipping: {item}")
+
+    toc_with_levels = normalized_levels
+    if not toc_with_levels:
+        logging.warning("No valid TOC items after normalization.")
+        return []
+
     # Merge structure and content (by index)
     prune = len(toc_with_levels) > 512
     max_lvl = "0"

rag/svr/task_executor.py

@@ -12,9 +12,13 @@
 #  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 argparse
 import time
 
+from rag.svr.task_executor_refactor.task_manager import TaskManager
+from rag.svr.task_executor_refactor.recording_context import timed_with_recording, get_recording_context, \
+    RecordingContext, set_recording_context, NullRecordingContext
+
 start_ts = time.time()
 
 # LiteLLM fetches a model cost map from GitHub during import unless this is set.
@@ -89,7 +93,13 @@ from rag.utils.redis_conn import REDIS_CONN, RedisDistributedLock
 from rag.graphrag.utils import chat_limiter
 from common.signal_utils import start_tracemalloc_and_snapshot, stop_tracemalloc
 from common.exceptions import TaskCanceledException
-from common.asyncio_utils import LoopLocalSemaphore
+from rag.svr.task_executor_limiter import (
+    task_limiter,
+    chunk_limiter,
+    embed_limiter,
+    minio_limiter,
+    kg_limiter,
+)
 from common import settings
 from common.constants import PAGERANK_FLD, TAG_FLD, SVR_CONSUMER_GROUP_NAME
 from rag.utils.table_es_metadata import (
@@ -97,6 +107,7 @@ from rag.utils.table_es_metadata import (
     merge_table_parser_config_from_kb,
     table_parser_strip_doc_metadata_keys,
 )
+from rag.nlp import search as nlp_search
 
 BATCH_SIZE = 64
 
@@ -129,9 +140,10 @@ TASK_TYPE_TO_PIPELINE_TASK_TYPE = {
 }
 
 UNACKED_ITERATOR = None
+# Task type and executor index (consistent with SAAS version)
+TASK_TYPE = "common"
+TE_IDX = "0"
 
-CONSUMER_NO = "0" if len(sys.argv) < 2 else sys.argv[1]
-CONSUMER_NAME = "task_executor_" + CONSUMER_NO
 BOOT_AT = datetime.now().astimezone().isoformat(timespec="milliseconds")
 PENDING_TASKS = 0
 LAG_TASKS = 0
@@ -140,18 +152,9 @@ FAILED_TASKS = 0
 
 CURRENT_TASKS = {}
 
-MAX_CONCURRENT_TASKS = int(os.environ.get('MAX_CONCURRENT_TASKS', "5"))
-MAX_CONCURRENT_CHUNK_BUILDERS = int(os.environ.get('MAX_CONCURRENT_CHUNK_BUILDERS', "1"))
-MAX_CONCURRENT_MINIO = int(os.environ.get('MAX_CONCURRENT_MINIO', '10'))
-task_limiter = LoopLocalSemaphore(MAX_CONCURRENT_TASKS)
-chunk_limiter = LoopLocalSemaphore(MAX_CONCURRENT_CHUNK_BUILDERS)
-embed_limiter = LoopLocalSemaphore(MAX_CONCURRENT_CHUNK_BUILDERS)
-minio_limiter = LoopLocalSemaphore(MAX_CONCURRENT_MINIO)
-kg_limiter = LoopLocalSemaphore(2)
 WORKER_HEARTBEAT_TIMEOUT = int(os.environ.get('WORKER_HEARTBEAT_TIMEOUT', '120'))
 stop_event = threading.Event()
 
-
 def signal_handler(sig, frame):
     logging.info("Received interrupt signal, shutting down...")
     stop_event.set()
@@ -197,7 +200,8 @@ async def collect():
     global CONSUMER_NAME, DONE_TASKS, FAILED_TASKS
     global UNACKED_ITERATOR
 
-    svr_queue_names = settings.get_svr_queue_names()
+    svr_queue_names = settings.get_svr_queue_names(TASK_TYPE)
+
     redis_msg = None
     try:
         if not UNACKED_ITERATOR:
@@ -261,12 +265,16 @@ async def get_storage_binary(bucket, name):
     return await thread_pool_exec(settings.STORAGE_IMPL.get, bucket, name)
 
 
+@timed_with_recording
 @timeout(60 * 80, 1)
 async def build_chunks(task, progress_callback):
     if task["size"] > settings.DOC_MAXIMUM_SIZE:
         set_progress(task["id"], prog=-1, msg="File size exceeds( <= %dMb )" %
                                               (int(settings.DOC_MAXIMUM_SIZE / 1024 / 1024)))
+        get_recording_context().record("file_size_exceeded", True)
         return []
+    get_recording_context().record("file_size_exceeded", False)
+    get_recording_context().record("parser_id", task["parser_id"])
 
     chunker = FACTORY[task["parser_id"].lower()]
     try:
@@ -299,6 +307,23 @@ async def build_chunks(task, progress_callback):
             f"roles_keys={list((parser_config_for_chunk.get('table_column_roles') or {}).keys())}"
         )
 
+    # Record chunk configuration for comparison
+    from common.float_utils import normalize_overlapped_percent
+    chunk_config = {
+        "parser_id": task["parser_id"],
+        "chunk_token_num": parser_config_for_chunk.get("chunk_token_num", 128),
+        "overlapped_percent": normalize_overlapped_percent(
+            parser_config_for_chunk.get("overlapped_percent", 0)
+        ),
+        "delimiter": parser_config_for_chunk.get("delimiter", "\n!?。；！？"),
+        "from_page": task["from_page"],
+        "to_page": task["to_page"],
+        "language": task["language"],
+        "layout_recognizer": parser_config_for_chunk.get("layout_recognizer"),
+    }
+    get_recording_context().record("chunk_config", chunk_config)
+    get_recording_context().record("parser_config_after_merge", parser_config_for_chunk)
+
     try:
         async with chunk_limiter:
             task_language = task.get("language") or "Chinese"
@@ -322,15 +347,22 @@ async def build_chunks(task, progress_callback):
         logging.exception("Chunking {}/{} got exception".format(task["location"], task["name"]))
         raise
 
+    # Record raw chunks for comparison
+    get_recording_context().record("raw_chunks", cks)
+
     # Extract and persist PDF outline if the parser attached it.
+    outline_data = cks[0].get("__outline__") if cks else None
+    get_recording_context().record("outline_data", outline_data)
+
     if cks and cks[0].get("__outline__"):
         outline = cks[0].pop("__outline__")
         try:
-            DocMetadataService.update_document_metadata(
+            ret = DocMetadataService.update_document_metadata(
                 task["doc_id"],
                 update_metadata_to({"outline": outline},
                                    DocMetadataService.get_document_metadata(task["doc_id"]) or {})
             )
+            get_recording_context().save_func_return_value("DocMetadataService.update_document_metadata", ret)
             logging.info("Persisted PDF outline (%d entries) for doc %s", len(outline), task["doc_id"])
         except Exception as e:
             logging.warning("Failed to persist PDF outline for doc %s: %s", task["doc_id"], e)
@@ -385,6 +417,9 @@ async def build_chunks(task, progress_callback):
     el = timer() - st
     logging.info("MINIO PUT({}) cost {:.3f} s".format(task["name"], el))
 
+    # Record docs after MinIO upload
+    get_recording_context().record("docs_after_prep", docs)
+
     if task["parser_config"].get("auto_keywords", 0):
         st = timer()
         progress_callback(msg="Start to generate keywords for every chunk ...")
@@ -419,6 +454,10 @@ async def build_chunks(task, progress_callback):
             raise
         progress_callback(msg="Keywords generation {} chunks completed in {:.2f}s".format(len(docs), timer() - st))
 
+    # Record keywords extraction count
+    keywords = [d for d in docs if d.get("important_kwd")]
+    get_recording_context().record("keywords_extracted", keywords)
+
     if task["parser_config"].get("auto_questions", 0):
         st = timer()
         progress_callback(msg="Start to generate questions for every chunk ...")
@@ -452,6 +491,10 @@ async def build_chunks(task, progress_callback):
             raise
         progress_callback(msg="Question generation {} chunks completed in {:.2f}s".format(len(docs), timer() - st))
 
+    # Record question generation
+    questions = [d for d in docs if d.get("question_kwd")]
+    get_recording_context().record("questions_generated", questions)
+
     if task["parser_config"].get("enable_metadata", False) and (task["parser_config"].get("metadata") or task["parser_config"].get("built_in_metadata")):
         st = timer()
         progress_callback(msg="Start to generate meta-data for every chunk ...")
@@ -510,9 +553,14 @@ async def build_chunks(task, progress_callback):
             existing_meta = DocMetadataService.get_document_metadata(task["doc_id"])
             existing_meta = existing_meta if isinstance(existing_meta, dict) else {}
             metadata = update_metadata_to(metadata, existing_meta)
-            DocMetadataService.update_document_metadata(task["doc_id"], metadata)
+            ret = DocMetadataService.update_document_metadata(task["doc_id"], metadata)
+            get_recording_context().save_func_return_value("DocMetadataService.update_document_metadata", ret)
         progress_callback(msg="Question generation {} chunks completed in {:.2f}s".format(len(docs), timer() - st))
 
+    # Record metadata generation count
+    metadata_list = [d for d in docs if d.get("metadata_obj")]
+    get_recording_context().record("metadata_list_generated", metadata_list)
+
     if task["kb_parser_config"].get("tag_kb_ids", []):
         progress_callback(msg="Start to tag for every chunk ...")
         kb_ids = task["kb_parser_config"]["tag_kb_ids"]
@@ -578,9 +626,19 @@ async def build_chunks(task, progress_callback):
             raise
         progress_callback(msg="Tagging {} chunks completed in {:.2f}s".format(len(docs), timer() - st))
 
+    # Record tags applied
+    tags_applied = [d for d in docs if d.get(TAG_FLD)]
+    get_recording_context().record("tags_applied", tags_applied)
+
+    # Record final chunks for comparison
+    get_recording_context().record("final_chunks", docs)
+    final_chunk_ids = [c.get("id") for c in docs if isinstance(c, dict) and "id" in c]
+    get_recording_context().record("final_chunk_ids_count", len(final_chunk_ids))
+
     return docs
 
 
+@timed_with_recording
 def build_TOC(task, docs, progress_callback):
     progress_callback(msg="Start to generate table of content ...")
     chat_model_config = get_model_config_by_type_and_name(task["tenant_id"], LLMType.CHAT, task["llm_id"])
@@ -634,6 +692,7 @@ def init_kb(row, vector_size: int):
     return settings.docStoreConn.create_idx(idxnm, row.get("kb_id", ""), vector_size, parser_id)
 
 
+@timed_with_recording
 async def embedding(docs, mdl, parser_config=None, callback=None):
     if parser_config is None:
         parser_config = {}
@@ -686,6 +745,7 @@ async def embedding(docs, mdl, parser_config=None, callback=None):
     return tk_count, vector_size
 
 
+@timed_with_recording
 async def run_dataflow(task: dict):
     from api.db.services.canvas_service import UserCanvasService
     from rag.flow.pipeline import Pipeline
@@ -708,32 +768,47 @@ async def run_dataflow(task: dict):
     pipeline = Pipeline(dsl, tenant_id=task["tenant_id"], doc_id=doc_id, task_id=task_id, flow_id=dataflow_id)
     chunks = await pipeline.run(file=task["file"]) if task.get("file") else await pipeline.run()
     if doc_id == CANVAS_DEBUG_DOC_ID:
+        get_recording_context().record("dataflow_debug_result", "canvas_debug_mode")
+        get_recording_context().record("dataflow_chunks", chunks)
         return
 
     if not chunks:
-        PipelineOperationLogService.create(document_id=doc_id, pipeline_id=dataflow_id,
+        get_recording_context().record("pipeline_output_count", 0)
+        get_recording_context().record("pipeline_output_type", "empty")
+        ret = PipelineOperationLogService.create(document_id=doc_id, pipeline_id=dataflow_id,
                                            task_type=PipelineTaskType.PARSE, dsl=str(pipeline))
+        get_recording_context().save_func_return_value("PipelineOperationLogService.create", ret)
         return
 
     embedding_token_consumption = chunks.get("embedding_token_consumption", 0)
     # The output key may exist with an empty payload; check presence, not truthiness.
     if "chunks" in chunks:
         chunks = copy.deepcopy(chunks["chunks"])
+        output_type = "chunks"
     elif "json" in chunks:
         chunks = copy.deepcopy(chunks["json"])
+        output_type = "json"
     elif "markdown" in chunks:
         chunks = [{"text": [chunks["markdown"]]}] if chunks["markdown"] else []
+        output_type = "markdown"
     elif "text" in chunks:
         chunks = [{"text": [chunks["text"]]}] if chunks["text"] else []
+        output_type = "text"
     elif "html" in chunks:
         chunks = [{"text": [chunks["html"]]}] if chunks["html"] else []
+        output_type = "html"
     else:
         chunks = []
+        output_type = "empty"
+
+    get_recording_context().record("pipeline_output_type", output_type)
+    get_recording_context().record("pipeline_output_count", len(chunks))
 
     # An empty normalized payload means "nothing parsed", so stop before embedding/indexing.
     if not chunks:
-        PipelineOperationLogService.create(document_id=doc_id, pipeline_id=dataflow_id,
+        ret = PipelineOperationLogService.create(document_id=doc_id, pipeline_id=dataflow_id,
                                            task_type=PipelineTaskType.PARSE, dsl=str(pipeline))
+        get_recording_context().save_func_return_value("PipelineOperationLogService.create", ret)
         return
 
     keys = [k for o in chunks for k in list(o.keys())]
@@ -763,6 +838,8 @@ async def run_dataflow(task: dict):
                 if i % (len(texts) // settings.EMBEDDING_BATCH_SIZE / 100 + 1) == 1:
                     set_progress(task_id, prog=prog, msg=f"{i + 1} / {len(texts) // settings.EMBEDDING_BATCH_SIZE}")
             vects = np.vstack(vects_batches) if vects_batches else np.array([])
+            get_recording_context().record("embedding_token_consumption", embedding_token_consumption)
+            get_recording_context().record("vector_size", len(vects[0]) if len(vects) > 0 else 0)
 
             assert len(vects) == len(chunks)
             for i, ck in enumerate(chunks):
@@ -772,8 +849,9 @@ async def run_dataflow(task: dict):
             raise
         except Exception as e:
             set_progress(task_id, prog=-1, msg=f"[ERROR]: {e}")
-            PipelineOperationLogService.create(document_id=doc_id, pipeline_id=dataflow_id,
+            ret = PipelineOperationLogService.create(document_id=doc_id, pipeline_id=dataflow_id,
                                                task_type=PipelineTaskType.PARSE, dsl=str(pipeline))
+            get_recording_context().save_func_return_value("PipelineOperationLogService.create", ret)
             return
 
     metadata = {}
@@ -814,26 +892,31 @@ async def run_dataflow(task: dict):
         existing_meta = DocMetadataService.get_document_metadata(doc_id)
         existing_meta = existing_meta if isinstance(existing_meta, dict) else {}
         metadata = update_metadata_to(metadata, existing_meta)
-        DocMetadataService.update_document_metadata(doc_id, metadata)
+        get_recording_context().record("run_dataflow_metadata", metadata)
+        ret = DocMetadataService.update_document_metadata(doc_id, metadata)
+        get_recording_context().save_func_return_value("DocMetadataService.update_document_metadata", ret)
 
     start_ts = timer()
     set_progress(task_id, prog=0.82, msg="[DOC Engine]:\nStart to index...")
     e = await insert_chunks(task_id, task["tenant_id"], task["kb_id"], chunks, partial(set_progress, task_id, 0, 100000000))
     if not e:
-        PipelineOperationLogService.create(document_id=doc_id, pipeline_id=dataflow_id,
+        ret = PipelineOperationLogService.create(document_id=doc_id, pipeline_id=dataflow_id,
                                            task_type=PipelineTaskType.PARSE, dsl=str(pipeline))
+        get_recording_context().save_func_return_value("PipelineOperationLogService.create", ret)
         return
 
     time_cost = timer() - start_ts
     task_time_cost = timer() - task_start_ts
     set_progress(task_id, prog=1., msg="Indexing done ({:.2f}s). Task done ({:.2f}s)".format(time_cost, task_time_cost))
-    DocumentService.increment_chunk_num(doc_id, task_dataset_id, embedding_token_consumption, len(chunks),
+    ret = DocumentService.increment_chunk_num(doc_id, task_dataset_id, embedding_token_consumption, len(chunks),
                                         task_time_cost)
+    get_recording_context().save_func_return_value("DocumentService.increment_chunk_num", ret)
     logging.info("[Done], chunks({}), token({}), elapsed:{:.2f}".format(len(chunks), embedding_token_consumption,
                                                                         task_time_cost))
-    PipelineOperationLogService.create(document_id=doc_id, pipeline_id=dataflow_id, task_type=PipelineTaskType.PARSE,
+    get_recording_context().record("dataflow_chunks", chunks)
+    ret = PipelineOperationLogService.create(document_id=doc_id, pipeline_id=dataflow_id, task_type=PipelineTaskType.PARSE,
                                        dsl=str(pipeline))
-
+    get_recording_context().save_func_return_value("PipelineOperationLogService.create", ret)
 
 RAPTOR_METHOD_SEARCH_LIMIT = 10000
 
@@ -901,19 +984,18 @@ async def has_raptor_chunks(doc_id: str, tenant_id: str, kb_id: str, tree_builde
 
 async def delete_raptor_chunks(doc_id: str, tenant_id: str, kb_id: str, keep_method: str | None = None):
     """Delete RAPTOR summaries for doc_id, optionally preserving one method."""
-    from rag.nlp import search as nlp_search
-
     if keep_method is None:
         logging.info(
             "delete_raptor_chunks: removing all RAPTOR summaries (doc=%s tenant=%s kb=%s)",
             doc_id, tenant_id, kb_id,
         )
-        await thread_pool_exec(
+        ret = await thread_pool_exec(
             settings.docStoreConn.delete,
             {"doc_id": doc_id, "raptor_kwd": ["raptor"]},
             nlp_search.index_name(tenant_id),
             kb_id,
         )
+        get_recording_context().save_func_return_value("docStoreConn.delete", ret)
         return 0
 
     field_map = await get_raptor_chunk_field_map(doc_id, tenant_id, kb_id)
@@ -929,12 +1011,13 @@ async def delete_raptor_chunks(doc_id: str, tenant_id: str, kb_id: str, keep_met
         "delete_raptor_chunks: removing %d stale RAPTOR chunks (doc=%s tenant=%s kb=%s keep=%s)",
         len(chunk_ids), doc_id, tenant_id, kb_id, keep_method,
     )
-    await thread_pool_exec(
+    ret = await thread_pool_exec(
         settings.docStoreConn.delete,
         {"id": list(chunk_ids)},
         nlp_search.index_name(tenant_id),
         kb_id,
     )
+    get_recording_context().save_func_return_value("docStoreConn.delete", ret)
     return len(chunk_ids)
 
 
@@ -1171,6 +1254,7 @@ async def delete_image(kb_id, chunk_id):
         raise
 
 
+@timed_with_recording
 async def insert_chunks(task_id, task_tenant_id, task_dataset_id, chunks, progress_callback):
     """
     Insert chunks into document store (Elasticsearch OR Infinity).
@@ -1205,8 +1289,9 @@ async def insert_chunks(task_id, task_tenant_id, task_dataset_id, chunks, progre
         mothers.append(mom_ck)
 
     for b in range(0, len(mothers), settings.DOC_BULK_SIZE):
-        await thread_pool_exec(settings.docStoreConn.insert, mothers[b:b + settings.DOC_BULK_SIZE],
+        ret = await thread_pool_exec(settings.docStoreConn.insert, mothers[b:b + settings.DOC_BULK_SIZE],
                                 search.index_name(task_tenant_id), task_dataset_id, )
+        get_recording_context().save_func_return_value("docStoreConn.insert", ret)
         task_canceled = has_canceled(task_id)
         if task_canceled:
             progress_callback(-1, msg="Task has been canceled.")
@@ -1215,6 +1300,7 @@ async def insert_chunks(task_id, task_tenant_id, task_dataset_id, chunks, progre
     for b in range(0, len(chunks), settings.DOC_BULK_SIZE):
         doc_store_result = await thread_pool_exec(settings.docStoreConn.insert, chunks[b:b + settings.DOC_BULK_SIZE],
                                                    search.index_name(task_tenant_id), task_dataset_id, )
+        get_recording_context().save_func_return_value("docStoreConn.insert", doc_store_result)
         task_canceled = has_canceled(task_id)
         if task_canceled:
             # Roll back partial RAPTOR summary inserts so the next run is not
@@ -1225,12 +1311,13 @@ async def insert_chunks(task_id, task_tenant_id, task_dataset_id, chunks, progre
             ]
             if raptor_ids_to_rollback:
                 try:
-                    await thread_pool_exec(
+                    ret = await thread_pool_exec(
                         settings.docStoreConn.delete,
                         {"id": raptor_ids_to_rollback},
                         search.index_name(task_tenant_id),
                         task_dataset_id,
                     )
+                    get_recording_context().save_func_return_value("docStoreConn.delete", ret)
                     logging.info(
                         "insert_chunks: rolled back %d partial RAPTOR chunks after cancellation (task=%s)",
                         len(raptor_ids_to_rollback), task_id,
@@ -1252,10 +1339,12 @@ async def insert_chunks(task_id, task_tenant_id, task_dataset_id, chunks, progre
         chunk_ids_str = " ".join(chunk_ids)
         try:
             TaskService.update_chunk_ids(task_id, chunk_ids_str)
+            get_recording_context().save_func_return_value("TaskService.update_chunk_ids", None)
         except DoesNotExist:
             logging.warning(f"do_handle_task update_chunk_ids failed since task {task_id} is unknown.")
             doc_store_result = await thread_pool_exec(settings.docStoreConn.delete, {"id": chunk_ids},
                                                        search.index_name(task_tenant_id), task_dataset_id, )
+            get_recording_context().save_func_return_value("docStoreConn.delete", doc_store_result)
             tasks = []
             for chunk_id in chunk_ids:
                 tasks.append(asyncio.create_task(delete_image(task_dataset_id, chunk_id)))
@@ -1277,7 +1366,8 @@ async def do_handle_task(task):
     task_type = task.get("task_type", "")
 
     if task_type == "memory":
-        await handle_save_to_memory_task(task)
+        result = await handle_save_to_memory_task(task)
+        get_recording_context().save_func_return_value("handle_save_to_memory_task", result)
         return
 
     if task_type == "dataflow" and task.get("doc_id", "") == CANVAS_DEBUG_DOC_ID:
@@ -1355,7 +1445,9 @@ async def do_handle_task(task):
                     },
                 }
             )
-            if not KnowledgebaseService.update_by_id(kb.id, {"parser_config": kb_parser_config}):
+            update_result = KnowledgebaseService.update_by_id(kb.id, {"parser_config": kb_parser_config})
+            get_recording_context().save_func_return_value("KnowledgebaseService.update_by_id", update_result)
+            if not update_result:
                 progress_callback(prog=-1.0, msg="Internal error: Invalid RAPTOR configuration")
                 return
 
@@ -1373,6 +1465,8 @@ async def do_handle_task(task):
                 callback=progress_callback,
                 doc_ids=task.get("doc_ids", []),
             )
+        get_recording_context().record("raptor_chunks", chunks)
+        get_recording_context().record("raptor_token_count", token_count)
         if fake_doc_ids := task.get("doc_ids", []):
             task_doc_id = fake_doc_ids[0]  # use the first document ID to represent this task for logging purposes
     # Either using graphrag or Standard chunking methods
@@ -1409,7 +1503,9 @@ async def do_handle_task(task):
                     }
                 }
             )
-            if not KnowledgebaseService.update_by_id(kb.id, {"parser_config": kb_parser_config}):
+            update_result = KnowledgebaseService.update_by_id(kb.id, {"parser_config": kb_parser_config})
+            get_recording_context().save_func_return_value("KnowledgebaseService.update_by_id", update_result)
+            if not update_result:
                 progress_callback(prog=-1.0, msg="Internal error: Invalid GraphRAG configuration")
                 return
 
@@ -1434,6 +1530,7 @@ async def do_handle_task(task):
                 with_community=with_community,
             )
             logging.info(f"GraphRAG task result for task {task}:\n{result}")
+        get_recording_context().record("graphrag_result", result)
         progress_callback(prog=1.0, msg="Knowledge Graph done ({:.2f}s)".format(timer() - start_ts))
         return
     elif task_type == "mindmap":
@@ -1445,6 +1542,11 @@ async def do_handle_task(task):
         task['llm_id'] = doc_task_llm_id
         start_ts = timer()
         chunks = await build_chunks(task, progress_callback)
+        get_recording_context().record("chunks", chunks)
+        # Record chunk_ids_count for comparison
+        chunk_ids = [c.get("id") for c in chunks if isinstance(c, dict) and "id" in c]
+        get_recording_context().record("chunk_ids_count", len(chunk_ids))
+        # Record chunks array for content comparison (first, middle, last, random)
         logging.info("Build document {}: {:.2f}s".format(task_document_name, timer() - start_ts))
         if not chunks:
             progress_callback(1., msg=f"No chunk built from {task_document_name}")
@@ -1461,6 +1563,8 @@ async def do_handle_task(task):
             logging.exception(error_message)
             token_count = 0
             raise
+        get_recording_context().record("token_count", token_count)
+        get_recording_context().record("vector_size", vector_size)
         progress_message = "Embedding chunks ({:.2f}s)".format(timer() - start_ts)
         logging.info(progress_message)
         progress_callback(msg=progress_message)
@@ -1479,7 +1583,9 @@ async def do_handle_task(task):
 
     try:
         if not await _maybe_insert_chunks(chunks):
+            get_recording_context().record("insertion_result", "failed")
             return
+        get_recording_context().record("insertion_result", "success")
         if has_canceled(task_id):
             progress_callback(-1, msg="Task has been canceled.")
             return
@@ -1487,12 +1593,15 @@ async def do_handle_task(task):
         if raptor_cleanup_chunks:
             cleaned_chunks = 0
             for cleanup_doc_id, keep_method in raptor_cleanup_chunks:
-                cleaned_chunks += await delete_raptor_chunks(
+                ret = await delete_raptor_chunks(
                     cleanup_doc_id,
                     task_tenant_id,
                     task_dataset_id,
                     keep_method=keep_method,
                 )
+                cleaned_chunks += ret
+                get_recording_context().save_func_return_value("delete_raptor_chunks", ret)
+
             if cleaned_chunks:
                 progress_callback(msg=f"Cleaned up {cleaned_chunks} stale RAPTOR chunks.")
 
@@ -1502,7 +1611,8 @@ async def do_handle_task(task):
             )
         )
 
-        DocumentService.increment_chunk_num(task_doc_id, task_dataset_id, token_count, chunk_count, 0)
+        ret = DocumentService.increment_chunk_num(task_doc_id, task_dataset_id, token_count, chunk_count, 0)
+        get_recording_context().save_func_return_value("DocumentService.increment_chunk_num", ret)
 
         # Table parser (manual): push metadata/both column values to document-level metadata for UI / chat filters
         if task.get("parser_id", "").lower() == "table":
@@ -1525,7 +1635,8 @@ async def do_handle_task(task):
                         f"table_strip_key_count={len(strip_keys)}, agg_keys={list(agg.keys())}"
                     )
                     try:
-                        DocMetadataService.update_document_metadata(task_doc_id, merged)
+                        ret = DocMetadataService.update_document_metadata(task_doc_id, merged)
+                        get_recording_context().save_func_return_value("DocMetadataService.update_document_metadata", ret)
                         logging.debug("[TABLE_META_DEBUG] update_document_metadata succeeded")
                     except Exception as ue:
                         logging.error(
@@ -1546,15 +1657,20 @@ async def do_handle_task(task):
         if toc_thread:
             d = await toc_thread
             if d:
+                get_recording_context().record("toc_chunk", [d])
                 if not await _maybe_insert_chunks([d]):
+                    get_recording_context().record("toc_inserted", False)
                     return
-                DocumentService.increment_chunk_num(task_doc_id, task_dataset_id, 0, 1, 0)
+                get_recording_context().record("toc_inserted", True)
+                ret = DocumentService.increment_chunk_num(task_doc_id, task_dataset_id, 0, 1, 0)
+                get_recording_context().save_func_return_value("DocumentService.increment_chunk_num", ret)
 
         if has_canceled(task_id):
             progress_callback(-1, msg="Task has been canceled.")
             return
 
         task_time_cost = timer() - task_start_ts
+        get_recording_context().record("task_status", "completed")
         progress_callback(prog=1.0, msg="Task done ({:.2f}s)".format(task_time_cost))
         logging.info(
             "Chunk doc({}), page({}-{}), chunks({}), token({}), elapsed:{:.2f}".format(
@@ -1573,12 +1689,13 @@ async def do_handle_task(task):
                     task_dataset_id,
                 )
                 if exists:
-                    await thread_pool_exec(
+                    ret = await thread_pool_exec(
                         settings.docStoreConn.delete,
                         {"doc_id": task_doc_id},
                         search.index_name(task_tenant_id),
                         task_dataset_id,
                     )
+                    get_recording_context().save_func_return_value("docStoreConn.delete", ret)
             except Exception as e:
                 logging.exception(
                     f"Remove doc({task_doc_id}) from docStore failed when task({task_id}) canceled, exception: {e}")
@@ -1596,9 +1713,28 @@ async def handle_task():
                                                              PipelineTaskType.PARSE) or PipelineTaskType.PARSE
     task_id = task["id"]
     try:
-        logging.info(f"handle_task begin for task {json.dumps(task)}")
         CURRENT_TASKS[task["id"]] = copy.deepcopy(task)
-        await do_handle_task(task)
+        run_mode = os.environ.get("TE_RUN_MODE", "0")
+        logging.info(f"TE_RUN_MODE is {run_mode}")
+
+        # Check if dry-run comparison is enabled via environment variable
+        if run_mode == "1": # dry run mode - compare
+            set_recording_context(RecordingContext())
+            await do_handle_task(task) # original execution
+            # dry run mode
+            logging.info(f"-----dry run task:{task_id}, {task.get('name', '')}, doc id:{task.get('doc_id', '')}")
+            await TaskManager.dry_run_task(task, get_recording_context(), chat_limiter, minio_limiter, chunk_limiter,
+                               embed_limiter,kg_limiter, set_progress, has_canceled)
+        elif run_mode == "0": # use refactor-ed version
+            # switch to refactor-ed version
+            logging.info(f"-----run refactor-ed task executor:{task_id}, {task.get('name', '')}, doc id:{task.get('doc_id', '')}")
+            await TaskManager.run_refactored_task(task, chat_limiter, minio_limiter, chunk_limiter,
+                               embed_limiter,kg_limiter, set_progress, has_canceled)
+        else: # original version
+            logging.info(f"-----run original task executor:{task_id}, {task.get('name', '')}, doc id:{task.get('doc_id', '')}")
+            set_recording_context(NullRecordingContext())
+            await do_handle_task(task)
+
         DONE_TASKS += 1
         CURRENT_TASKS.pop(task_id, None)
         logging.info(f"handle_task done for task {json.dumps(task)}")
@@ -1626,9 +1762,10 @@ async def handle_task():
             referred_document_id = None
             if task_type in ["graphrag", "raptor", "mindmap"]:
                 referred_document_id = task["doc_ids"][0]
-            PipelineOperationLogService.record_pipeline_operation(document_id=task["doc_id"], pipeline_id="",
+            ret = PipelineOperationLogService.record_pipeline_operation(document_id=task["doc_id"], pipeline_id="",
                                                                   task_type=pipeline_task_type,
                                                                   task_id=task_id, referred_document_id=referred_document_id)
+            get_recording_context().save_func_return_value("PipelineOperationLogService.record_pipeline_operation", ret)
 
     redis_msg.ack()
 
@@ -1685,7 +1822,8 @@ async def report_status():
         except Exception as e:
             logging.warning(f"Failed to report heartbeat: {e}")
         else:
-            logging.info(f"{CONSUMER_NAME} reported heartbeat: {heartbeat}")
+            logging.debug(f"{CONSUMER_NAME} reported heartbeat: {heartbeat}")
+            pass
 
         # Clean up own expired heartbeat
         try:
@@ -1752,6 +1890,7 @@ async def main():
           /____/
     """)
     logging.info(f'RAGFlow ingestion version: {get_ragflow_version()}')
+    logging.info(f"ENABLE_DRY_RUN_COMPARISON: {os.environ.get("ENABLE_DRY_RUN_COMPARISON", "0")}")
     show_configs()
     settings.init_settings()
     settings.check_and_install_torch()
@@ -1786,6 +1925,17 @@ async def main():
 
 
 if __name__ == "__main__":
+    # Parse command line arguments (consistent with SAAS version)
+    parser = argparse.ArgumentParser(description='Task Executor')
+    parser.add_argument("-i", "--index", type=str, default='0')
+    parser.add_argument("-t", "--type", type=str, default="common", help="[common, graphrag, raptor, resume]")
+    args = parser.parse_args()
+    
+    # Update global variables
+    TASK_TYPE = args.type
+    TE_IDX = args.index
+    CONSUMER_NAME = f"task_executor_{TASK_TYPE}_{TE_IDX}"
+    
     faulthandler.enable()
     init_root_logger(CONSUMER_NAME)
     try:

rag/svr/task_executor_limiter.py

@@ -0,0 +1,28 @@
+#
+#  Copyright 2025 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 os
+
+from common.asyncio_utils import LoopLocalSemaphore
+
+MAX_CONCURRENT_TASKS = int(os.environ.get("MAX_CONCURRENT_TASKS", "5"))
+MAX_CONCURRENT_CHUNK_BUILDERS = int(os.environ.get("MAX_CONCURRENT_CHUNK_BUILDERS", "1"))
+MAX_CONCURRENT_MINIO = int(os.environ.get("MAX_CONCURRENT_MINIO", "10"))
+
+task_limiter = LoopLocalSemaphore(MAX_CONCURRENT_TASKS)
+chunk_limiter = LoopLocalSemaphore(MAX_CONCURRENT_CHUNK_BUILDERS)
+embed_limiter = LoopLocalSemaphore(MAX_CONCURRENT_CHUNK_BUILDERS)
+minio_limiter = LoopLocalSemaphore(MAX_CONCURRENT_MINIO)
+kg_limiter = LoopLocalSemaphore(2)

rag/svr/task_executor_refactor/chunk_builder.py

@@ -0,0 +1,136 @@
+#
+#  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.
+
+"""
+Chunk Builder Module.
+
+Provides parser factory and document chunking logic:
+- Parser module registration and selection
+- Document chunking via parser
+- PDF outline extraction
+"""
+
+import logging
+from timeit import default_timer as timer
+from typing import Dict, List
+
+from common.constants import ParserType
+from common.misc_utils import thread_pool_exec
+from rag.svr.task_executor_refactor.task_context import TaskContext
+
+from api.db.services.doc_metadata_service import DocMetadataService
+from common.metadata_utils import update_metadata_to
+from rag.utils.table_es_metadata import merge_table_parser_config_from_kb
+
+
+def get_parser(parser_id: str):
+    """Get parser module by ID.
+
+    Args:
+        parser_id: The parser identifier.
+
+    Returns:
+        The parser module for the given parser ID.
+    """
+    from rag.app import laws, paper, presentation, manual, qa, table, book, resume, picture, naive, one, audio, email, tag
+
+    factory = {
+        "general": naive,
+        ParserType.NAIVE.value: naive,
+        ParserType.PAPER.value: paper,
+        ParserType.BOOK.value: book,
+        ParserType.PRESENTATION.value: presentation,
+        ParserType.MANUAL.value: manual,
+        ParserType.LAWS.value: laws,
+        ParserType.QA.value: qa,
+        ParserType.TABLE.value: table,
+        ParserType.RESUME.value: resume,
+        ParserType.PICTURE.value: picture,
+        ParserType.ONE.value: one,
+        ParserType.AUDIO.value: audio,
+        ParserType.EMAIL.value: email,
+        ParserType.KG.value: naive,
+        ParserType.TAG.value: tag,
+    }
+    return factory[parser_id.lower()]
+
+
+async def run_chunking(
+    chunker,
+    binary: bytes,
+    ctx: TaskContext,
+) -> List[Dict]:
+    """Run document chunking via parser.
+
+    Args:
+        chunker: The parser module to use.
+        binary: Binary content of the document.
+        ctx: TaskContext containing task configuration.
+
+    Returns:
+        List of chunk dictionaries.
+    """
+    st = timer()
+    try:
+        # Merge table parser config
+        parser_config = merge_table_parser_config_from_kb(ctx.raw_task)
+
+        async with ctx.chunk_limiter:
+            cks = await thread_pool_exec(
+                chunker.chunk,
+                ctx.name,
+                binary=binary,
+                from_page=ctx.from_page,
+                to_page=ctx.to_page,
+                lang=ctx.language,
+                callback=ctx.progress_cb,
+                kb_id=ctx.kb_id,
+                parser_config=parser_config,
+                tenant_id=ctx.tenant_id,
+            )
+        logging.info("Chunking({}) {}/{} done".format(timer() - st, ctx.location, ctx.name))
+        ctx.recording_context.record("parser_config_after_merge", parser_config)
+        return cks
+    except Exception as e:
+        ctx.progress_cb(-1, msg="Internal server error while chunking: %s" % str(e).replace("'", ""))
+        logging.exception("Chunking {}/{} got exception".format(ctx.location, ctx.name))
+        raise
+
+
+async def extract_outline(cks: List[Dict], ctx: TaskContext) -> None:
+    """Extract and persist PDF outline if present.
+
+    Args:
+        cks: List of chunk dictionaries.
+        ctx: TaskContext containing task configuration.
+    """
+    outline_data = cks[0].get("__outline__") if cks else None
+    ctx.recording_context.record("outline_data", outline_data)
+
+    if cks and cks[0].get("__outline__"):
+        outline = cks[0].pop("__outline__")
+        try:
+            if ctx.write_interceptor:
+                ctx.write_interceptor.intercept("DocMetadataService.update_document_metadata")
+            else:
+                temp_doc = DocMetadataService.get_document_metadata(ctx.doc_id) or {}
+                DocMetadataService.update_document_metadata(
+                    ctx.doc_id,
+                    update_metadata_to({"outline": outline}, temp_doc)
+                )
+
+            logging.info("Persisted PDF outline (%d entries) for doc %s", len(outline), ctx.doc_id)
+        except Exception as e:
+            logging.warning("Failed to persist PDF outline for doc %s: %s", ctx.doc_id, e)

rag/svr/task_executor_refactor/chunk_post_processor.py

@@ -0,0 +1,308 @@
+#
+#  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.
+
+"""
+Chunk Post-Processor Module.
+
+Provides post-processing functions for chunks:
+- Keyword extraction
+- Question generation
+- Metadata generation
+- Content tagging
+"""
+
+import asyncio
+import json
+import logging
+import random
+import re
+from timeit import default_timer as timer
+from typing import Dict, List
+
+from common.constants import TAG_FLD, LLMType
+from common.metadata_utils import turn2jsonschema, update_metadata_to
+from common import settings
+from rag.nlp import rag_tokenizer
+from rag.svr.task_executor_refactor.task_context import TaskContext
+
+from api.db.services.doc_metadata_service import DocMetadataService
+from api.db.services.llm_service import LLMBundle
+from api.db.joint_services.tenant_model_service import get_model_config_by_type_and_name
+from rag.prompts.generator import gen_metadata, keyword_extraction, question_proposal, content_tagging
+from rag.graphrag.utils import get_llm_cache, set_llm_cache
+
+
+async def extract_keywords(docs: List[Dict], ctx: TaskContext) -> None:
+    """Extract keywords for chunks.
+
+    Args:
+        docs: List of chunk dictionaries to process.
+        ctx: TaskContext containing task configuration.
+    """
+    chat_limiter = ctx.chat_limiter
+
+    st = timer()
+    ctx.progress_cb(msg="Start to generate keywords for every chunk ...")
+    chat_model_config = get_model_config_by_type_and_name(ctx.tenant_id, LLMType.CHAT, ctx.llm_id)
+    with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_model:
+
+        async def doc_keyword_extraction(chat_mdl, d, topn):
+            cached = get_llm_cache(chat_mdl.llm_name, d["content_with_weight"], "keywords", {"topn": topn})
+            if not cached:
+                if ctx.has_canceled_func(ctx.id):
+                    ctx.progress_cb(-1, msg="Task has been canceled.")
+                    return
+                async with chat_limiter:
+                    cached = await keyword_extraction(chat_mdl, d["content_with_weight"], topn)
+                set_llm_cache(chat_mdl.llm_name, d["content_with_weight"], cached, "keywords", {"topn": topn})
+            if cached:
+                d["important_kwd"] = [k for k in re.split(r"[,，;；、\r\n]+", cached) if k.strip()]
+                d["important_tks"] = rag_tokenizer.tokenize(" ".join(d["important_kwd"]))
+            return
+
+        tasks = []
+        for doc in docs:
+            tasks.append(
+                asyncio.create_task(doc_keyword_extraction(chat_model, doc, ctx.parser_config["auto_keywords"])))
+        try:
+            await asyncio.gather(*tasks, return_exceptions=False)
+        except Exception as e:
+            logging.error("Error in doc_keyword_extraction: {}".format(e))
+            for t in tasks:
+                t.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            raise
+        ctx.progress_cb(msg="Keywords generation {} chunks completed in {:.2f}s".format(len(docs), timer() - st))
+
+
+async def generate_questions(docs: List[Dict], ctx: TaskContext) -> None:
+    """Generate questions for chunks.
+
+    Args:
+        docs: List of chunk dictionaries to process.
+        ctx: TaskContext containing task configuration.
+    """
+    chat_limiter = ctx.chat_limiter
+
+    st = timer()
+    ctx.progress_cb(msg="Start to generate questions for every chunk ...")
+    chat_model_config = get_model_config_by_type_and_name(ctx.tenant_id, LLMType.CHAT, ctx.llm_id)
+    with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_model:
+
+        async def doc_question_proposal(chat_mdl, d, topn):
+            cached = get_llm_cache(chat_mdl.llm_name, d["content_with_weight"], "question", {"topn": topn})
+            if not cached:
+                if ctx.has_canceled_func(ctx.id):
+                    ctx.progress_cb(-1, msg="Task has been canceled.")
+                    return
+                async with chat_limiter:
+                    cached = await question_proposal(chat_mdl, d["content_with_weight"], topn)
+                set_llm_cache(chat_mdl.llm_name, d["content_with_weight"], cached, "question", {"topn": topn})
+            if cached:
+                d["question_kwd"] = cached.split("\n")
+                d["question_tks"] = rag_tokenizer.tokenize("\n".join(d["question_kwd"]))
+
+        tasks = []
+        for doc in docs:
+            tasks.append(
+                asyncio.create_task(doc_question_proposal(chat_model, doc, ctx.parser_config["auto_questions"])))
+        try:
+            await asyncio.gather(*tasks, return_exceptions=False)
+        except Exception as e:
+            logging.error("Error in doc_question_proposal", exc_info=e)
+            for t in tasks:
+                t.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            raise
+        ctx.progress_cb(msg="Question generation {} chunks completed in {:.2f}s".format(len(docs), timer() - st))
+
+
+def build_metadata_config(parser_config: dict) -> list:
+    """Build the metadata configuration from parser_config.
+
+    Extracts and normalizes ``metadata`` and ``built_in_metadata`` from the
+    parser configuration into a single list or dict that is passed to the LLM
+    cache and generation functions.
+
+    This should be called once per ``generate_metadata`` invocation — the result
+    is identical for every chunk within the same document parse session so
+    extracting it avoids rebuilding inside the per-chunk async task.
+
+    Args:
+        parser_config: Configuration dict from the parser, expected to contain
+            ``metadata`` (dict or list) and optionally ``built_in_metadata``
+            (list of metadata item dicts).
+
+    Returns:
+        A list or dict representing the merged metadata configuration.
+    """
+    metadata_conf = parser_config.get("metadata", [])
+    built_in_metadata = list(parser_config.get("built_in_metadata") or [])
+    if isinstance(metadata_conf, dict):
+        if not isinstance(metadata_conf.get("properties"), dict):
+            metadata_conf = {"type": "object", "properties": {}}
+        if built_in_metadata:
+            metadata_conf = {
+                **metadata_conf,
+                "properties": {
+                    **metadata_conf.get("properties", {}),
+                    **turn2jsonschema(built_in_metadata).get("properties", {}),
+                },
+            }
+    elif isinstance(metadata_conf, list):
+        metadata_conf = metadata_conf + built_in_metadata
+    else:
+        metadata_conf = built_in_metadata
+    return metadata_conf
+
+
+async def generate_metadata(docs: List[Dict], ctx: TaskContext) -> None:
+    """Generate metadata for chunks.
+
+    Args:
+        docs: List of chunk dictionaries to process.
+        ctx: TaskContext containing task configuration.
+    """
+    chat_limiter = ctx.chat_limiter
+
+    st = timer()
+    ctx.progress_cb(msg="Start to generate meta-data for every chunk ...")
+    chat_model_config = get_model_config_by_type_and_name(ctx.tenant_id, LLMType.CHAT, ctx.llm_id)
+    with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_model:
+        metadata_conf = build_metadata_config(ctx.parser_config)
+
+        async def gen_metadata_task(chat_mdl, d):
+            cached = get_llm_cache(chat_mdl.llm_name, d["content_with_weight"], "metadata",
+                                   metadata_conf)
+            if not cached:
+                if ctx.has_canceled_func(ctx.id):
+                    ctx.progress_cb(-1, msg="Task has been canceled.")
+                    return
+                async with chat_limiter:
+                    cached = await gen_metadata(chat_mdl,
+                                                turn2jsonschema(metadata_conf),
+                                                d["content_with_weight"])
+                set_llm_cache(chat_mdl.llm_name, d["content_with_weight"], cached, "metadata",
+                              metadata_conf)
+            if cached:
+                d["metadata_obj"] = cached
+
+        tasks = []
+        for doc in docs:
+            tasks.append(asyncio.create_task(gen_metadata_task(chat_model, doc)))
+        try:
+            await asyncio.gather(*tasks, return_exceptions=False)
+        except Exception as e:
+            logging.error("Error in gen_metadata", exc_info=e)
+            for t in tasks:
+                t.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            raise
+
+        metadata = {}
+        for doc in docs:
+            if "metadata_obj" in doc:
+                metadata = update_metadata_to(metadata, doc["metadata_obj"])
+                del doc["metadata_obj"]
+        if metadata:
+            existing_meta = DocMetadataService.get_document_metadata(ctx.doc_id)
+            existing_meta = existing_meta if isinstance(existing_meta, dict) else {}
+            metadata = update_metadata_to(metadata, existing_meta)
+            if ctx.write_interceptor:
+                ctx.write_interceptor.intercept("DocMetadataService.update_document_metadata")
+            else:
+                DocMetadataService.update_document_metadata(ctx.doc_id, metadata)
+        ctx.progress_cb(msg="Metadata generation {} chunks completed in {:.2f}s".format(len(docs), timer() - st))
+
+
+async def apply_tags(docs: List[Dict], ctx: TaskContext) -> None:
+    """Apply tags to chunks.
+
+    Args:
+        docs: List of chunk dictionaries to process.
+        ctx: TaskContext containing task configuration.
+    """
+    chat_limiter = ctx.chat_limiter
+
+    ctx.progress_cb(msg="Start to tag for every chunk ...")
+    kb_ids = ctx.kb_parser_config["tag_kb_ids"]
+    tenant_id = ctx.tenant_id
+    topn_tags = ctx.kb_parser_config.get("topn_tags", 3)
+    S = 1000
+    st = timer()
+    examples = []
+    all_tags = settings.retriever.all_tags_in_portion(tenant_id, kb_ids, S)
+    chat_model_config = get_model_config_by_type_and_name(tenant_id, LLMType.CHAT, ctx.llm_id)
+    with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_model:
+
+        docs_to_tag = []
+        for doc in docs:
+            if ctx.has_canceled_func(ctx.id):
+                ctx.progress_cb(-1, msg="Task has been canceled.")
+                return
+            if settings.retriever.tag_content(tenant_id, kb_ids, doc, all_tags, topn_tags=topn_tags, S=S) and len(
+                    doc.get(TAG_FLD, [])) > 0:
+                examples.append({"content": doc["content_with_weight"], TAG_FLD: doc[TAG_FLD]})
+            else:
+                docs_to_tag.append(doc)
+
+        async def doc_content_tagging(chat_mdl, d, topn_tags):
+            cached = get_llm_cache(chat_mdl.llm_name, d["content_with_weight"], all_tags, {"topn": topn_tags})
+            if not cached:
+                if ctx.has_canceled_func(ctx.id):
+                    ctx.progress_cb(-1, msg="Task has been canceled.")
+                    return
+                picked_examples = random.choices(examples, k=2) if len(examples) > 2 else examples
+                if not picked_examples:
+                    picked_examples.append({"content": "This is an example", TAG_FLD: {'example': 1}})
+                async with chat_limiter:
+                    cached = await content_tagging(
+                        chat_mdl,
+                        d["content_with_weight"],
+                        all_tags,
+                        picked_examples,
+                        topn_tags,
+                    )
+                if cached:
+                    cached = json.dumps(cached)
+            if cached:
+                set_llm_cache(chat_mdl.llm_name, d["content_with_weight"], cached, all_tags, {"topn": topn_tags})
+                d[TAG_FLD] = json.loads(cached)
+
+        tasks = []
+        for doc in docs_to_tag:
+            tasks.append(asyncio.create_task(doc_content_tagging(chat_model, doc, topn_tags)))
+        try:
+            await asyncio.gather(*tasks, return_exceptions=False)
+        except Exception as e:
+            logging.error("Error tagging docs: {}".format(e))
+            for t in tasks:
+                t.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            raise
+        ctx.progress_cb(msg="Tagging {} chunks completed in {:.2f}s".format(len(docs), timer() - st))
+
+
+def count_with_key(docs: List[Dict], key: str) -> int:
+    """Count docs that have a specific key.
+
+    Args:
+        docs: List of chunk dictionaries.
+        key: The key to check for.
+
+    Returns:
+        Count of docs that have the key.
+    """
+    return sum(1 for d in docs if d.get(key))

rag/svr/task_executor_refactor/chunk_service.py

@@ -0,0 +1,479 @@
+#
+#  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.
+
+"""
+Chunk Service Module.
+
+Provides [`ChunkService`](rag/svr/task_executor_refactor/chunk_service.py:50) for document chunking,
+post-processing (keywords, questions, metadata, tags), MinIO upload, and chunk insertion into document store.
+
+This module orchestrates the chunk building pipeline by delegating to:
+- [`chunk_builder`](rag/svr/task_executor_refactor/chunk_builder.py): Parser selection and document chunking
+- [`chunk_post_processor`](rag/svr/task_executor_refactor/chunk_post_processor.py): Post-processing functions
+"""
+
+import asyncio
+import copy
+import logging
+from datetime import datetime
+from functools import partial
+from timeit import default_timer as timer
+from typing import Any, Dict, List
+
+import xxhash
+from common import settings
+from common.constants import PAGERANK_FLD, TAG_FLD
+from common.misc_utils import thread_pool_exec
+from common.float_utils import normalize_overlapped_percent
+from rag.nlp import search
+from rag.svr.task_executor_refactor.task_context import TaskContext
+from rag.utils.base64_image import image2id
+
+from api.db.services.task_service import TaskService
+from rag.svr.task_executor_refactor.constants import GRAPH_RAPTOR_FAKE_DOC_ID
+
+# Re-export for backward compatibility
+from rag.svr.task_executor_refactor.chunk_builder import (
+    get_parser,
+    run_chunking,
+    extract_outline,
+)
+from rag.svr.task_executor_refactor.chunk_post_processor import (
+    extract_keywords,
+    generate_questions,
+    generate_metadata,
+    apply_tags,
+)
+
+
+class ChunkService:
+    """Service for document chunking and post-processing.
+
+    This service handles:
+    - Document chunking via parser modules (delegated to chunk_builder)
+    - MinIO upload of chunk images
+    - Keyword extraction (delegated to chunk_post_processor)
+    - Question generation (delegated to chunk_post_processor)
+    - Metadata generation (delegated to chunk_post_processor)
+    - Content tagging (delegated to chunk_post_processor)
+    - Table of contents generation
+    - Chunk insertion into document store
+
+    All intermediate results are recorded via RecordingContext for comparison.
+    """
+
+    def __init__(
+        self,
+        ctx: TaskContext,
+    ):
+        """Initialize ChunkService.
+
+        Args:
+            ctx: TaskContext containing task configuration and execution resources.
+        """
+        self._task_context = ctx
+
+    async def build_chunks(
+        self,
+        storage_binary: bytes,
+    ) -> List[Dict[str, Any]]:
+        """Build chunks from document binary.
+
+        This is the main entry point for chunk building. It orchestrates:
+        1. File size validation
+        2. Parser selection and chunking (delegated to chunk_builder)
+        3. Outline extraction (delegated to chunk_builder)
+        4. MinIO upload
+        5. Post-processing (delegated to chunk_post_processor)
+
+        Args:
+            storage_binary: Binary content of the document.
+
+        Returns:
+            List of chunk dictionaries ready for embedding.
+        """
+        ctx = self._task_context
+        # Validate file size
+        if ctx.size > settings.DOC_MAXIMUM_SIZE:
+            self._progress(prog=-1, msg="File size exceeds( <= %dMb )" %
+                          (int(settings.DOC_MAXIMUM_SIZE / 1024 / 1024)))
+            self._task_context.recording_context.record("file_size_exceeded", True)
+            return []
+        ctx.recording_context.record("file_size_exceeded", False)
+        ctx.recording_context.record("parser_id", ctx.parser_id)
+
+        # Get parser
+        chunker = get_parser(ctx.parser_id)
+
+        # record config for compare
+        chunk_config = {
+            "parser_id": ctx.parser_id,
+            "chunk_token_num": ctx.parser_config.get("chunk_token_num", 128),
+            "overlapped_percent": normalize_overlapped_percent(
+                ctx.parser_config.get("overlapped_percent", 0)
+            ),
+            "delimiter": ctx.parser_config.get("delimiter", "\n!?。；！？"),
+            "from_page": ctx.from_page,
+            "to_page": ctx.to_page,
+            "language": ctx.language,
+            "layout_recognizer": ctx.parser_config.get("layout_recognizer"),
+        }
+        ctx.recording_context.record("chunk_config", chunk_config)
+
+        # Run chunking (delegated)
+        cks = await run_chunking(chunker, storage_binary, ctx)
+
+        # Record raw chunks
+        self._task_context.recording_context.record("raw_chunks", cks)
+
+        # Extract outline (delegated)
+        await extract_outline(cks, ctx)
+
+        # Prepare docs and upload to MinIO
+        docs = await self._prepare_docs_and_upload(cks)
+
+        # Record docs after prep
+        self._task_context.recording_context.record("docs_after_prep", docs)
+
+        # Post-processing (delegated to chunk_post_processor)
+        if ctx.parser_config.get("auto_keywords", 0):
+            await extract_keywords(docs, ctx)
+        keywords = [d for d in docs if d.get("important_kwd")]
+        self._task_context.recording_context.record("keywords_extracted", keywords)
+
+        if ctx.parser_config.get("auto_questions", 0):
+            await generate_questions(docs, ctx)
+        questions = [d for d in docs if d.get("question_kwd")]
+        self._task_context.recording_context.record("questions_generated", questions)
+
+        if ctx.parser_config.get("enable_metadata", False) and (
+            ctx.parser_config.get("metadata") or ctx.parser_config.get("built_in_metadata")
+        ):
+            await generate_metadata(docs, ctx)
+        metadata_list = [d for d in docs if d.get("metadata_obj")]
+        self._task_context.recording_context.record("metadata_list_generated", metadata_list)
+
+        if ctx.kb_parser_config.get("tag_kb_ids", []):
+            await apply_tags(docs, ctx)
+        tags_applied = [d for d in docs if d.get(TAG_FLD)]
+        self._task_context.recording_context.record("tags_applied", tags_applied)
+
+        # Record final chunks
+        self._task_context.recording_context.record("final_chunks", docs)
+        final_chunk_ids = [c.get("id") for c in docs if isinstance(c, dict) and "id" in c]
+        self._task_context.recording_context.record("final_chunk_ids_count", len(final_chunk_ids))
+
+        return docs
+
+    async def _prepare_docs_and_upload(self, cks: List[Dict]) -> List[Dict]:
+        """Prepare docs and upload images to MinIO."""
+        ctx = self._task_context
+        docs = []
+        doc = {
+            "doc_id": ctx.doc_id,
+            "kb_id": str(ctx.kb_id)
+        }
+        if ctx.pagerank:
+            doc[PAGERANK_FLD] = int(ctx.pagerank)
+
+        st = timer()
+
+        async def upload_to_minio(document, chunk):
+            try:
+                d = copy.deepcopy(document)
+                d.update(chunk)
+                d["id"] = xxhash.xxh64(
+                    (chunk["content_with_weight"] + str(d["doc_id"])).encode("utf-8", "surrogatepass")).hexdigest()
+                d["create_time"] = str(datetime.now()).replace("T", " ")[:19]
+                d["create_timestamp_flt"] = datetime.now().timestamp()
+
+                if d.get("img_id"):
+                    docs.append(d)
+                    return
+
+                if not d.get("image"):
+                    _ = d.pop("image", None)
+                    d["img_id"] = ""
+                    docs.append(d)
+                    return
+
+                await image2id(d, partial(settings.STORAGE_IMPL.put, tenant_id=ctx.tenant_id), d["id"], ctx.kb_id)
+                docs.append(d)
+            except Exception:
+                logging.exception(
+                    "Saving image of chunk {}/{}/{} got exception".format(ctx.location, ctx.name, d["id"]))
+                raise
+
+        tasks = []
+        for ck in cks:
+            tasks.append(asyncio.create_task(upload_to_minio(doc, ck)))
+        try:
+            await asyncio.gather(*tasks, return_exceptions=False)
+        except Exception as e:
+            logging.error(f"MINIO PUT({ctx.name}) got exception: {e}")
+            for t in tasks:
+                t.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            raise
+
+        el = timer() - st
+        logging.info("MINIO PUT({}) cost {:.3f} s".format(ctx.name, el))
+        return docs
+
+    def _progress(self, prog=None, msg=None):
+        """Progress callback helper."""
+        if prog is not None or msg is not None:
+            self._task_context.progress_cb(prog=prog, msg=msg)
+
+    # =========================================================================
+    # Insert Service Methods (merged from insert_service.py)
+    # =========================================================================
+
+    async def insert_chunks(
+        self,
+        task_id: str,
+        task_tenant_id: str,
+        task_dataset_id: str,
+        chunks: List[Dict[str, Any]],
+        doc_bulk_size: int = None,
+    ) -> bool:
+        """Insert chunks into document store.
+
+        Args:
+            task_id: Task identifier.
+            task_tenant_id: Tenant ID.
+            task_dataset_id: Dataset/knowledge base ID.
+            chunks: List of chunk dictionaries to insert.
+            doc_bulk_size: Batch size for document store inserts.
+
+        Returns:
+            True if all chunks were inserted successfully, False otherwise.
+        """
+        doc_bulk_size = doc_bulk_size or settings.DOC_BULK_SIZE
+
+        # Create mother chunks (summary chunks)
+        mothers = self._create_mother_chunks(chunks)
+
+        # Insert mother chunks
+        if not await self._insert_mother_chunks(task_id, task_tenant_id, task_dataset_id, mothers, doc_bulk_size):
+            return False
+
+        # Insert main chunks
+        return await self._insert_main_chunks(task_id, task_tenant_id, task_dataset_id, chunks, doc_bulk_size)
+
+    @classmethod
+    def _create_mother_chunks(cls, chunks: List[Dict]) -> List[Dict]:
+        """Create mother chunks from summary fields.
+
+        Mother chunks are summary/abstract chunks that are stored separately.
+        """
+        mothers = []
+        mother_ids = set()
+
+        for ck in chunks:
+            mom = ck.get("mom") or ck.get("mom_with_weight") or ""
+            if not mom:
+                continue
+
+            mom_id = xxhash.xxh64(mom.encode("utf-8")).hexdigest()
+            ck["mom_id"] = mom_id
+
+            if mom_id in mother_ids:
+                continue
+
+            mother_ids.add(mom_id)
+            mom_ck = copy.deepcopy(ck)
+            mom_ck["id"] = mom_id
+            mom_ck["content_with_weight"] = mom
+            mom_ck["available_int"] = 0
+
+            # Keep only essential fields
+            allowed_fields = [
+                "id", "content_with_weight", "doc_id", "docnm_kwd",
+                "kb_id", "available_int", "position_int",
+                "create_timestamp_flt", "page_num_int", "top_int"
+            ]
+            for fld in list(mom_ck.keys()):
+                if fld not in allowed_fields:
+                    del mom_ck[fld]
+
+            mothers.append(mom_ck)
+
+        return mothers
+
+    async def _insert_mother_chunks(
+        self,
+        task_id: str,
+        task_tenant_id: str,
+        task_dataset_id: str,
+        mothers: List[Dict],
+        doc_bulk_size: int,
+    ) -> bool:
+        """Insert mother chunks in batches."""
+        for b in range(0, len(mothers), doc_bulk_size):
+            await self._intercept_doc_store_insert(
+                mothers[b:b + doc_bulk_size],
+                search.index_name(task_tenant_id),
+                task_dataset_id
+            )
+
+            if self._task_context.has_canceled_func(task_id):
+                self._task_context.progress_cb(-1, msg="Task has been canceled.")
+                return False
+
+        return True
+
+    async def _intercept_doc_store_delete(self, condition: dict, index_name: str, task_dataset_id: str) -> Any:
+        if self._task_context.write_interceptor:
+            return self._task_context.write_interceptor.intercept("docStoreConn.delete")
+        else:
+            return await thread_pool_exec(settings.docStoreConn.delete, condition, index_name, task_dataset_id)
+
+    async def _intercept_doc_store_insert(self, chunks: list, index_name: str, task_dataset_id: str) -> Any:
+        if self._task_context.write_interceptor:
+            if self._task_context.doc_id == GRAPH_RAPTOR_FAKE_DOC_ID: # raptor - non-determinisic
+                return self._task_context.write_interceptor.intercept("docStoreConn.insert", [])
+            return self._task_context.write_interceptor.intercept("docStoreConn.insert")
+        else:
+            return await thread_pool_exec(settings.docStoreConn.insert, chunks, index_name, task_dataset_id)
+
+    async def _insert_main_chunks(
+        self,
+        task_id: str,
+        task_tenant_id: str,
+        task_dataset_id: str,
+        chunks: List[Dict],
+        doc_bulk_size: int,
+    ) -> bool:
+        """Insert main chunks in batches with cancellation handling."""
+        for b in range(0, len(chunks), doc_bulk_size):
+            doc_store_result = await self._intercept_doc_store_insert(
+                chunks[b:b + doc_bulk_size],
+                search.index_name(task_tenant_id),
+                task_dataset_id
+            )
+
+            if self._task_context.has_canceled_func(task_id):
+                # Roll back partial RAPTOR summary inserts
+                await self._rollback_raptor_chunks(
+                    task_id, task_tenant_id, task_dataset_id, chunks, b, doc_bulk_size
+                )
+                self._task_context.progress_cb(-1, msg="Task has been canceled.")
+                return False
+
+            if b % 128 == 0:
+                self._task_context.progress_cb(prog=0.8 + 0.1 * (b + 1) / len(chunks),msg="")
+
+            if doc_store_result:
+                error_message = (
+                    f"Insert chunk error: {doc_store_result}, "
+                    "please check log file and Elasticsearch/Infinity status!"
+                )
+                self._task_context.progress_cb(-1, msg=error_message)
+                raise Exception(error_message)
+
+            # Update chunk IDs in task
+            chunk_ids = [chunk["id"] for chunk in chunks[:b + doc_bulk_size]]
+            if not await self._update_task_chunk_ids(task_id, chunk_ids):
+                # Roll back on failure
+                await self._rollback_insertion(task_tenant_id, task_dataset_id, chunk_ids)
+                self._task_context.progress_cb(
+                    -1,
+                    msg=f"Chunk updates failed since task {task_id} is unknown."
+                )
+                return False
+
+        return True
+
+    async def _rollback_raptor_chunks(
+        self,
+        task_id: str,
+        task_tenant_id: str,
+        task_dataset_id: str,
+        chunks: List[Dict],
+        up_to_batch: int,
+        doc_bulk_size: int,
+    ):
+        """Roll back partial RAPTOR summary inserts after cancellation."""
+        raptor_ids = [
+            c["id"] for c in chunks[:up_to_batch + doc_bulk_size]
+            if c.get("raptor_kwd") == "raptor"
+        ]
+
+        if raptor_ids:
+            try:
+                await self._intercept_doc_store_delete(
+                    {"id": raptor_ids}, search.index_name(task_tenant_id), task_dataset_id
+                )
+                logging.info(
+                    "insert_chunks: rolled back %d partial RAPTOR chunks after cancellation (task=%s)",
+                    len(raptor_ids), task_id,
+                )
+            except Exception:
+                logging.exception(
+                    "insert_chunks: failed to roll back partial RAPTOR chunks after cancellation (task=%s)",
+                    task_id,
+                )
+
+    async def _update_task_chunk_ids(self, task_id: str, chunk_ids: List[str]) -> bool:
+        """Update chunk IDs in the task record."""
+        from peewee import DoesNotExist
+
+        try:
+            if self._task_context.write_interceptor:
+                if self._task_context.doc_id == GRAPH_RAPTOR_FAKE_DOC_ID:
+                    self._task_context.write_interceptor.intercept("TaskService.update_chunk_ids", True)
+                else:
+                    self._task_context.write_interceptor.intercept("TaskService.update_chunk_ids")
+            else:
+                TaskService.update_chunk_ids(task_id, " ".join(chunk_ids))
+            return True
+        except DoesNotExist:
+            logging.warning(f"do_handle_task update_chunk_ids failed since task {task_id} is unknown.")
+            return False
+
+    async def _rollback_insertion(
+        self,
+        task_tenant_id: str,
+        task_dataset_id: str,
+        chunk_ids: List[str],
+    ):
+        """Roll back an insertion by deleting chunks and images."""
+        await self._intercept_doc_store_delete(
+            {"id": chunk_ids}, search.index_name(task_tenant_id), task_dataset_id
+        )
+
+        # Delete associated images
+        tasks = []
+        for chunk_id in chunk_ids:
+            tasks.append(asyncio.create_task(self._delete_image(task_dataset_id, chunk_id)))
+
+        try:
+            await asyncio.gather(*tasks, return_exceptions=False)
+        except Exception as e:
+            logging.error(f"delete_image failed: {e}")
+            for t in tasks:
+                t.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            raise
+
+    async def _delete_image(self, kb_id: str, chunk_id: str):
+        """Delete a chunk's image from storage."""
+        try:
+            async with self._task_context.minio_limiter:
+                settings.STORAGE_IMPL.delete(kb_id, chunk_id)
+        except Exception:
+            logging.exception(f"Deleting image of chunk {chunk_id} got exception")
+            raise

rag/svr/task_executor_refactor/comparator.py

@@ -0,0 +1,570 @@
+#
+#  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.
+
+"""
+Comparison Logic Module.
+
+This module provides the [`ContextComparator`](rag/svr/task_executor_refactor/comparator.py:100) class, which compares
+intermediate results from two [`RecordingContext`](rag/svr/task_executor_refactor/recording_context.py:54) instances:
+one from production execution and one from dry-run execution.
+
+The comparison supports various data types with appropriate strategies:
+- Basic types (int, str, bool): Direct equality comparison
+- Float numbers: Configurable tolerance range
+- Lists: Length comparison + ID set comparison + full content comparison (all chunks)
+- Dicts: Key set comparison + recursive value comparison
+- None: Equality comparison
+"""
+
+import logging
+from typing import Any, List, Optional, Set
+
+from rag.svr.task_executor_refactor.recording_context import BaseRecordingContext
+from rag.svr.task_executor_refactor.report_generator import (
+    ComparisonResult,
+    ComparisonReport,
+)
+from rag.svr.task_executor_refactor.write_operation_interceptor import ALLOWED_METHOD_NAMES
+
+
+class ContextComparator:
+    """Compare two RecordingContext instances for intermediate results.
+
+    This class compares the recorded data from production execution against
+    dry-run execution, generating a detailed report of matches and mismatches.
+
+    Usage:
+        comparator = ContextComparator()
+        report = comparator.compare("task_123", ctx_production, ctx_dry_run)
+        print(report.summary())
+    """
+
+    # Default tolerance for float comparison
+    DEFAULT_FLOAT_TOLERANCE = 1e-6
+
+    # Keys to strip from dict values before comparison (non-deterministic values)
+    DICT_KEYS_TO_STRIP = {"seconds", "_created_time", "_elapsed_time"}
+
+    # Keys that represent counts and should be compared as numbers
+    COUNT_KEYS = {
+        "outline_entry_count",
+        "tags_applied_count",
+        "final_chunk_count",
+        "final_chunk_ids_count",
+        "chunk_count",
+        "chunk_ids_count",
+        "token_count",
+        "raptor_token_count",
+    }
+
+    # Keys that contain chunk data for comparison
+    CHUNK_KEYS = {
+        "toc_chunk",
+        "raw_chunks",
+        "final_chunks",
+        "chunks",
+        "raptor_chunks",
+        "docs_after_prep",
+        "dataflow_chunks",
+    }
+
+    def __init__(self, float_tolerance: float = None):
+        """Initialize the Comparator.
+
+        Args:
+            float_tolerance: Tolerance for float comparison.
+                Defaults to DEFAULT_FLOAT_TOLERANCE.
+        """
+        self.float_tolerance = self.DEFAULT_FLOAT_TOLERANCE if float_tolerance is None else float_tolerance
+
+    def _strip_non_deterministic_fields(self, data: dict) -> dict:
+        """Remove non-deterministic fields (like 'seconds') from dict values.
+
+        This creates a shallow copy of the data dict with specified keys
+        removed from any nested dict values.
+
+        Args:
+            data: The input dictionary to process.
+
+        Returns:
+            A new dictionary with non-deterministic fields removed.
+        """
+        import copy
+        result = copy.copy(data)
+        for key, value in result.items():
+            if isinstance(value, dict):
+                # Create a new dict without the non-deterministic keys
+                cleaned = {
+                    k: v for k, v in value.items()
+                    if k not in self.DICT_KEYS_TO_STRIP
+                }
+                result[key] = cleaned
+        return result
+
+    @staticmethod
+    def _get_key_values_to_compare(prod_data_all:dict):
+        prod_data = dict()
+        for key, value in prod_data_all.items():
+            if key in ALLOWED_METHOD_NAMES:
+                continue
+            if key.endswith("_time"):
+                continue
+            if key.startswith("settings.docStoreConn."):
+                continue
+            prod_data[key] = value
+        return prod_data
+
+    def compare(
+        self,
+        task_id: str,
+        ctx_production: BaseRecordingContext,
+        ctx_dry_run: BaseRecordingContext,
+        comparison_keys: List[str] = None,
+    ) -> ComparisonReport:
+        """Compare two RecordingContext instances.
+
+        Args:
+            task_id: The task identifier.
+            ctx_production: RecordingContext from production execution.
+            ctx_dry_run: RecordingContext from dry-run execution.
+            comparison_keys: Optional list of keys to compare.
+                If None, all keys from both contexts will be compared.
+
+        Returns:
+            A ComparisonReport with the comparison results.
+        """
+        report = ComparisonReport(task_id=task_id)
+
+        # Get all keys from both contexts
+        prod_data_all = ctx_production.get_all_func_return_values() if ctx_production else {}
+        prod_data = self._get_key_values_to_compare(prod_data_all)
+        dry_run_data_all = ctx_dry_run.get_all_func_return_values() if ctx_dry_run else {}
+        dry_run_data = self._get_key_values_to_compare(dry_run_data_all)
+
+        # Strip non-deterministic fields (like 'seconds') from dict values
+        prod_data = self._strip_non_deterministic_fields(prod_data)
+        dry_run_data = self._strip_non_deterministic_fields(dry_run_data)
+
+        # Determine keys to compare
+        if comparison_keys:
+            keys_to_compare = set(comparison_keys)
+        else:
+            keys_to_compare = set(prod_data.keys()) | set(dry_run_data.keys())
+
+        # Find missing keys
+        prod_keys = set(prod_data.keys())
+        dry_run_keys = set(dry_run_data.keys())
+
+        report.missing_in_production = sorted(dry_run_keys - prod_keys)
+        report.missing_in_dry_run = sorted(prod_keys - dry_run_keys)
+
+        # Compare each key
+        for key in sorted(keys_to_compare):
+            if key in prod_data and key in dry_run_data:
+                result = self.compare_value(key, prod_data[key], dry_run_data[key])
+                report.details.append(result)
+                if result.match:
+                    report.matched_keys += 1
+                else:
+                    report.mismatched_keys += 1
+                    logging.info(f"---prod:{prod_data[key]} diff with dry run:{dry_run_data[key]}")
+
+        report.total_keys = report.matched_keys + report.mismatched_keys
+        return report
+
+    def compare_value(
+        self,
+        key: str,
+        prod_value: Any,
+        dry_run_value: Any,
+    ) -> ComparisonResult:
+        """Compare a single value with appropriate strategy.
+
+        Args:
+            key: The key being compared.
+            prod_value: Value from production context.
+            dry_run_value: Value from dry-run context.
+
+        Returns:
+            A ComparisonResult with the comparison.
+        """
+        # Handle None cases
+        if prod_value is None and dry_run_value is None:
+            return ComparisonResult(key=key, match=True)
+        if prod_value is None or dry_run_value is None:
+            return ComparisonResult(
+                key=key,
+                match=False,
+                production_value=prod_value,
+                dry_run_value=dry_run_value,
+                diff_details="One value is None",
+            )
+
+        # Handle booleans
+        if isinstance(prod_value, bool) and isinstance(dry_run_value, bool):
+            match = prod_value == dry_run_value
+            return ComparisonResult(
+                key=key,
+                match=match,
+                production_value=prod_value,
+                dry_run_value=dry_run_value,
+                diff_details=None if match else "Boolean values differ",
+            )
+
+        # Handle lists (chunks)
+        if isinstance(prod_value, list) and isinstance(dry_run_value, list):
+            if key in self.CHUNK_KEYS:
+                return self._compare_chunks(key, prod_value, dry_run_value)
+            return self._compare_lists(key, prod_value, dry_run_value)
+
+        # Handle dicts
+        if isinstance(prod_value, dict) and isinstance(dry_run_value, dict):
+            return self._compare_dicts(key, prod_value, dry_run_value)
+
+        # Handle numbers
+        if isinstance(prod_value, (int, float)) and isinstance(dry_run_value, (int, float)):
+            return self._compare_numbers(key, prod_value, dry_run_value)
+
+        # Handle strings
+        if isinstance(prod_value, str) and isinstance(dry_run_value, str):
+            match = prod_value == dry_run_value
+            return ComparisonResult(
+                key=key,
+                match=match,
+                production_value=prod_value,
+                dry_run_value=dry_run_value,
+                diff_details=None if match else "String values differ",
+            )
+
+        # Default: try direct equality
+        match = prod_value == dry_run_value
+        return ComparisonResult(
+            key=key,
+            match=match,
+            production_value=prod_value,
+            dry_run_value=dry_run_value,
+            diff_details=None if match else "Values differ",
+        )
+
+    @classmethod
+    def _compare_lists(cls, key: str, prod_list: list, dry_run_list: list) -> ComparisonResult:
+        """Compare two lists.
+
+        Args:
+            key: The key being compared.
+            prod_list: List from production context.
+            dry_run_list: List from dry-run context.
+
+        Returns:
+            A ComparisonResult with the comparison.
+        """
+        if len(prod_list) != len(dry_run_list):
+            return ComparisonResult(
+                key=key,
+                match=False,
+                production_value=len(prod_list),
+                dry_run_value=len(dry_run_list),
+                diff_details=f"Length differs: {len(prod_list)} vs {len(dry_run_list)}",
+            )
+
+        # Try element-wise comparison
+        for i, (p, d) in enumerate(zip(prod_list, dry_run_list)):
+            if p != d:
+                return ComparisonResult(
+                    key=key,
+                    match=False,
+                    production_value=len(prod_list),
+                    dry_run_value=len(dry_run_list),
+                    diff_details=f"Element {i} differs",
+                )
+
+        return ComparisonResult(
+            key=key,
+            match=True,
+            production_value=len(prod_list),
+            dry_run_value=len(dry_run_list),
+        )
+
+    def _compare_chunks(
+        self,
+        key: str,
+        prod_chunks: list,
+        dry_run_chunks: list,
+    ) -> ComparisonResult:
+        """Compare chunk lists with multi-level strategy.
+
+        Comparison levels:
+        1. Length comparison
+        2. ID set comparison
+        3. Full content comparison (all chunks)
+
+        Args:
+            key: The key being compared.
+            prod_chunks: Chunks from production context.
+            dry_run_chunks: Chunks from dry-run context.
+
+        Returns:
+            A ComparisonResult with the comparison.
+        """
+        # Level 1: Length comparison
+        if len(prod_chunks) != len(dry_run_chunks):
+            return ComparisonResult(
+                key=key,
+                match=False,
+                production_value=len(prod_chunks),
+                dry_run_value=len(dry_run_chunks),
+                diff_details=f"Chunk count differs: {len(prod_chunks)} vs {len(dry_run_chunks)}",
+            )
+
+        # Level 2: ID set comparison
+        prod_ids = self._extract_chunk_ids(prod_chunks)
+        dry_run_ids = self._extract_chunk_ids(dry_run_chunks)
+
+        if prod_ids != dry_run_ids:
+            missing_ids = prod_ids - dry_run_ids
+            extra_ids = dry_run_ids - prod_ids
+            details = f"Chunk IDs differ, total prod:{len(prod_ids)}, dry run:{len(dry_run_ids)}"
+            if missing_ids:
+                details += f", missing in dry-run: {len(missing_ids)}"
+            if extra_ids:
+                details += f", extra in dry-run: {len(extra_ids)}"
+            return ComparisonResult(
+                key=key,
+                match=False,
+                production_value=len(prod_ids),
+                dry_run_value=len(dry_run_ids),
+                diff_details=details,
+            )
+
+        # Level 3: Full content comparison (all chunks)
+        content_diffs = self._compare_all_chunks(prod_chunks, dry_run_chunks)
+        if content_diffs:
+            return ComparisonResult(
+                key=key,
+                match=False,
+                production_value=len(prod_chunks),
+                dry_run_value=len(dry_run_chunks),
+                diff_details=f"Content differs in samples: {'; '.join(content_diffs[:3])}",
+            )
+
+        return ComparisonResult(
+            key=key,
+            match=True,
+            production_value=len(prod_chunks),
+            dry_run_value=len(dry_run_chunks),
+        )
+
+    def _compare_all_chunks(
+        self,
+        prod_chunks: list,
+        dry_run_chunks: list,
+    ) -> List[str]:
+        """Compare ALL chunks from both lists.
+
+        Args:
+            prod_chunks: Chunks from production context.
+            dry_run_chunks: Chunks from dry-run context.
+
+        Returns:
+            List of difference descriptions.
+        """
+        if not prod_chunks or not dry_run_chunks:
+            return []
+
+        diffs = []
+        n = len(prod_chunks)
+
+        # Check if chunks have valid IDs
+        prod_has_id = any(self._get_chunk_id(c) for c in prod_chunks)
+        dry_run_has_id = any(self._get_chunk_id(c) for c in dry_run_chunks)
+        use_index_matching = not prod_has_id or not dry_run_has_id
+
+        # Build index by chunk ID for matching (only if IDs are available)
+        if not use_index_matching:
+            dry_run_by_id = {self._get_chunk_id(c): c for c in dry_run_chunks}
+        else:
+            dry_run_by_id = None
+
+        # Compare ALL chunks
+        for idx in range(n):
+            prod_chunk = prod_chunks[idx]
+            chunk_id = self._get_chunk_id(prod_chunk)
+
+            if use_index_matching:
+                # Use index position for matching
+                if idx < len(dry_run_chunks):
+                    dry_run_chunk = dry_run_chunks[idx]
+                else:
+                    dry_run_chunk = None
+            else:
+                # Use ID for matching
+                dry_run_chunk = dry_run_by_id.get(chunk_id)
+
+            if dry_run_chunk is None:
+                diffs.append(f"Chunk {idx} (id={chunk_id}) not found in dry-run")
+                continue
+
+            # Compare content
+            content_diff = self._compare_chunk_content(prod_chunk, dry_run_chunk)
+            if content_diff:
+                diffs.append(f"Chunk {idx} (id={chunk_id}): {content_diff}")
+
+        return diffs
+
+    @classmethod
+    def _compare_chunk_content(cls, prod_chunk: dict, dry_run_chunk: dict) -> Optional[str]:
+        """Compare content of two chunks.
+
+        Args:
+            prod_chunk: Chunk from production context.
+            dry_run_chunk: Chunk from dry-run context.
+
+        Returns:
+            Difference description or None if matched.
+        """
+        # Compare key fields
+        key_fields = ["content_with_weight", "content_ltks", "doc_id", "kb_id"]
+        for fld in key_fields:
+            if prod_chunk.get(fld) != dry_run_chunk.get(fld):
+                return f"Field '{fld}' differs, prod_chunk:{prod_chunk.get(fld)}, dry_run_chunk:{dry_run_chunk}"
+
+        # Compare vector fields
+        prod_vec_keys = {k for k in prod_chunk if k.startswith("q_") and k.endswith("_vec")}
+        dry_run_vec_keys = {k for k in dry_run_chunk if k.startswith("q_") and k.endswith("_vec")}
+
+        if prod_vec_keys != dry_run_vec_keys:
+            return f"Vector fields differ: {prod_vec_keys} vs {dry_run_vec_keys}"
+
+        for vec_key in prod_vec_keys:
+            p_vec = prod_chunk.get(vec_key)
+            d_vec = dry_run_chunk.get(vec_key)
+            if p_vec != d_vec:
+                return f"Vector '{vec_key}' differs"
+
+        return None
+
+    @classmethod
+    def _extract_chunk_ids(cls, chunks: list) -> Set[str]:
+        """Extract chunk IDs from a list of chunks.
+
+        Args:
+            chunks: List of chunk dictionaries.
+
+        Returns:
+            Set of chunk IDs.
+        """
+        ids = set()
+        for c in chunks:
+            if isinstance(c, dict) and "id" in c:
+                ids.add(str(c["id"]))
+        return ids
+
+    @classmethod
+    def _get_chunk_id(cls, chunk: Any) -> str:
+        """Get chunk ID from a chunk dictionary.
+
+        Args:
+            chunk: A chunk dictionary.
+
+        Returns:
+            Chunk ID as string, or empty string if not found.
+        """
+        if isinstance(chunk, dict):
+            return str(chunk.get("id", ""))
+        return ""
+
+    @classmethod
+    def _compare_dicts(cls, key: str, prod_dict: dict, dry_run_dict: dict) -> ComparisonResult:
+        """Compare two dictionaries.
+
+        Args:
+            key: The key being compared.
+            prod_dict: Dict from production context.
+            dry_run_dict: Dict from dry-run context.
+
+        Returns:
+            A ComparisonResult with the comparison.
+        """
+        prod_keys = set(prod_dict.keys())
+        dry_run_keys = set(dry_run_dict.keys())
+
+        if prod_keys != dry_run_keys:
+            missing = prod_keys - dry_run_keys
+            extra = dry_run_keys - prod_keys
+            details = "Keys differ"
+            if missing:
+                details += f", missing in dry-run: {missing}"
+            if extra:
+                details += f", extra in dry-run: {extra}"
+            return ComparisonResult(
+                key=key,
+                match=False,
+                production_value=sorted(prod_keys),
+                dry_run_value=sorted(dry_run_keys),
+                diff_details=details,
+            )
+
+        # Compare values for each key
+        for k in prod_keys:
+            p_val = prod_dict[k]
+            d_val = dry_run_dict[k]
+            if p_val != d_val:
+                return ComparisonResult(
+                    key=key,
+                    match=False,
+                    production_value=prod_dict,
+                    dry_run_value=dry_run_dict,
+                    diff_details=f"Value for key '{k}' differs",
+                )
+
+        return ComparisonResult(
+            key=key,
+            match=True,
+            production_value=prod_dict,
+            dry_run_value=dry_run_dict,
+        )
+
+    def _compare_numbers(
+        self,
+        key: str,
+        prod_value: float,
+        dry_run_value: float,
+    ) -> ComparisonResult:
+        """Compare two numbers with tolerance.
+
+        Args:
+            key: The key being compared.
+            prod_value: Number from production context.
+            dry_run_value: Number from dry-run context.
+
+        Returns:
+            A ComparisonResult with the comparison.
+        """
+        diff = abs(prod_value - dry_run_value)
+        if diff <= self.float_tolerance:
+            return ComparisonResult(
+                key=key,
+                match=True,
+                production_value=prod_value,
+                dry_run_value=dry_run_value,
+            )
+
+        return ComparisonResult(
+            key=key,
+            match=False,
+            production_value=prod_value,
+            dry_run_value=dry_run_value,
+            diff_details=f"Difference {diff} exceeds tolerance {self.float_tolerance}",
+        )

rag/svr/task_executor_refactor/constants.py

@@ -0,0 +1,24 @@
+#
+#  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.
+
+"""
+Shared constants for task executor modules.
+
+This module exists to break circular imports between task_executor.py and
+task_executor_refactor modules.
+"""
+
+CANVAS_DEBUG_DOC_ID = "dataflow_x"
+GRAPH_RAPTOR_FAKE_DOC_ID = "graph_raptor_x"

rag/svr/task_executor_refactor/dataflow_service.py

@@ -0,0 +1,389 @@
+#
+#  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.
+
+"""
+Dataflow Service Module.
+
+Provides [`DataflowService`](rag/svr/task_executor_refactor/dataflow_service.py:42) for dataflow
+pipeline execution.
+"""
+
+import abc
+import copy
+import logging
+import re
+from datetime import datetime
+from timeit import default_timer as timer
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import xxhash
+from common import settings
+from rag.svr.task_executor_refactor.embedding_utils import EmbeddingUtils
+from rag.flow.pipeline import Pipeline
+
+from api.db.services.canvas_service import UserCanvasService
+from api.db.services.document_service import DocumentService
+from api.db.services.doc_metadata_service import DocMetadataService
+from api.db.services.pipeline_operation_log_service import PipelineOperationLogService
+from api.db.joint_services.tenant_model_service import get_model_config_by_type_and_name
+from common.constants import LLMType, PipelineTaskType
+from common.metadata_utils import update_metadata_to
+from common.misc_utils import thread_pool_exec
+from rag.nlp import rag_tokenizer, add_positions
+from rag.svr.task_executor_refactor.constants import CANVAS_DEBUG_DOC_ID
+from rag.svr.task_executor_refactor.task_context import TaskContext
+
+
+class BillingHook(abc.ABC):
+    """Abstract base for billing hooks on pipeline success/error.
+
+    Implementations override the no-op methods to integrate with billing
+    systems (e.g., consume quota on success, release hold on error).
+    """
+
+    async def on_pipeline_success(self) -> None:
+        """Called when the dataflow pipeline completes successfully."""
+
+    async def on_pipeline_error(self) -> None:
+        """Called when the dataflow pipeline encounters an error."""
+
+
+class DataflowService:
+    """Service for dataflow pipeline execution.
+
+    This service handles:
+    - Dataflow DSL loading and execution
+    - Chunk embedding for dataflow output
+    - Chunk metadata processing and indexing
+    """
+
+    def __init__(
+        self,
+        ctx: TaskContext,
+        billing_hook: Optional[BillingHook] = None,
+        embedding_batch_size: int = None,
+        doc_bulk_size: int = None,
+    ):
+        """Initialize DataflowService.
+
+        Args:
+            ctx: TaskContext containing task configuration and execution resources.
+            billing_hook: Optional billing hook for pipeline success/error callbacks.
+            embedding_batch_size: Batch size for embedding operations.
+            doc_bulk_size: Batch size for document store inserts.
+        """
+        self._task_context = ctx
+        self._billing_hook = billing_hook
+        self._embedding_batch_size = embedding_batch_size or self._get_default_embedding_batch_size()
+        self._doc_bulk_size = doc_bulk_size or self._get_default_bulk_size()
+
+    async def run_dataflow(self) -> None:
+        """Run a dataflow pipeline."""
+        ctx = self._task_context
+        pipeline = None
+        try:
+            task_start_ts = timer()
+            dataflow_id = ctx.dataflow_id
+            doc_id = ctx.doc_id
+            task_id = ctx.id
+            task_dataset_id = ctx.kb_id
+
+            # Load DSL
+            dsl = await self._load_dsl(dataflow_id)
+            if dsl is None:
+                return
+
+            # Run pipeline
+            pipeline = Pipeline(
+                dsl, tenant_id=ctx.tenant_id, doc_id=doc_id,
+                task_id=task_id, flow_id=dataflow_id
+            )
+            chunks = await pipeline.run(file=ctx.file) if ctx.file else await pipeline.run()
+
+            if doc_id == CANVAS_DEBUG_DOC_ID:
+                ctx.recording_context.record("dataflow_debug_result", "canvas_debug_mode")
+                ctx.recording_context.record("dataflow_chunks", chunks)
+                return
+
+            if not chunks:
+                ctx.recording_context.record("pipeline_output_count", 0)
+                ctx.recording_context.record("pipeline_output_type", "empty")
+                self._record_pipeline_log(doc_id, dataflow_id, pipeline)
+                return
+
+            embedding_token_consumption = chunks.get("embedding_token_consumption", 0)
+            output_type = DataflowService._get_output_type(chunks)
+            chunks = self._normalize_chunks(chunks)
+
+            ctx.recording_context.record("pipeline_output_type", output_type)
+            ctx.recording_context.record("pipeline_output_count", len(chunks))
+
+            if not chunks:
+                self._record_pipeline_log(doc_id, dataflow_id, pipeline)
+                return
+
+            # Embed chunks if needed
+            keys = [k for o in chunks for k in list(o.keys())]
+            if not any([re.match(r"q_[0-9]+_vec", k) for k in keys]):
+                chunks, embedding_token_consumption = await self._embed_chunks(
+                    chunks, embedding_token_consumption
+                )
+                if chunks is None:
+                    self._record_pipeline_log(doc_id, dataflow_id, pipeline)
+                    return
+
+            # Process chunks
+            metadata = self._process_chunks(chunks)
+
+            # Update document metadata
+            if metadata:
+                self._update_document_metadata(doc_id, metadata)
+
+            # Insert chunks
+            start_ts = timer()
+            self._progress(prog=0.82, msg="[DOC Engine]:\nStart to index...")
+            e = await self._insert_chunks(
+                task_id, ctx.tenant_id, ctx.kb_id, chunks
+            )
+            if not e:
+                self._record_pipeline_log(doc_id, dataflow_id, pipeline)
+                return
+
+            time_cost = timer() - start_ts
+            task_time_cost = timer() - task_start_ts
+            self._progress(
+                prog=1.,
+                msg="Indexing done ({:.2f}s). Task done ({:.2f}s)".format(time_cost, task_time_cost)
+            )
+
+            # Update document stats
+            if ctx.write_interceptor:
+                ctx.write_interceptor.intercept("DocumentService.increment_chunk_num")
+            else:
+                DocumentService.increment_chunk_num(
+                    doc_id, task_dataset_id, embedding_token_consumption, len(chunks), task_time_cost
+                )
+
+            logging.info(
+                "[Done], chunks({}), token({}), elapsed:{:.2f}".format(
+                    len(chunks), embedding_token_consumption, task_time_cost
+                )
+            )
+            ctx.recording_context.record("dataflow_chunks", chunks)
+            self._record_pipeline_log(doc_id, dataflow_id, pipeline)
+
+            # Billing hook: pipeline succeeded
+            if self._billing_hook:
+                await self._billing_hook.on_pipeline_success()
+        except Exception:
+            if self._billing_hook:
+                await self._billing_hook.on_pipeline_error()
+            raise
+
+    async def _load_dsl(self, dataflow_id: str) -> Optional[str]:
+        """Load dataflow DSL from service."""
+        ctx = self._task_context
+        if ctx.task_type == "dataflow":
+            e, cvs = UserCanvasService.get_by_id(dataflow_id)
+            assert e, "User pipeline not found."
+            return cvs.dsl
+        else:
+            e, pipeline_log = PipelineOperationLogService.get_by_id(dataflow_id)
+            assert e, "Pipeline log not found."
+            return pipeline_log.dsl
+
+    @staticmethod
+    def _get_output_type(chunks: Dict) -> str:
+        """Determine output type from chunks dict."""
+        if "chunks" in chunks:
+            return "chunks"
+        elif "json" in chunks:
+            return "json"
+        elif "markdown" in chunks:
+            return "markdown"
+        elif "text" in chunks:
+            return "text"
+        elif "html" in chunks:
+            return "html"
+        return "empty"
+
+    @classmethod
+    def _normalize_chunks(cls, chunks: Dict) -> List[Dict]:
+        """Normalize chunks from various output formats."""
+        if "chunks" in chunks:
+            return copy.deepcopy(chunks["chunks"])
+        elif "json" in chunks:
+            return copy.deepcopy(chunks["json"])
+        elif "markdown" in chunks:
+            return [{"text": [chunks["markdown"]]}] if chunks["markdown"] else []
+        elif "text" in chunks:
+            return [{"text": [chunks["text"]]}] if chunks["text"] else []
+        elif "html" in chunks:
+            return [{"text": [chunks["html"]]}] if chunks["html"] else []
+        return []
+
+    async def _embed_chunks(
+        self, chunks: List[Dict], token_consumption: int
+    ) -> Tuple[Optional[List[Dict]], int]:
+        """Embed chunks using the embedding model."""
+        ctx = self._task_context
+        try:
+            self._progress(prog=0.82, msg="\n-------------------------------------\nStart to embedding...")
+            e, kb = self._get_kb_by_id(ctx.kb_id)
+            embedding_id = kb.embd_id
+            embd_model_config = get_model_config_by_type_and_name(
+                ctx.tenant_id, LLMType.EMBEDDING, embedding_id
+            )
+            from api.db.services.llm_service import LLMBundle
+            with LLMBundle(ctx.tenant_id, embd_model_config) as embedding_model:
+
+                # Prepare texts for embedding using EmbeddingUtils
+                texts = EmbeddingUtils.prepare_texts_for_dataflow_embedding(chunks)
+                delta = 0.20 / (len(texts) // self._embedding_batch_size + 1)
+                prog = 0.8
+
+                # Batch encode using EmbeddingUtils
+                vects_batches = []
+                for i in range(0, len(texts), self._embedding_batch_size):
+                    batch = texts[i: i + self._embedding_batch_size]
+                    async with ctx.embed_limiter:
+                        vts, c = await thread_pool_exec(
+                            self._encode_batch, batch, embedding_model
+                        )
+                    vects_batches.append(vts)
+                    token_consumption += c
+                    prog += delta
+                    if i % (len(texts) // self._embedding_batch_size / 100 + 1) == 1:
+                        self._progress(
+                            prog=prog,
+                            msg=f"{i + 1} / {len(texts) // self._embedding_batch_size}"
+                        )
+
+                # Stack vectors using EmbeddingUtils
+                vects = EmbeddingUtils.stack_vectors(vects_batches)
+                if len(vects) != len(chunks):
+                    raise ValueError(f"Vector count mismatch: {len(vects)} vs {len(chunks)}")
+
+                # Attach vectors using EmbeddingUtils
+                EmbeddingUtils.attach_vectors(chunks, vects)
+
+                return chunks, token_consumption
+
+        except Exception as e:
+            ctx.progress_cb(prog=-1, msg=f"[ERROR]: {e}")
+            return None, token_consumption
+
+    @classmethod
+    async def _encode_batch(cls, txts: List[str], embedding_model) -> Tuple[np.ndarray, int]:
+        """Batch encode texts using the embedding model with truncation."""
+        truncated = EmbeddingUtils.truncate_texts(txts, embedding_model.max_length)
+        return embedding_model.encode(truncated)
+
+    def _process_chunks(self, chunks: List[Dict]) -> Dict:
+        """Process chunks for metadata and indexing."""
+        ctx = self._task_context
+        metadata = {}
+        for ck in chunks:
+            ck["doc_id"] = ctx.doc_id
+            ck["kb_id"] = [str(ctx.kb_id)]
+            ck["docnm_kwd"] = ctx.name
+            ck["create_time"] = str(datetime.now()).replace("T", " ")[:19]
+            ck["create_timestamp_flt"] = datetime.now().timestamp()
+
+            if not ck.get("id"):
+                ck["id"] = xxhash.xxh64((ck["text"] + str(ck["doc_id"])).encode("utf-8")).hexdigest()
+
+            if "questions" in ck:
+                if "question_tks" not in ck:
+                    ck["question_kwd"] = ck["questions"].split("\n")
+                    ck["question_tks"] = rag_tokenizer.tokenize(str(ck["questions"]))
+                del ck["questions"]
+
+            if "keywords" in ck:
+                if "important_tks" not in ck:
+                    ck["important_kwd"] = [k for k in re.split(r"[,，;；、\r\n]+", ck["keywords"]) if k.strip()]
+                    ck["important_tks"] = rag_tokenizer.tokenize(str(ck["keywords"]))
+                del ck["keywords"]
+
+            if "summary" in ck:
+                if "content_ltks" not in ck:
+                    ck["content_ltks"] = rag_tokenizer.tokenize(str(ck["summary"]))
+                    ck["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(ck["content_ltks"])
+                del ck["summary"]
+
+            if "metadata" in ck:
+                metadata = update_metadata_to(metadata, ck["metadata"])
+                del ck["metadata"]
+
+            if "content_with_weight" not in ck:
+                ck["content_with_weight"] = ck["text"]
+            del ck["text"]
+
+            if "positions" in ck:
+                add_positions(ck, ck["positions"])
+                del ck["positions"]
+
+        return metadata
+
+    def _update_document_metadata(self, doc_id: str, metadata: Dict) -> None:
+        """Update document metadata."""
+        existing_meta = DocMetadataService.get_document_metadata(doc_id)
+        existing_meta = existing_meta if isinstance(existing_meta, dict) else {}
+        metadata = update_metadata_to(metadata, existing_meta)
+        self._task_context.recording_context.record("run_dataflow_metadata", metadata)
+        if self._task_context.write_interceptor:
+            self._task_context.write_interceptor.intercept("DocMetadataService.update_document_metadata")
+        else:
+            DocMetadataService.update_document_metadata(doc_id, metadata)
+
+    async def _insert_chunks(
+        self, task_id: str, tenant_id: str, kb_id: str, chunks: List[Dict]
+    ) -> bool:
+        """Insert chunks into document store."""
+        from rag.svr.task_executor_refactor.chunk_service import ChunkService
+        chunk_service = ChunkService(self._task_context)
+        return await chunk_service.insert_chunks(task_id, tenant_id, kb_id, chunks)
+
+    def _record_pipeline_log(self, doc_id: str, dataflow_id: str, pipeline) -> None:
+        """Record pipeline operation log."""
+        if self._task_context.write_interceptor:
+            self._task_context.write_interceptor.intercept("PipelineOperationLogService.create")
+        else:
+            PipelineOperationLogService.create(
+                document_id=doc_id, pipeline_id=dataflow_id,
+                task_type=PipelineTaskType.PARSE, dsl=str(pipeline)
+            )
+
+    @classmethod
+    def _get_kb_by_id(cls, kb_id: str):
+        """Get knowledge base by ID."""
+        from api.db.services.knowledgebase_service import KnowledgebaseService
+        return KnowledgebaseService.get_by_id(kb_id)
+
+    def _progress(self, prog=None, msg=None):
+        """Progress callback helper."""
+        if prog is not None or msg is not None:
+            self._task_context.progress_cb(prog=prog, msg=msg)
+
+    @classmethod
+    def _get_default_embedding_batch_size(cls) -> int:
+        """Get default embedding batch size."""
+        return settings.EMBEDDING_BATCH_SIZE
+
+    @classmethod
+    def _get_default_bulk_size(cls) -> int:
+        """Get default bulk size."""
+        return settings.DOC_BULK_SIZE

rag/svr/task_executor_refactor/embedding_service.py

@@ -0,0 +1,127 @@
+#
+#  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.
+
+"""
+Embedding Service Module.
+
+Provides [`EmbeddingService`](rag/svr/task_executor_refactor/embedding_service.py:42) for vector embedding operations.
+"""
+
+import asyncio
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+from common import settings
+from rag.svr.task_executor_refactor.embedding_utils import EmbeddingUtils
+from rag.svr.task_executor_refactor.task_context import TaskContext
+
+
+class EmbeddingService:
+    """Service for vector embedding operations.
+
+    This service handles:
+    - Batch encoding of text chunks
+    - Title + content vector combination
+    - Embedding model rate limiting
+
+    All intermediate results are recorded via RecordingContext for comparison.
+    """
+
+    def __init__(
+        self,
+        ctx: TaskContext,
+        embedding_batch_size: int = None,
+    ):
+        """Initialize EmbeddingService.
+
+        Args:
+            ctx: TaskContext containing task configuration and execution resources.
+            embedding_batch_size: Batch size for embedding operations.
+        """
+        self._task_context = ctx
+
+        self._embedding_batch_size = embedding_batch_size or settings.EMBEDDING_BATCH_SIZE
+
+    def embed_chunks(
+        self,
+        docs: List[Dict[str, Any]],
+        embedding_model,
+        parser_config: Dict = None,
+    ) -> Tuple[int, int]:
+        """Embed a list of chunks.
+
+        Args:
+            docs: List of chunk dictionaries to embed.
+            embedding_model: The embedding model bundle (LLMBundle).
+            parser_config: Parser configuration for filename embedding weight.
+
+        Returns:
+            Tuple of (token_count, vector_size).
+        """
+        if parser_config is None:
+            parser_config = {}
+
+        # Prepare text for embedding using EmbeddingUtils
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs)
+
+        # Encode titles using EmbeddingUtils for truncation
+        tk_count = 0
+        if len(titles) > 0 and len(titles) == len(contents):
+            vts, c = self._encode_single([titles[0]], embedding_model)
+            tts = np.tile(vts[0], (len(contents), 1))
+            tk_count += c
+        else:
+            tts = None
+
+        # Batch encode contents using EmbeddingUtils
+        vects_batches = []
+        for i in range(0, len(contents), self._embedding_batch_size):
+            batch = contents[i: i + self._embedding_batch_size]
+            vts, c = self._encode_batch(batch, embedding_model)
+            vects_batches.append(vts)
+            tk_count += c
+            if self._task_context.progress_cb:
+                self._task_context.progress_cb(prog=0.7 + 0.2 * (i + 1) / len(contents), msg="")
+
+        # Stack vectors using EmbeddingUtils
+        cnts = EmbeddingUtils.stack_vectors(vects_batches)
+
+        # Combine title and content vectors using EmbeddingUtils
+        title_weight = parser_config.get("filename_embd_weight", EmbeddingUtils.DEFAULT_TITLE_WEIGHT)
+        vects = EmbeddingUtils.combine_title_content_vectors(tts, cnts, title_weight)
+
+        assert len(vects) == len(docs)
+
+        # Attach vectors to docs using EmbeddingUtils
+        vector_size = EmbeddingUtils.attach_vectors(docs, vects)
+
+        return tk_count, vector_size
+
+    def _encode_single(self, texts: List[str], model) -> Tuple[np.ndarray, int]:
+        """Encode a single batch of texts."""
+        return self._run_encode(texts, model)
+
+    def _encode_batch(self, texts: List[str], model) -> Tuple[np.ndarray, int]:
+        """Encode a batch of texts with rate limiting and truncation."""
+        # Use EmbeddingUtils for truncation
+        truncated = EmbeddingUtils.truncate_texts(texts, model.max_length)
+        return self._run_encode(truncated, model)
+
+    def _run_encode(self, texts: List[str], model) -> Tuple[np.ndarray, int]:
+        """Run encoding with rate limiting."""
+        async def _encode():
+            async with self._task_context.embed_limiter:
+                return model.encode(texts)
+        return asyncio.get_event_loop().run_until_complete(_encode())

rag/svr/task_executor_refactor/embedding_utils.py

@@ -0,0 +1,223 @@
+#
+#  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.
+
+"""
+Embedding Utils Module.
+
+Provides utility functions for vector embedding operations to avoid code duplication
+across different services (e.g., [`EmbeddingService`](rag/svr/task_executor_refactor/embedding_service.py),
+[`DataflowService`](rag/svr/task_executor_refactor/dataflow_service.py)).
+
+This module centralizes:
+- Batch encoding of texts with truncation
+- Vector stacking from multiple batches
+- Vector attachment to chunk dictionaries
+- Title and content vector combination with configurable weights
+"""
+
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+
+from common.token_utils import truncate
+
+
+class EmbeddingUtils:
+    """Utility class for common embedding operations.
+
+    This class provides static methods for:
+    - Preparing texts for embedding (title/content extraction, HTML normalization)
+    - Batch encoding with truncation
+    - Stacking vector batches
+    - Attaching vectors to chunk dictionaries
+    - Combining title and content vectors with weights
+    """
+
+    DEFAULT_TITLE_WEIGHT = 0.1
+    DEFAULT_TITLE_PLACEHOLDER = "Title"
+    CONTENT_PLACEHOLDER_FOR_WHITESPACE = "None"
+
+    @classmethod
+    def prepare_texts_for_embedding(
+        cls,
+        docs: List[Dict[str, Any]],
+        use_question_kwd: bool = True,
+    ) -> Tuple[List[str], List[str]]:
+        """Prepare title and content texts for embedding.
+
+        Extracts titles from 'docnm_kwd' field and contents from 'question_kwd'
+        (if available and use_question_kwd is True) or 'content_with_weight'.
+        Table HTML tags are normalized to spaces.
+
+        Args:
+            docs: List of chunk dictionaries.
+            use_question_kwd: Whether to use 'question_kwd' as content if available.
+
+        Returns:
+            Tuple of (titles, contents) lists.
+        """
+        titles = []
+        contents = []
+        for d in docs:
+            title = d.get("docnm_kwd", cls.DEFAULT_TITLE_PLACEHOLDER)
+            titles.append(title)
+
+            content = cls._extract_content(d, use_question_kwd=use_question_kwd)
+            content = cls._normalize_table_html(content)
+            content = cls._handle_whitespace(content)
+
+            contents.append(content)
+        return titles, contents
+
+    @classmethod
+    def prepare_texts_for_dataflow_embedding(
+        cls,
+        chunks: List[Dict[str, Any]],
+    ) -> List[str]:
+        """Prepare texts for dataflow embedding.
+
+        Extracts content from 'questions', 'summary', or 'text' fields
+        (in priority order).
+
+        Args:
+            chunks: List of chunk dictionaries from dataflow output.
+
+        Returns:
+            List of text strings for embedding.
+        """
+        texts = []
+        for chunk in chunks:
+            text = chunk.get("questions", chunk.get("summary", chunk.get("text", "")))
+            texts.append(text)
+        return texts
+
+    @classmethod
+    def truncate_texts(cls, texts: List[str], max_length: int) -> List[str]:
+        """Truncate texts to the specified maximum length.
+
+        Args:
+            texts: List of text strings to truncate.
+            max_length: Maximum length for each text (will subtract 10 for safety margin).
+
+        Returns:
+            List of truncated text strings.
+        """
+        safe_max_length = max_length - 10
+        return [truncate(text, safe_max_length) for text in texts]
+
+    @classmethod
+    def stack_vectors(cls, vects_batches: List[np.ndarray]) -> np.ndarray:
+        """Stack a list of vector batches into a single array.
+
+        Args:
+            vects_batches: List of numpy arrays from batch encoding.
+
+        Returns:
+            Stacked numpy array, or empty array if no batches provided.
+        """
+        return np.vstack(vects_batches) if vects_batches else np.array([])
+
+    @classmethod
+    def attach_vectors(
+        cls,
+        docs: List[Dict[str, Any]],
+        vectors: np.ndarray,
+        vector_key_template: str = "q_%d_vec",
+    ) -> int:
+        """Attach vectors to chunk dictionaries.
+
+        Args:
+            docs: List of chunk dictionaries to modify in-place.
+            vectors: Numpy array of vectors to attach.
+            vector_key_template: Format string for the vector key (default: "q_%d_vec").
+
+        Returns:
+            The size of each vector (assumes uniform size).
+        """
+        vector_size = 0
+        if len(vectors) != len(docs):
+            raise ValueError(f"vectors/docs length mismatch: {len(vectors)} != {len(docs)}")
+        for i, doc in enumerate(docs):
+            vector = vectors[i].tolist()
+            vector_size = len(vector)
+            key = vector_key_template % vector_size
+            doc[key] = vector
+        return vector_size
+
+    @classmethod
+    def combine_title_content_vectors(
+        cls,
+        title_vecs: Optional[np.ndarray],
+        content_vecs: np.ndarray,
+        title_weight: Optional[float] = None,
+    ) -> np.ndarray:
+        """Combine title and content vectors with a configurable weight.
+
+        Args:
+            title_vecs: Title embedding vectors (may be None).
+            content_vecs: Content embedding vectors.
+            title_weight: Weight for title vectors (0.0 to 1.0). Defaults to 0.1.
+
+        Returns:
+            Combined vector array. If title_vecs is None or shapes don't match,
+            returns content_vecs unchanged.
+        """
+        if title_weight is None:
+            title_weight = cls.DEFAULT_TITLE_WEIGHT
+        if not title_weight:
+            title_weight = cls.DEFAULT_TITLE_WEIGHT
+
+        if (
+            title_vecs is not None
+            and content_vecs.ndim == 2
+            and title_vecs.shape == content_vecs.shape
+        ):
+            return title_weight * title_vecs + (1 - title_weight) * content_vecs
+        return content_vecs
+
+    @classmethod
+    def _extract_content(
+        cls,
+        doc: Dict[str, Any],
+        use_question_kwd: bool = True,
+    ) -> str:
+        """Extract content from a chunk dictionary.
+
+        Priority: question_kwd (joined by newline) -> content_with_weight.
+        """
+        if use_question_kwd:
+            question_kwd = doc.get("question_kwd", [])
+            if question_kwd:
+                return "\n".join(question_kwd)
+        return doc.get("content_with_weight", "")
+
+    @classmethod
+    def _normalize_table_html(cls, text: str) -> str:
+        """Normalize table HTML tags to spaces.
+
+        Replaces table-related HTML tags (table, td, caption, tr, th) with spaces.
+        """
+        return re.sub(r"</?(table|td|caption|tr|th)( [^<>]{0,12})?>", " ", text)
+
+    @classmethod
+    def _handle_whitespace(cls, text: str) -> str:
+        """Replace whitespace-only content with a placeholder.
+
+        Prevents embedding models from receiving empty or meaningless input.
+        """
+        if not text.strip():
+            return cls.CONTENT_PLACEHOLDER_FOR_WHITESPACE
+        return text

rag/svr/task_executor_refactor/post_processor.py

@@ -0,0 +1,156 @@
+#
+#  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.
+
+"""
+Post Processor Module.
+
+Provides [`PostProcessor`](rag/svr/task_executor_refactor/post_processor.py:42) for post-indexing
+operations like table parser metadata aggregation and TOC insertion.
+"""
+
+import logging
+from typing import Dict, List, Optional
+
+from api.db.services.document_service import DocumentService
+from api.db.services.doc_metadata_service import DocMetadataService
+from common.metadata_utils import update_metadata_to
+from rag.svr.task_executor_refactor.task_context import TaskContext
+from rag.utils.table_es_metadata import (
+    aggregate_table_manual_doc_metadata,
+    merge_table_parser_config_from_kb,
+    table_parser_strip_doc_metadata_keys,
+)
+
+class PostProcessor:
+    """Service for post-indexing operations.
+
+    This service handles:
+    - Table parser metadata aggregation
+    - Document metadata updates
+    - TOC (Table of Contents) chunk insertion
+    """
+
+    def __init__(
+        self,
+        ctx: TaskContext,
+    ):
+        """Initialize PostProcessor.
+
+        Args:
+            ctx: TaskContext containing task configuration and execution resources.
+        """
+        self._task_context = ctx
+
+    async def process_table_parser_metadata(
+        self,
+        task_doc_id: str,
+        chunks: List[Dict],
+    ) -> None:
+        """Process table parser metadata aggregation.
+
+        Args:
+            task_doc_id: Document ID.
+            chunks: List of chunk dictionaries.
+        """
+        ctx = self._task_context
+        if ctx.parser_id.lower() != "table":
+            return
+
+        eff_pc = merge_table_parser_config_from_kb(ctx.raw_task)
+        logging.debug(
+            f"[TABLE_META_DEBUG] table post-index: table_column_mode={eff_pc.get('table_column_mode')!r}"
+        )
+
+        if eff_pc.get("table_column_mode") != "manual":
+            return
+
+        try:
+            agg = aggregate_table_manual_doc_metadata(chunks, ctx.raw_task)
+            logging.debug(f"[TABLE_META_DEBUG] aggregated metadata: {agg}")
+
+            strip_keys = table_parser_strip_doc_metadata_keys(eff_pc)
+            existing = DocMetadataService.get_document_metadata(task_doc_id)
+            existing = existing if isinstance(existing, dict) else {}
+
+            preserved = {k: v for k, v in existing.items() if k not in strip_keys}
+            merged = update_metadata_to(dict(preserved), agg)
+
+            logging.debug(
+                f"[TABLE_META_DEBUG] calling update_document_metadata for doc_id={task_doc_id}, "
+                f"meta_fields keys={list(merged.keys())}, "
+                f"table_strip_key_count={len(strip_keys)}, agg_keys={list(agg.keys())}"
+            )
+
+            try:
+                if self._task_context.write_interceptor:
+                    self._task_context.write_interceptor.intercept("DocMetadataService.update_document_metadata")
+                else:
+                    DocMetadataService.update_document_metadata(task_doc_id, merged)
+                logging.debug("[TABLE_META_DEBUG] update_document_metadata succeeded")
+            except Exception as ue:
+                logging.error(
+                    "update_document_metadata failed (table parser, doc_id=%s): %s",
+                    task_doc_id,
+                    ue,
+                    exc_info=True,
+                )
+        except Exception as e:
+            logging.exception(
+                "Table parser document metadata aggregation failed (doc_id=%s): %s",
+                task_doc_id,
+                e,
+            )
+
+    async def insert_toc_chunk(
+        self,
+        toc_chunk: Optional[Dict],
+        chunk_service,
+    ) -> bool:
+        """Insert TOC chunk into document store.
+
+        Args:
+            toc_chunk: TOC chunk dictionary or None.
+            chunk_service: ChunkService instance for chunk insertion.
+
+        Returns:
+            True if TOC chunk was inserted successfully, False otherwise.
+        """
+        ctx = self._task_context
+        if toc_chunk is None:
+            return False
+
+        if self._task_context.has_canceled_func(ctx.id):
+            self._task_context.progress_cb(-1, msg="Task has been canceled.")
+            return False
+
+        insert_result = await chunk_service.insert_chunks(ctx.id, ctx.tenant_id, ctx.kb_id, [toc_chunk])
+
+        if not insert_result:
+            self._task_context.recording_context.record("toc_inserted", False)
+            return False
+
+        self._task_context.recording_context.record("toc_inserted", True)
+
+        if self._task_context.write_interceptor:
+            self._task_context.write_interceptor.intercept("DocumentService.increment_chunk_num")
+        else:
+            DocumentService.increment_chunk_num(ctx.doc_id, ctx.kb_id, 0, 1, 0)
+
+        return True
+
+    def _progress(self, prog=None, msg=None):
+        """Progress callback helper."""
+        if prog is not None or msg is not None:
+            self._task_context.progress_cb(prog=prog, msg=msg)

rag/svr/task_executor_refactor/raptor_service.py

@@ -0,0 +1,468 @@
+#
+#  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.
+
+"""
+Raptor Service Module.
+
+Provides [`RaptorService`](rag/svr/task_executor_refactor/raptor_service.py:48) for RAPTOR
+(Recursive Abstractive Processing for Tree-Organized Retrieval) summary generation.
+"""
+
+import copy
+import logging
+import os
+from datetime import datetime
+from typing import Dict, List, Optional, Set, Tuple
+
+import numpy as np
+
+from api.db.services.document_service import DocumentService
+from api.db.services.task_service import GRAPH_RAPTOR_FAKE_DOC_ID
+from common import settings
+from common.constants import PAGERANK_FLD
+from common.misc_utils import thread_pool_exec
+from common.token_utils import num_tokens_from_string
+from rag.nlp import rag_tokenizer, search
+from rag.utils.raptor_utils import (
+    collect_raptor_chunk_ids,
+    collect_raptor_methods,
+    get_raptor_clustering_method,
+    get_raptor_tree_builder,
+    get_skip_reason,
+    make_raptor_summary_chunk_id,
+    should_skip_raptor,
+)
+from rag.svr.task_executor_refactor.task_context import TaskContext
+
+
+class RaptorService:
+    """Service for RAPTOR summary generation.
+
+    This service handles:
+    - RAPTOR chunk method detection (checkpoint)
+    - RAPTOR summary generation per document or dataset-level
+    - Stale RAPTOR chunk cleanup
+    - Auto-disable rules for certain file types
+    """
+
+    def __init__(
+        self,
+        ctx: TaskContext,
+    ):
+        """Initialize RaptorService.
+
+        Args:
+            ctx: TaskContext containing task configuration and execution resources.
+        """
+        self._task_context = ctx
+
+    async def run_raptor_for_kb(
+        self,
+        kb_parser_config: Dict,
+        chat_mdl,
+        embd_mdl,
+        vector_size: int,
+        doc_ids: List[str],
+    ) -> Tuple[List[Dict], int, List[Tuple[str, Optional[str]]]]:
+        """Generate RAPTOR summaries for selected documents.
+
+        Args:
+            kb_parser_config: Knowledge base parser configuration.
+            chat_mdl: Chat model bundle for RAPTOR.
+            embd_mdl: Embedding model bundle for RAPTOR.
+            vector_size: Vector dimension size.
+            doc_ids: List of document IDs to process.
+
+        Returns:
+            Tuple of (chunks, token_count, cleanup_raptor_chunks).
+        """
+        raptor_config = kb_parser_config.get("raptor", {})
+        tree_builder = get_raptor_tree_builder(raptor_config)
+        clustering_method = get_raptor_clustering_method(raptor_config)
+        vctr_nm = "q_%d_vec" % vector_size
+
+        res = []
+        tk_count = 0
+        cleanup_raptor_chunks = []
+        max_errors = int(os.environ.get("RAPTOR_MAX_ERRORS", 3))
+
+        # Collect document info
+        doc_info_by_id = self._collect_doc_info(doc_ids)
+
+        # Determine scope
+        if raptor_config.get("scope", "file") == "file":
+            res, tk_count = await self._run_file_level_raptor(
+                raptor_config, tree_builder, clustering_method,
+                chat_mdl, embd_mdl, vctr_nm, doc_ids, doc_info_by_id,
+                max_errors, res, tk_count, cleanup_raptor_chunks
+            )
+        else:
+            res, tk_count = await self._run_dataset_level_raptor(
+                raptor_config, tree_builder, clustering_method,
+                chat_mdl, embd_mdl, vctr_nm, doc_ids, doc_info_by_id,
+                max_errors, res, tk_count, cleanup_raptor_chunks
+            )
+
+        return res, tk_count, cleanup_raptor_chunks
+
+    @classmethod
+    def _collect_doc_info(cls, doc_ids: List[str]) -> Dict[str, Dict]:
+        """Collect document info for all doc_ids."""
+        doc_info_by_id = {}
+        for doc_id in set(doc_ids):
+            ok, source_doc = DocumentService.get_by_id(doc_id)
+            if not ok or not source_doc:
+                continue
+            doc_info_by_id[doc_id] = {
+                "name": getattr(source_doc, "name", ""),
+                "type": getattr(source_doc, "type", ""),
+                "parser_id": getattr(source_doc, "parser_id", ""),
+                "parser_config": getattr(source_doc, "parser_config", {}) or {},
+            }
+        return doc_info_by_id
+
+    async def _run_file_level_raptor(
+        self, raptor_config, tree_builder, clustering_method,
+        chat_mdl, embd_mdl, vctr_nm, doc_ids, doc_info_by_id,
+        max_errors, res, tk_count, cleanup_raptor_chunks
+    ):
+        """Run RAPTOR at file level (per document)."""
+        ctx = self._task_context
+        fake_doc_id = GRAPH_RAPTOR_FAKE_DOC_ID
+        if self._task_context.write_interceptor: # dry run mode
+            dataset_methods = set()
+        else:
+            dataset_methods = await self._get_raptor_chunk_methods(fake_doc_id, ctx.tenant_id, ctx.kb_id)
+        remove_dataset_summaries = bool(dataset_methods)
+        has_file_level_target = False
+
+        if dataset_methods:
+            self._task_context.progress_cb(msg="[RAPTOR] will remove dataset-level summaries after file-level summaries are available.")
+
+        for x, doc_id in enumerate(doc_ids):
+            if self._should_skip_raptor(doc_id, doc_info_by_id, raptor_config):
+                self._task_context.progress_cb(prog=(x + 1.) / len(doc_ids))
+                continue
+            if self._task_context.write_interceptor:
+                existing_methods = set()
+            else:
+                existing_methods = await self._get_raptor_chunk_methods(doc_id, ctx.tenant_id, ctx.kb_id)
+            if tree_builder in existing_methods:
+                has_file_level_target = True
+                if existing_methods != {tree_builder}:
+                    self._schedule_raptor_cleanup(
+                        doc_id, tree_builder, cleanup_raptor_chunks
+                    )
+                    self._task_context.progress_cb(msg=f"[RAPTOR] doc:{doc_id} will remove old RAPTOR summaries after insert.")
+                    self._task_context.progress_cb(msg=f"[RAPTOR] doc:{doc_id} already has {tree_builder} RAPTOR chunks, skipping.")
+                    self._task_context.progress_cb(prog=(x + 1.) / len(doc_ids))
+                continue
+
+            if existing_methods:
+                self._task_context.progress_cb(msg=f"[RAPTOR] doc:{doc_id} will migrate RAPTOR summaries to {tree_builder} after insert.")
+
+            chunks = self._load_doc_chunks(doc_id, vctr_nm)
+            if not chunks:
+                continue
+
+            before_generate = len(res)
+            new_chunks, new_tk_count = await self._generate_raptor(
+                chunks, doc_id, raptor_config, chat_mdl, embd_mdl,
+                tree_builder, clustering_method, max_errors, doc_info_by_id
+            )
+            res.extend(new_chunks)
+            tk_count += new_tk_count
+
+            if len(res) > before_generate:
+                has_file_level_target = True
+                if existing_methods:
+                    self._schedule_raptor_cleanup(
+                        doc_id, tree_builder, cleanup_raptor_chunks
+                    )
+            self._task_context.progress_cb(prog=(x + 1.) / len(doc_ids))
+
+        if remove_dataset_summaries:
+            if has_file_level_target:
+                self._schedule_raptor_cleanup(
+                    fake_doc_id, None, cleanup_raptor_chunks
+                )
+            else:
+                self._task_context.progress_cb(msg="[RAPTOR] kept dataset-level summaries because no file-level summaries were built.")
+
+        return res, tk_count
+
+    async def _run_dataset_level_raptor(
+        self, raptor_config, tree_builder, clustering_method,
+        chat_mdl, embd_mdl, vctr_nm, doc_ids, doc_info_by_id,
+        max_errors, res, tk_count, cleanup_raptor_chunks
+    ):
+        """Run RAPTOR at dataset level (all documents combined)."""
+        ctx = self._task_context
+        fake_doc_id = GRAPH_RAPTOR_FAKE_DOC_ID
+        migrated_file_docs = 0
+        file_cleanup_doc_ids = []
+        skipped_doc_ids = set()
+
+        for doc_id in set(doc_ids):
+            if self._should_skip_raptor(doc_id, doc_info_by_id, raptor_config):
+                skipped_doc_ids.add(doc_id)
+                continue
+            if self._task_context.write_interceptor:
+                existing_methods = set()
+            else:
+                existing_methods = await self._get_raptor_chunk_methods(doc_id, ctx.tenant_id, ctx.kb_id)
+            if existing_methods:
+                file_cleanup_doc_ids.append(doc_id)
+                migrated_file_docs += 1
+
+        if migrated_file_docs:
+            self._task_context.progress_cb(
+                msg=f"[RAPTOR] will remove file-level summaries for {migrated_file_docs} docs after dataset-level build succeeds."
+            )
+
+        if self._task_context.write_interceptor:
+            existing_methods = set()
+        else:
+            existing_methods = await self._get_raptor_chunk_methods(fake_doc_id, ctx.tenant_id, ctx.kb_id)
+        if tree_builder in existing_methods:
+            if existing_methods != {tree_builder}:
+                self._schedule_raptor_cleanup(
+                    fake_doc_id, tree_builder, cleanup_raptor_chunks
+                )
+                self._task_context.progress_cb(msg="[RAPTOR] will remove old dataset-level RAPTOR summaries after insert.")
+            for doc_id in file_cleanup_doc_ids:
+                self._schedule_raptor_cleanup(doc_id, None, cleanup_raptor_chunks)
+            self._task_context.progress_cb(msg=f"[RAPTOR] dataset-level {tree_builder} summaries already exist, skipping.")
+            return res, tk_count
+
+        migrate_dataset_summaries = bool(existing_methods)
+        if migrate_dataset_summaries:
+            self._task_context.progress_cb(msg=f"[RAPTOR] will migrate dataset-level RAPTOR summaries to {tree_builder} after insert.")
+
+        chunks = self._load_all_doc_chunks(doc_ids, vctr_nm, skipped_doc_ids)
+        if not chunks:
+            if skipped_doc_ids and len(skipped_doc_ids) == len(set(doc_ids)):
+                self._task_context.progress_cb(msg="[RAPTOR] all documents were skipped by RAPTOR auto-disable rules.")
+                return res, tk_count
+            self._task_context.progress_cb(msg="[ERROR] No valid chunks with vectors found. Please ensure documents are parsed with the current embedding model.")
+            return res, tk_count
+
+        before_generate = len(res)
+        new_chunks, new_tk_count = await self._generate_raptor(
+            chunks, fake_doc_id, raptor_config, chat_mdl, embd_mdl,
+            tree_builder, clustering_method, max_errors, doc_info_by_id
+        )
+        res.extend(new_chunks)
+        tk_count += new_tk_count
+
+        if len(res) > before_generate:
+            for doc_id in file_cleanup_doc_ids:
+                self._schedule_raptor_cleanup(doc_id, None, cleanup_raptor_chunks)
+            if migrate_dataset_summaries:
+                self._schedule_raptor_cleanup(
+                    fake_doc_id, tree_builder, cleanup_raptor_chunks
+                )
+
+        return res, tk_count
+
+    def _should_skip_raptor(
+        self, doc_id: str, doc_info_by_id: Dict, raptor_config: Dict
+    ) -> bool:
+        """Check if RAPTOR should be skipped for a document."""
+        ctx = self._task_context
+        doc_info = doc_info_by_id.get(doc_id, {})
+        file_type = doc_info.get("type") or ctx.raw_task.get("type", "")
+        parser_id = doc_info.get("parser_id") or ctx.parser_id
+        parser_config = doc_info.get("parser_config") or ctx.parser_config
+
+        if should_skip_raptor(file_type, parser_id, parser_config, raptor_config):
+            skip_reason = get_skip_reason(file_type, parser_id, parser_config)
+            doc_name = doc_info.get("name") or doc_id
+            logging.info("Skipping Raptor for document %s: %s", doc_name, skip_reason)
+            self._task_context.progress_cb(msg=f"[RAPTOR] doc:{doc_id} skipped: {skip_reason}")
+            return True
+        return False
+
+    def _load_doc_chunks(self, doc_id: str, vctr_nm: str) -> List[Tuple[str, np.ndarray]]:
+        """Load chunks for a single document."""
+        ctx = self._task_context
+        chunks = []
+        skipped_chunks = 0
+
+        fields = ["content_with_weight", vctr_nm]
+        for d in settings.retriever.chunk_list(
+            doc_id, ctx.tenant_id, [str(ctx.kb_id)],
+            fields=fields,
+            sort_by_position=True
+        ):
+            if vctr_nm not in d or d[vctr_nm] is None:
+                skipped_chunks += 1
+                logging.warning(f"RAPTOR: Chunk missing vector field '{vctr_nm}' in doc {doc_id}, skipping")
+                continue
+            chunks.append((d["content_with_weight"], np.array(d[vctr_nm])))
+
+        if skipped_chunks > 0:
+            self._task_context.progress_cb(
+                msg=f"[WARN] Skipped {skipped_chunks} chunks without vector field '{vctr_nm}' for doc {doc_id}."
+            )
+        if not chunks:
+            logging.warning(f"RAPTOR: No valid chunks with vectors found for doc {doc_id}")
+            self._task_context.progress_cb(msg=f"[WARN] No valid chunks with vectors found for doc {doc_id}, skipping")
+
+        return chunks
+
+    def _load_all_doc_chunks(
+        self, doc_ids: List[str], vctr_nm: str, skipped_doc_ids: Set[str]
+    ) -> List[Tuple[str, np.ndarray]]:
+        """Load chunks for all documents."""
+        ctx = self._task_context
+        chunks = []
+        skipped_chunks = 0
+
+        fields = ["content_with_weight", vctr_nm]
+        for doc_id in doc_ids:
+            if doc_id in skipped_doc_ids:
+                continue
+            for d in settings.retriever.chunk_list(
+                doc_id, ctx.tenant_id, [str(ctx.kb_id)],
+                fields=fields,
+                sort_by_position=True
+            ):
+                if vctr_nm not in d or d[vctr_nm] is None:
+                    skipped_chunks += 1
+                    logging.warning(f"RAPTOR: Chunk missing vector field '{vctr_nm}' in doc {doc_id}, skipping")
+                    continue
+                chunks.append((d["content_with_weight"], np.array(d[vctr_nm])))
+
+        if skipped_chunks > 0:
+            self._task_context.progress_cb(
+                msg=f"[WARN] Skipped {skipped_chunks} chunks without vector field '{vctr_nm}'."
+            )
+
+        return chunks
+
+    async def _generate_raptor(
+        self,
+        chunks: List[Tuple[str, np.ndarray]],
+        doc_id: str,
+        raptor_config: Dict,
+        chat_mdl,
+        embd_mdl,
+        tree_builder: str,
+        clustering_method: str,
+        max_errors: int,
+        doc_info_by_id: Dict,
+    ) -> Tuple[List[Dict], int]:
+        """Run RAPTOR and generate summary chunks."""
+        ctx = self._task_context
+        from rag.raptor import RecursiveAbstractiveProcessing4TreeOrganizedRetrieval as Raptor
+
+        raptor_ext_config = raptor_config.get("ext") or {}
+        vctr_nm = "q_%d_vec" % len(chunks[0][1]) if chunks else "q_768_vec"
+
+        raptor = Raptor(
+            raptor_config.get("max_cluster", 64),
+            chat_mdl,
+            embd_mdl,
+            raptor_config["prompt"],
+            raptor_config["max_token"],
+            raptor_config["threshold"],
+            max_errors=max_errors,
+            tree_builder=tree_builder,
+            clustering_method=clustering_method,
+            psi_exact_max_leaves=raptor_ext_config.get("psi_exact_max_leaves", 4096),
+            psi_bucket_size=raptor_ext_config.get("psi_bucket_size", 1024),
+        )
+
+        original_length = len(chunks)
+        processed_chunks, layers = await raptor(
+            chunks, raptor_config["random_seed"], self._task_context.progress_cb, ctx.id
+        )
+
+        effective_doc_name = ctx.name if doc_id == GRAPH_RAPTOR_FAKE_DOC_ID else doc_info_by_id.get(doc_id, {}).get("name") or ctx.name
+
+        doc = {
+            "doc_id": doc_id,
+            "kb_id": [str(ctx.kb_id)],
+            "docnm_kwd": effective_doc_name,
+            "title_tks": rag_tokenizer.tokenize(effective_doc_name),
+            "raptor_kwd": "raptor",
+            "extra": {"raptor_method": tree_builder},
+        }
+        if ctx.pagerank:
+            doc[PAGERANK_FLD] = int(ctx.pagerank)
+
+        # Build index→layer mapping
+        chunk_layer = {}
+        for layer_idx, (layer_start, layer_end) in enumerate(layers):
+            if layer_idx == 0:
+                continue
+            for ci in range(layer_start, layer_end):
+                chunk_layer[ci] = layer_idx
+
+        res = []
+        tk_count = 0
+        for idx, (content, vctr) in enumerate(processed_chunks[original_length:], start=original_length):
+            d = copy.deepcopy(doc)
+            d["id"] = make_raptor_summary_chunk_id(content, doc_id)
+            d["create_time"] = str(datetime.now()).replace("T", " ")[:19]
+            d["create_timestamp_flt"] = datetime.now().timestamp()
+            d[vctr_nm] = vctr.tolist()
+            d["content_with_weight"] = content
+            d["content_ltks"] = rag_tokenizer.tokenize(content)
+            d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
+            d["raptor_layer_int"] = chunk_layer.get(idx, 1)
+            res.append(d)
+            tk_count += num_tokens_from_string(content)
+
+        return res, tk_count
+
+    @classmethod
+    def _schedule_raptor_cleanup(cls, doc_id: str, keep_method: Optional[str], cleanup_list: List):
+        """Queue stale RAPTOR summaries for deletion."""
+        cleanup_plan = (doc_id, keep_method)
+        if cleanup_plan not in cleanup_list:
+            cleanup_list.append(cleanup_plan)
+
+    @classmethod
+    async def _get_raptor_chunk_methods(cls, doc_id: str, tenant_id: str, kb_id: str) -> Set[str]:
+        """Get RAPTOR chunk methods for a document."""
+        from common.doc_store.doc_store_base import OrderByExpr
+
+        async def search_fields(fields: list, condition: dict, order_by=None):
+            res = await thread_pool_exec(
+                settings.docStoreConn.search,
+                fields, [], condition, [], order_by or OrderByExpr(),
+                0, 10000, search.index_name(tenant_id), [kb_id]
+            )
+            return settings.docStoreConn.get_fields(res, fields)
+
+        try:
+            primary = await search_fields(
+                ["raptor_kwd", "extra"], {"doc_id": doc_id, "raptor_kwd": ["raptor"]}
+            )
+            if collect_raptor_chunk_ids(primary):
+                return collect_raptor_methods(primary)
+
+            return collect_raptor_methods(
+                await search_fields(
+                    ["raptor_kwd", "extra"],
+                    {"doc_id": doc_id},
+                    OrderByExpr().desc("create_timestamp_flt"),
+                )
+            )
+        except Exception:
+            logging.exception("Failed to check RAPTOR chunks for doc %s", doc_id)
+            raise

rag/svr/task_executor_refactor/raptor_utils.py

@@ -0,0 +1,97 @@
+#
+#  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.
+
+"""
+RAPTOR chunk management utilities.
+
+Provides functions for managing RAPTOR summary chunks,
+including detection, retrieval, and deletion.
+"""
+
+import logging
+
+from common.misc_utils import thread_pool_exec
+from common import settings
+from rag.nlp import search as nlp_search
+from rag.utils.raptor_utils import (
+    collect_raptor_chunk_ids,
+)
+
+RAPTOR_METHOD_SEARCH_LIMIT = 10000
+
+
+async def get_raptor_chunk_field_map(doc_id: str, tenant_id: str, kb_id: str) -> dict:
+    """Return stored RAPTOR marker fields for a document."""
+    from common.doc_store.doc_store_base import OrderByExpr
+
+    async def search_fields(fields: list[str], condition: dict, order_by=None):
+        """Search chunk fields in the current knowledge base."""
+        res = await thread_pool_exec(
+            settings.docStoreConn.search,
+            fields, [], condition, [], order_by or OrderByExpr(),
+            0, RAPTOR_METHOD_SEARCH_LIMIT, nlp_search.index_name(tenant_id), [kb_id]
+        )
+        return settings.docStoreConn.get_fields(res, fields)
+
+    primary = await search_fields(["raptor_kwd", "extra"], {"doc_id": doc_id, "raptor_kwd": ["raptor"]})
+    if collect_raptor_chunk_ids(primary):
+        return primary
+
+    try:
+        return await search_fields(
+            ["raptor_kwd", "extra"],
+            {"doc_id": doc_id},
+            OrderByExpr().desc("create_timestamp_flt"),
+        )
+    except Exception:
+        logging.debug("RAPTOR fallback method lookup with extra field failed for doc %s", doc_id, exc_info=True)
+        return primary
+
+
+async def delete_raptor_chunks(doc_id: str, tenant_id: str, kb_id: str, keep_method: str | None = None) -> int:
+    """Delete RAPTOR summaries for doc_id, optionally preserving one method."""
+    if keep_method is None:
+        logging.info(
+            "delete_raptor_chunks: removing all RAPTOR summaries (doc=%s tenant=%s kb=%s)",
+            doc_id, tenant_id, kb_id,
+        )
+        await thread_pool_exec(
+            settings.docStoreConn.delete,
+            {"doc_id": doc_id, "raptor_kwd": ["raptor"]},
+            nlp_search.index_name(tenant_id),
+            kb_id,
+        )
+        return 0
+
+    field_map = await get_raptor_chunk_field_map(doc_id, tenant_id, kb_id)
+    chunk_ids = collect_raptor_chunk_ids(field_map, exclude_methods={keep_method})
+    if not chunk_ids:
+        logging.debug(
+            "delete_raptor_chunks: no stale RAPTOR chunks to remove (doc=%s tenant=%s kb=%s keep=%s)",
+            doc_id, tenant_id, kb_id, keep_method,
+        )
+        return 0
+
+    logging.info(
+        "delete_raptor_chunks: removing %d stale RAPTOR chunks (doc=%s tenant=%s kb=%s keep=%s)",
+        len(chunk_ids), doc_id, tenant_id, kb_id, keep_method,
+    )
+    await thread_pool_exec(
+        settings.docStoreConn.delete,
+        {"id": list(chunk_ids)},
+        nlp_search.index_name(tenant_id),
+        kb_id,
+    )
+    return len(chunk_ids)

rag/svr/task_executor_refactor/recording_context.py

@@ -0,0 +1,419 @@
+#
+#  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.
+
+"""Recording Context Module.
+
+This module provides the [`BaseRecordingContext`](rag/svr/task_executor_refactor/recording_context.py:48) abstract base class,
+[`RecordingContext`](rag/svr/task_executor_refactor/recording_context.py:89) concrete class, and
+[`NullRecordingContext`](rag/svr/task_executor_refactor/recording_context.py:204) no-op class, which capture
+actual execution results from the production code path (e.g., [`do_handle_task()`](rag/svr/task_executor.py))
+for later comparison with dry-run results.
+
+The recording context is used throughout the task execution pipeline to collect
+intermediate metrics and final results at various stages:
+
+1. **File validation**: Records file size check results and parser ID
+2. **Chunking**: Records raw chunks after document splitting
+3. **Outline extraction**: Records whether outline was extracted and entry count
+4. **MinIO upload**: Records document count after image upload
+5. **Post-processing**: Records counts for keywords, questions, metadata, and tags
+6. **Final results**: Records final chunks and their IDs for comparison
+
+The module also provides context variable management functions and a timing
+decorator that automatically integrates with the current recording context.
+
+Usage example::
+
+    from rag.svr.task_executor_refactor.recording_context import RecordingContext
+
+    ctx = RecordingContext()
+    ctx.record("raw_chunk_count", 42)
+    ctx.record("final_chunks", chunks)
+
+    # Later, in comparison:
+    comparator.compare(task_id, ctx, dry_run_records)
+"""
+
+import contextvars
+import functools
+import time
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from typing import Any, Callable, Dict, List, Tuple
+
+
+class BaseRecordingContext(ABC):
+    """Abstract base class for recording context implementations.
+
+    Defines the common interface shared by
+    [`RecordingContext`](rag/svr/task_executor_refactor/recording_context.py:89) and
+    [`NullRecordingContext`](rag/svr/task_executor_refactor/recording_context.py:204).
+
+    Variables typed as ``BaseRecordingContext`` can hold either implementation,
+    enabling production/dry-run polymorphism without conditional branches.
+    """
+
+    @abstractmethod
+    def record(self, key: str, value: Any) -> None:
+        """Record a value with the given key."""
+
+    @abstractmethod
+    def save_func_return_value(self, func_name: str, return_value: Any) -> None:
+        """Record a function's return value into a list associated with func_name."""
+
+    @abstractmethod
+    def get_func_return_values(self, func_name: str) -> List[Any]:
+        """Get the list of recorded return values for a function."""
+
+    @abstractmethod
+    def get(self, key: str, default: Any = None) -> Any:
+        """Get a recorded value by key."""
+
+    @abstractmethod
+    def get_all_func_return_values(self) -> Dict[str, Any]:
+        """Get all recorded data."""
+
+    @abstractmethod
+    def has(self, key: str) -> bool:
+        """Check if a key exists in recorded data."""
+
+    @abstractmethod
+    def clear(self) -> None:
+        """Clear all recorded data."""
+
+    @abstractmethod
+    def reset(self) -> None:
+        """Clear all recorded data and timing records."""
+
+    @abstractmethod
+    @contextmanager
+    def measure(self, name: str):
+        """Timing context manager to record execution duration."""
+
+    @abstractmethod
+    def __repr__(self) -> str:
+        """Return a string representation."""
+
+
+class RecordingContext(BaseRecordingContext):
+    """Captures actual execution results from production code for comparison.
+
+    This class acts as a dictionary-like container that stores key-value pairs
+    representing various metrics and intermediate results collected during
+    the production execution of a document processing task. It also supports
+    timing measurements via the [`measure()`](rag/svr/task_executor_refactor/recording_context.py:78) context manager.
+
+    The recorded data is later consumed by the [`Comparator`](rag/svr/task_executor_refactor/comparator.py:130)
+    to compare against dry-run execution results.
+
+    Example:
+        >>> ctx = RecordingContext()
+        >>> ctx.record("chunk_count", 100)
+        >>> ctx.get("chunk_count")
+        100
+        >>> ctx.get("missing_key", "default")
+        'default'
+    """
+
+    def __init__(self) -> None:
+        """Initialize a new RecordingContext."""
+        self._data: Dict[str, Any] = {}
+        self.records: List[Tuple[str, float]] = []
+
+    def record(self, key: str, value: Any) -> None:
+        """Record a value with the given key.
+
+        This method stores the provided value under the specified key in the
+        internal data dictionary. If the key already exists, the value will
+        be overwritten.
+
+        Args:
+            key: The key to store the value under. Should be a descriptive
+                string that identifies the metric or result being recorded.
+            value: The value to record. Can be any Python object, including
+                primitives, lists, dicts, or complex objects.
+        """
+        self._data[key] = value
+
+    def save_func_return_value(self, func_name: str, return_value: Any) -> None:
+        """Record a function's return value into a list associated with func_name.
+
+        Each func_name has a corresponding return_values_list. This method appends
+        the return_value to the list for the given func_name. If the list does not
+        exist, it will be created.
+
+        Args:
+            func_name: The name of the function whose return value is being recorded.
+            return_value: The return value to record.
+        """
+        if func_name not in self._data:
+            self._data[func_name] = []
+        self._data[func_name].append(return_value)
+
+    def get_func_return_values(self, func_name: str) -> List[Any]:
+        """Get the list of recorded return values for a function.
+
+        Args:
+            func_name: The name of the function.
+
+        Returns:
+            A list of recorded return values, or an empty list if not found.
+        """
+        return self._data.get(func_name, [])
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Get a recorded value by key.
+
+        Retrieves the value associated with the given key. If the key does
+        not exist, returns the provided default value.
+
+        Args:
+            key: The key to look up in the recorded data.
+            default: Default value to return if the key is not found.
+                Defaults to None.
+
+        Returns:
+            The recorded value associated with the key, or the default value
+            if the key does not exist.
+        """
+        return self._data.get(key, default)
+
+    def get_all_func_return_values(self) -> Dict[str, Any]:
+        """Get all recorded data.
+
+        Returns a shallow copy of all recorded data as a dictionary.
+        Modifications to the returned dictionary will not affect the
+        internal state of this context.
+
+        Returns:
+            A new dictionary containing all recorded key-value pairs.
+        """
+        return dict(self._data)
+
+    def has(self, key: str) -> bool:
+        """Check if a key exists in recorded data.
+
+        Args:
+            key: The key to check for existence.
+
+        Returns:
+            True if the key exists in the recorded data, False otherwise.
+        """
+        return key in self._data
+
+    def clear(self) -> None:
+        """Clear all recorded data.
+
+        Removes all key-value pairs from the internal data dictionary
+        and clears all timing records, resetting the context to its
+        initial empty state.
+        """
+        self._data.clear()
+        self.records.clear()
+
+    @contextmanager
+    def measure(self, name: str):
+        """Timing context manager to record execution duration.
+
+        Records the elapsed time (in seconds) for the operation specified
+        by `name`.
+
+        Usage::
+
+            with ctx.measure("build_chunks"):
+                ...
+
+        Args:
+            name: A descriptive name for the timed operation.
+        """
+        start = time.perf_counter()
+        try:
+            yield
+        finally:
+            elapsed = time.perf_counter() - start
+            self.records.append((name, elapsed))
+
+    def reset(self) -> None:
+        """Clear all recorded data and timing records."""
+        self.clear()
+
+    def __repr__(self) -> str:
+        """Return a string representation of the RecordingContext.
+
+        Returns:
+            A string showing the class name and all recorded data.
+        """
+        return f"RecordingContext({self._data})"
+
+
+class NullRecordingContext(BaseRecordingContext):
+    """No-op RecordingContext for production mode.
+
+    Accepts all RecordingContext API calls but performs no allocation.
+    Eliminates memory overhead in production where recorded data is unused.
+
+    Uses __slots__ for zero instance memory footprint.
+
+    Usage:
+        >>> ctx = NullRecordingContext()
+        >>> ctx.record("chunks", large_list)  # no-op, no memory allocated
+        >>> ctx.get("chunks")                 # always returns None
+    """
+
+    __slots__ = ()
+
+    def record(self, key: str, value: Any) -> None:
+        pass
+
+    def save_func_return_value(self, func_name: str, return_value: Any) -> None:
+        pass
+
+    def get_func_return_values(self, func_name: str) -> List[Any]:
+        return []
+
+    def get(self, key: str, default: Any = None) -> Any:
+        return default
+
+    def get_all_func_return_values(self) -> Dict[str, Any]:
+        return {}
+
+    def has(self, key: str) -> bool:
+        return False
+
+    def clear(self) -> None:
+        pass
+
+    def reset(self) -> None:
+        pass
+
+    @contextmanager
+    def measure(self, name: str):
+        yield
+
+    def __repr__(self) -> str:
+        return "NullRecordingContext()"
+
+
+# Module-level singleton to avoid repeated allocations
+_NULL_RECORDING_CONTEXT = NullRecordingContext()
+
+
+# Context variable for coroutine / thread isolation
+_recording_ctx_var: contextvars.ContextVar[BaseRecordingContext] = contextvars.ContextVar("recording_context")
+
+
+def get_recording_context() -> BaseRecordingContext:
+    """Get the BaseRecordingContext for the current execution context.
+
+    Returns the BaseRecordingContext bound to the current coroutine / thread.
+    If no context has been bound, raise RuntimeError.
+
+    Returns:
+        The current BaseRecordingContext, raise RuntimeError if not set.
+    """
+    context = _recording_ctx_var.get(None)
+    if context is None:
+        raise RuntimeError("no context")
+    return context
+
+
+def set_recording_context(ctx: BaseRecordingContext) -> None:
+    """Bind a BaseRecordingContext to the current execution context.
+
+    Args:
+        ctx: The BaseRecordingContext to bind, or None to unbind.
+    """
+    _recording_ctx_var.set(ctx)
+
+
+@contextmanager
+def recording_context_manager(ctx: BaseRecordingContext = None):
+    """Context manager that sets and restores the BaseRecordingContext.
+
+    Usage::
+
+        with recording_context_manager(RecordingContext()) as ctx:
+            ctx.record("key", "value")
+
+    Args:
+        ctx: The BaseRecordingContext to use. If None, a new one is created.
+
+    Yields:
+        The BaseRecordingContext that was set.
+    """
+    if ctx is None:
+        ctx = RecordingContext()
+    token = _recording_ctx_var.set(ctx)
+    try:
+        yield ctx
+    finally:
+        _recording_ctx_var.reset(token)
+
+
+def timed_with_recording(
+    func: Callable = None,
+    *,
+    recording_context: BaseRecordingContext = None,
+) -> Callable:
+    """Decorator that automatically uses the current BaseRecordingContext for timing.
+
+    Supports two usage forms:
+
+    1. Direct decoration (automatically uses context variable):
+
+        @timed_with_recording
+        def foo(): ...
+
+    2. Parameterized decoration with explicit BaseRecordingContext:
+
+        @timed_with_recording(recording_context=my_ctx)
+        def foo(): ...
+
+    The decorator records the execution time of the decorated function
+    into the BaseRecordingContext's timing records.
+
+    Args:
+        func: The function to decorate (used when called without parentheses).
+        recording_context: Optional BaseRecordingContext to use for timing.
+            If not provided, uses the context variable's current value.
+
+    Returns:
+        The decorated function.
+    """
+    from common.decorator import timing
+
+    if func is not None and callable(func):
+        # Used as @timed_with_recording without parentheses
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            ctx = recording_context or get_recording_context()
+            if ctx is not None:
+                return timing(context=ctx)(func)(*args, **kwargs)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    # Used as @timed_with_recording(...) with parentheses
+    def decorator(the_func: Callable) -> Callable:
+        @functools.wraps(the_func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            ctx = recording_context or get_recording_context()
+            if ctx is not None:
+                return timing(context=ctx)(the_func)(*args, **kwargs)
+            return the_func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

rag/svr/task_executor_refactor/report_generator.py

@@ -0,0 +1,140 @@
+#
+#  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.
+
+"""
+Report Generator Module.
+
+Provides data classes for comparison result reporting:
+- [`ComparisonResult`](rag/svr/task_executor_refactor/report_generator.py:40): Single key comparison result
+- [`ComparisonReport`](rag/svr/task_executor_refactor/report_generator.py:66): Full comparison report with serialization
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, List, Optional
+
+
+@dataclass
+class ComparisonResult:
+    """Result of comparing a single key between two contexts.
+
+    Attributes:
+        key: The key being compared.
+        match: Whether the values match.
+        production_value: Value from production context.
+        dry_run_value: Value from dry-run context.
+        diff_details: Optional description of the difference.
+    """
+
+    key: str
+    match: bool
+    production_value: Any = None
+    dry_run_value: Any = None
+    diff_details: Optional[str] = None
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for serialization."""
+        return {
+            "key": self.key,
+            "match": self.match,
+            "diff_details": self.diff_details,
+        }
+
+
+@dataclass
+class ComparisonReport:
+    """Report of comparing two RecordingContext instances.
+
+    Attributes:
+        task_id: The task identifier.
+        total_keys: Total number of keys compared.
+        matched_keys: Number of keys that matched.
+        mismatched_keys: Number of keys that mismatched.
+        missing_in_production: Keys missing in production context.
+        missing_in_dry_run: Keys missing in dry-run context.
+        details: List of individual comparison results.
+    """
+
+    task_id: str
+    total_keys: int = 0
+    matched_keys: int = 0
+    mismatched_keys: int = 0
+    missing_in_production: List[str] = field(default_factory=list)
+    missing_in_dry_run: List[str] = field(default_factory=list)
+    details: List["ComparisonResult"] = field(default_factory=list)
+
+    def summary(self) -> str:
+        """Generate a summary string.
+
+        Returns:
+            A human-readable summary of the comparison.
+        """
+        if self.total_keys == 0:
+            return f"Task {self.task_id}: No keys to compare"
+        match_rate = (self.matched_keys / self.total_keys) * 100
+        return (
+            f"Task {self.task_id}: {self.matched_keys}/{self.total_keys} "
+            f"keys matched ({match_rate:.1f}%)"
+        )
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for serialization.
+
+        Returns:
+            A dictionary representation of the report.
+        """
+        return {
+            "task_id": self.task_id,
+            "total_keys": self.total_keys,
+            "matched_keys": self.matched_keys,
+            "mismatched_keys": self.mismatched_keys,
+            "missing_in_production": self.missing_in_production,
+            "missing_in_dry_run": self.missing_in_dry_run,
+            "details": [d.to_dict() for d in self.details],
+            "summary": self.summary(),
+        }
+
+    def to_markdown(self) -> str:
+        """Generate a mark-down report.
+
+        Returns:
+            A markdown-formatted report string.
+        """
+        lines = [
+            f"# Comparison Report: {self.task_id}",
+            "",
+            "## Summary",
+            "",
+            f"- **Total keys**: {self.total_keys}",
+            f"- **Matched**: {self.matched_keys}",
+            f"- **Mismatched**: {self.mismatched_keys}",
+            f"- **Missing in production**: {', '.join(self.missing_in_production) or 'None'}",
+            f"- **Missing in dry-run**: {', '.join(self.missing_in_dry_run) or 'None'}",
+            "",
+            "## Details",
+            "",
+        ]
+
+        if self.details:
+            lines.append("| Key | Match | Details |")
+            lines.append("|-----|-------|---------|")
+            for d in self.details:
+                match_str = "✅" if d.match else "❌"
+                details_str = d.diff_details or "-"
+                lines.append(f"| {d.key} | {match_str} | {details_str} |")
+        else:
+            lines.append("No comparison details available.")
+
+        lines.append("")
+        return "\n".join(lines)

rag/svr/task_executor_refactor/task_context.py

@@ -0,0 +1,520 @@
+#
+#  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.
+
+"""
+Task Context Module.
+
+Provides [`TaskContext`](rag/svr/task_executor_refactor/task_context.py) as a typed wrapper
+around the task dictionary, providing convenient property accessors for all
+commonly used task attributes throughout the task executor refactor codebase.
+
+This module defines:
+- [`TaskDict`](rag/svr/task_executor_refactor/task_context.py): TypedDict for the raw task dictionary.
+- [`TaskLimiters`](rag/svr/task_executor_refactor/task_context.py): Dataclass encapsulating all rate limiters.
+- [`TaskCallbacks`](rag/svr/task_executor_refactor/task_context.py): Dataclass encapsulating all callback functions.
+- [`TaskContext`](rag/svr/task_executor_refactor/task_context.py): Main facade combining the above components.
+
+Usage example::
+
+    from rag.svr.task_executor_refactor.task_context import TaskContext, TaskLimiters, TaskCallbacks
+
+    ctx = TaskContext(
+        task=task_dict,
+        limiters=TaskLimiters(
+            chat=chat_limiter,
+            minio=minio_limiter,
+            chunk=chunk_limiter,
+            embed=embed_limiter,
+            kg=kg_limiter,
+        ),
+        callbacks=TaskCallbacks(
+            progress=progress_callback,
+            has_canceled=has_canceled_func,
+        ),
+        write_interceptor=write_interceptor,
+        recording_context=recording_context,
+    )
+
+    # Access task properties directly
+    task_id = ctx.id
+    tenant_id = ctx.tenant_id
+    kb_id = ctx.kb_id
+"""
+
+import asyncio
+from dataclasses import dataclass, field
+from functools import partial
+from typing import Any, Callable, Dict, List, Optional, Required, TypedDict
+
+from rag.svr.task_executor_refactor.recording_context import BaseRecordingContext
+from rag.svr.task_executor_refactor.write_operation_interceptor import WriteOperationInterceptor
+
+
+# ============================================================================
+# Type Definitions
+# ============================================================================
+
+
+class TaskDict(TypedDict, total=False):
+    """TypedDict defining the structure of the raw task dictionary.
+
+    All fields are optional except 'id' and 'tenant_id' which are required.
+    """
+
+    id: Required[str]
+    """Task identifier (required)."""
+
+    tenant_id: Required[str]
+    """Tenant identifier (required)."""
+
+    kb_id: str
+    """Knowledge base / dataset identifier."""
+
+    doc_id: str
+    """Document identifier."""
+
+    doc_ids: List[str]
+    """List of document identifiers (for batch tasks like RAPTOR/GraphRAG)."""
+
+    name: str
+    """Document name."""
+
+    location: str
+    """Document location/path."""
+
+    size: int
+    """Document file size in bytes."""
+
+    parser_id: str
+    """Parser identifier (e.g., 'naive', 'table', 'paper')."""
+
+    parser_config: Dict[str, Any]
+    """Document-level parser configuration."""
+
+    kb_parser_config: Dict[str, Any]
+
+    """Knowledge base level parser configuration."""
+
+    language: str
+    """Document language (e.g., 'en', 'zh')."""
+
+    llm_id: str
+    """LLM model identifier."""
+
+    embd_id: str
+    """Embedding model identifier."""
+
+    from_page: int
+    """Starting page number for processing (0-based)."""
+
+    to_page: int
+    """Ending page number for processing (-1 means all pages)."""
+
+    task_type: str
+    """Task type (e.g., 'dataflow', 'raptor', 'graphrag', 'memory')."""
+
+    dataflow_id: str
+    """Dataflow/pipeline identifier."""
+
+    pagerank: int
+    """PageRank value for document scoring."""
+
+    file: Any
+    """File object for dataflow processing."""
+
+    memory_id: str
+    """Memory identifier for memory tasks."""
+
+    source_id: str
+    """Source identifier for memory tasks."""
+
+    message_dict: Dict[str, Any]
+    """Message dictionary for memory tasks."""
+
+# ============================================================================
+# Data Classes
+# ============================================================================
+
+
+@dataclass
+class TaskLimiters:
+    """Encapsulates all rate limiters for task execution.
+
+    Each limiter is an asyncio.Semaphore used to control concurrency
+    for different types of operations.
+    """
+
+    chat: asyncio.Semaphore = None
+    """Asyncio semaphore for chat model rate limiting."""
+
+    minio: asyncio.Semaphore = None
+    """Asyncio semaphore for MinIO rate limiting."""
+
+    chunk: asyncio.Semaphore = None
+    """Asyncio semaphore for chunk building rate limiting."""
+
+    embed: asyncio.Semaphore = None
+    """Asyncio semaphore for embedding rate limiting."""
+
+    kg: asyncio.Semaphore = None
+    """Asyncio semaphore for knowledge graph rate limiting."""
+
+
+def _noop_progress(**kwargs: Any) -> None:
+    """No-op progress callback."""
+    pass
+
+
+def _not_canceled(task_id: str) -> bool:
+    """Default cancellation check - always returns False."""
+    return False
+
+
+@dataclass
+class TaskCallbacks:
+    """Encapsulates all callback functions for task execution."""
+
+    progress: Callable = field(default_factory=lambda: _noop_progress)
+    """Callback function for progress updates (raw, requires task_id, from_page, to_page)."""
+
+    has_canceled: Callable = field(default_factory=lambda: _not_canceled)
+    """Function to check if task is canceled."""
+
+
+# ============================================================================
+# Main Class
+# ============================================================================
+
+
+class TaskContext:
+    """Typed wrapper around the task dictionary providing convenient property accessors.
+
+    This class uses composition to encapsulate:
+    1. The raw task dictionary (TaskDict)
+    2. Execution limiters (TaskLimiters)
+    3. Callback functions (TaskCallbacks)
+    4. Optional write operation interceptor
+    5. Optional recording context for intermediate results
+
+    The properties provide a clean interface for accessing task attributes
+    without needing to use dictionary access with string keys throughout
+    the codebase.
+    """
+
+    # Default values for optional task fields
+    _DEFAULTS: Dict[str, Any] = {
+        "kb_id": "",
+        "doc_id": "",
+        "doc_ids": [],
+        "name": "",
+        "location": "",
+        "size": 0,
+        "parser_id": "",
+        "parser_config": {},
+        "kb_parser_config": {},
+        "language": "en",
+        "llm_id": "",
+        "embd_id": "",
+        "from_page": 0,
+        "to_page": -1,
+        "task_type": "",
+        "dataflow_id": "",
+        "pagerank": 0,
+        "memory_id": "",
+        "source_id": "",
+        "message_dict": {},
+    }
+
+    def __init__(
+        self,
+        task: TaskDict,
+        limiters: TaskLimiters,
+        callbacks: TaskCallbacks,
+        write_interceptor: WriteOperationInterceptor = None,
+        recording_context: BaseRecordingContext = None,
+    ):
+        """Initialize TaskContext.
+
+        Args:
+            task: The raw task dictionary containing all task attributes.
+            limiters: TaskLimiters dataclass containing all rate limiters.
+            callbacks: TaskCallbacks dataclass containing all callback functions.
+            write_interceptor: Optional interceptor for write operations.
+            recording_context: Optional BaseRecordingContext for intermediate result
+                capture. Must be injected via constructor.
+
+        Raises:
+            ValueError: If required fields ('id', 'tenant_id') are missing from task.
+        """
+        # Validate required fields
+        if "id" not in task:
+            raise ValueError("Task must contain 'id'")
+        if "tenant_id" not in task:
+            raise ValueError("Task must contain 'tenant_id'")
+
+        self._task = task
+        self.limiters = limiters
+        self.callbacks = callbacks
+        self._write_interceptor = write_interceptor
+        self._recording_context = recording_context
+
+
+        # Prepare progress callback and set it on the context
+        progress_cb = partial(
+            callbacks.progress,
+            self.id,
+            self.from_page,
+            self.to_page,
+        )
+        self._progress_cb = progress_cb
+
+    # =========================================================================
+    # Core task identity properties
+    # =========================================================================
+
+    @property
+    def id(self) -> str:
+        """Task identifier."""
+        return self._task["id"]
+
+    @property
+    def tenant_id(self) -> str:
+        """Tenant identifier."""
+        return self._task["tenant_id"]
+
+    @property
+    def kb_id(self) -> str:
+        """Knowledge base / dataset identifier."""
+        return self._task.get("kb_id", self._DEFAULTS["kb_id"])
+
+    @property
+    def doc_id(self) -> str:
+        """Document identifier."""
+        return self._task.get("doc_id", self._DEFAULTS["doc_id"])
+
+    @property
+    def doc_ids(self) -> List[str]:
+        """List of document identifiers (for batch tasks like RAPTOR/GraphRAG)."""
+        return self._task.get("doc_ids", list(self._DEFAULTS["doc_ids"]))
+
+    # =========================================================================
+    # Document metadata properties
+    # =========================================================================
+
+    @property
+    def name(self) -> str:
+        """Document name."""
+        return self._task.get("name", self._DEFAULTS["name"])
+
+    @property
+    def location(self) -> str:
+        """Document location/path."""
+        return self._task.get("location", self._DEFAULTS["location"])
+
+    @property
+    def size(self) -> int:
+        """Document file size in bytes."""
+        return self._task.get("size", self._DEFAULTS["size"])
+
+    # =========================================================================
+    # Parser configuration properties
+    # =========================================================================
+
+    @property
+    def parser_id(self) -> str:
+        """Parser identifier (e.g., 'naive', 'table', 'paper')."""
+        return self._task.get("parser_id", self._DEFAULTS["parser_id"])
+
+    @property
+    def parser_config(self) -> Dict[str, Any]:
+        """Document-level parser configuration."""
+        return self._task.get("parser_config", {})
+
+    @property
+    def kb_parser_config(self) -> Dict[str, Any]:
+        """Knowledge base level parser configuration."""
+        return self._task.get("kb_parser_config", {})
+
+    # =========================================================================
+    # Language and model properties
+    # =========================================================================
+
+    @property
+    def language(self) -> str:
+        """Document language (e.g., 'en', 'zh')."""
+        return self._task.get("language", self._DEFAULTS["language"])
+
+    @property
+    def llm_id(self) -> str:
+        """LLM model identifier."""
+        return self._task.get("llm_id", self._DEFAULTS["llm_id"])
+
+    @property
+    def embd_id(self) -> str:
+        """Embedding model identifier."""
+        return self._task.get("embd_id", self._DEFAULTS["embd_id"])
+
+    # =========================================================================
+    # Page range properties
+    # =========================================================================
+
+    @property
+    def from_page(self) -> int:
+        """Starting page number for processing (0-based)."""
+        return self._task.get("from_page", self._DEFAULTS["from_page"])
+
+    @property
+    def to_page(self) -> int:
+        """Ending page number for processing (-1 means all pages)."""
+        return self._task.get("to_page", self._DEFAULTS["to_page"])
+
+    # =========================================================================
+    # Task type and routing properties
+    # =========================================================================
+
+    @property
+    def task_type(self) -> str:
+        """Task type (e.g., 'dataflow', 'raptor', 'graphrag', 'memory')."""
+        return self._task.get("task_type", self._DEFAULTS["task_type"])
+
+    @property
+    def dataflow_id(self) -> str:
+        """Dataflow/pipeline identifier."""
+        return self._task.get("dataflow_id", self._DEFAULTS["dataflow_id"])
+
+    # =========================================================================
+    # Additional properties
+    # =========================================================================
+
+    @property
+    def pagerank(self) -> int:
+        """PageRank value for document scoring."""
+        return self._task.get("pagerank", self._DEFAULTS["pagerank"])
+
+    @property
+    def file(self) -> Optional[Any]:
+        """File object for dataflow processing."""
+        return self._task.get("file")
+
+    # =========================================================================
+    # Memory task specific properties
+    # =========================================================================
+
+    @property
+    def memory_id(self) -> str:
+        """Memory identifier for memory tasks."""
+        return self._task.get("memory_id", self._DEFAULTS["memory_id"])
+
+    @property
+    def source_id(self) -> str:
+        """Source identifier for memory tasks."""
+        return self._task.get("source_id", self._DEFAULTS["source_id"])
+
+    @property
+    def message_dict(self) -> Dict[str, Any]:
+        """Message dictionary for memory tasks."""
+        return self._task.get("message_dict", {})
+
+    # =========================================================================
+    # Raw task dictionary access
+    # =========================================================================
+
+    @property
+    def raw_task(self) -> Dict[str, Any]:
+        """Return the raw task dictionary."""
+        return self._task
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Get a value from the task dictionary with a default.
+
+        Args:
+            key: The key to look up.
+            default: Default value if key is not found.
+
+        Returns:
+            The value associated with the key, or default if not found.
+        """
+        return self._task.get(key, default)
+
+    # =========================================================================
+    # Limiter properties (proxies to TaskLimiters dataclass)
+    # =========================================================================
+
+    @property
+    def chat_limiter(self) -> asyncio.Semaphore:
+        """Asyncio semaphore for chat model rate limiting."""
+        return self.limiters.chat or asyncio.Semaphore(1)
+
+    @property
+    def minio_limiter(self) -> asyncio.Semaphore:
+        """Asyncio semaphore for MinIO rate limiting."""
+        return self.limiters.minio or asyncio.Semaphore(1)
+
+    @property
+    def chunk_limiter(self) -> asyncio.Semaphore:
+        """Asyncio semaphore for chunk building rate limiting."""
+        return self.limiters.chunk or asyncio.Semaphore(1)
+
+    @property
+    def embed_limiter(self) -> asyncio.Semaphore:
+        """Asyncio semaphore for embedding rate limiting."""
+        return self.limiters.embed or asyncio.Semaphore(1)
+
+    @property
+    def kg_limiter(self) -> asyncio.Semaphore:
+        """Asyncio semaphore for knowledge graph rate limiting."""
+        return self.limiters.kg or asyncio.Semaphore(1)
+
+    # =========================================================================
+    # Context and interceptor properties
+    # =========================================================================
+
+    @property
+    def recording_context(self) -> BaseRecordingContext:
+        """BaseRecordingContext for this task.
+
+        Must be injected via constructor. Raises RuntimeError if accessed
+        before initialization or if no context was provided.
+        """
+        if self._recording_context is None:
+            raise RuntimeError("recording_context accessed but not injected into TaskContext")
+        return self._recording_context
+
+    @property
+    def write_interceptor(self) -> WriteOperationInterceptor:
+        """Write operation interceptor for comparison mode."""
+        return self._write_interceptor
+
+    # =========================================================================
+    # Callback properties (proxies to TaskCallbacks dataclass)
+    # =========================================================================
+
+    @property
+    def has_canceled_func(self) -> Callable:
+        """Function to check if task is canceled."""
+        return self.callbacks.has_canceled
+
+    # =========================================================================
+    # Pre-bound progress callback
+    # =========================================================================
+
+    @property
+    def progress_cb(self) -> Callable:
+        """Pre-bound progress callback (task_id, from_page, to_page already bound).
+
+        Use this property in services for progress updates.
+        Falls back to progress_callback if progress_cb is not set.
+        """
+        return self._progress_cb

rag/svr/task_executor_refactor/task_executor_refactoring_plan.md

@@ -0,0 +1,492 @@
+# Task Executor Refactoring Plan
+
+## 1. Current State Analysis
+
+### 1.1 Original File
+- **File Location**: `rag/svr/task_executor.py`
+- **Lines of Code**: Approximately 1,780 lines
+- **Primary Responsibilities**: Task consumption, document chunking, vectorization, index building, RAPTOR/GraphRAG processing, heartbeat reporting
+
+### 1.2 Identified Issues
+
+| Issue Type | Specific Manifestation |
+|------------|------------------------|
+| Single Responsibility Violation | One file handles 7+ different responsibilities |
+| Global State | Global variables like `DONE_TASKS`, `FAILED_TASKS`, `CURRENT_TASKS` |
+| Tight Coupling | Direct dependencies on `TaskService`, `DocumentService`, `REDIS_CONN`, etc. |
+| Untestable | Functions depend on global state and external services, difficult to mock |
+| Hardcoded Configuration | `BATCH_SIZE`, `FACTORY`, etc. hardcoded in the file |
+
+---
+
+## 2. Implemented Architecture
+
+### 2.1 Actual Module Structure
+
+```
+rag/svr/task_executor_refactor/
+├── task_context.py              # Task context encapsulation (~450 lines)
+├── recording_context.py         # Execution result recording context (~330 lines)
+├── write_operation_interceptor.py # Write operation interceptor (~130 lines)
+├── chunk_service.py             # Document chunking service (~430 lines)
+├── chunk_builder.py             # Chunk building logic (~130 lines)
+├── chunk_post_processor.py      # Post-chunking logic (~350 lines)
+├── embedding_service.py         # Embedding service (~130 lines)
+├── embedding_utils.py           # Embedding utility functions (~210 lines)
+├── raptor_service.py            # RAPTOR processing service (~520 lines)
+├── raptor_utils.py              # RAPTOR utility functions (~100 lines)
+├── dataflow_service.py          # Dataflow pipeline service (~430 lines)
+├── post_processor.py            # Post-processing service (~150 lines)
+├── comparator.py                # Comparator (~550 lines)
+├── report_generator.py          # Report generator (~130 lines)
+├── task_handler.py              # Task handler entry point (~630 lines)
+├── task_manager.py              # Task manager (~200 lines)
+├── constants.py                 # Constant definitions (~25 lines)
+└── insert_service.py            # Insert service (~150 lines)
+
+test/unit_test/rag/svr/task_executor_refactor/
+├── conftest.py                  # Shared test fixtures (~260 lines)
+├── test_task_context.py         # TaskContext tests (~410 lines)
+├── test_recording_context.py    # RecordingContext tests (~330 lines)
+├── test_write_operation_interceptor.py # Interceptor tests (~450 lines)
+├── test_chunk_service.py        # ChunkService tests (~560 lines)
+├── test_chunk_builder.py        # ChunkBuilder tests (~290 lines)
+├── test_chunk_post_processor.py # ChunkPostProcessor tests (~550 lines)
+├── test_embedding_service.py    # EmbeddingService tests (~190 lines)
+├── test_embedding_utils.py      # EmbeddingUtils tests (~370 lines)
+├── test_raptor_service.py       # RaptorService tests (~350 lines)
+├── test_dataflow_service.py     # DataflowService tests (~250 lines)
+├── test_post_processor.py       # PostProcessor tests (~120 lines)
+├── test_comparator.py           # Comparator tests (~570 lines)
+├── test_task_handler.py         # TaskHandler unit tests (~800 lines)
+├── test_task_handler_integration.py # TaskHandler integration tests (~1400 lines)
+└── test_constants.py            # Constants tests (~40 lines)
+```
+
+### 2.2 Layered Architecture Design
+
+```
+┌─────────────────────────────────────────────────────────────────┐
+│                        Business Layer                            │
+│                        task_handler.py                           │
+│  ┌───────────────────────────────────────────────────────────┐  │
+│  │  TaskHandler Class                                         │  │
+│  │  ├── handle_task()     # Entry point, handles cancellation and exceptions │  │
+│  │  ├── handle()          # Task type routing dispatch        │  │
+│  │  ├── _run_dataflow()   # Dataflow pipeline execution       │  │
+│  │  ├── _run_raptor()     # RAPTOR summary generation         │  │
+│  │  ├── _run_graphrag()   # GraphRAG knowledge graph          │  │
+│  │  └── _run_standard_chunking() # Standard chunking flow     │  │
+│  └───────────────────────────────────────────────────────────┘  │
+│                                                                 │
+│  Entry Functions:                                                │
+│  ├── run_refactored_task()   # Refactored version entry          │
+│  └── dry_run_task()          # Comparison mode entry             │
+├─────────────────────────────────────────────────────────────────┤
+│                        Service Layer                             │
+│                                                                 │
+│  ┌─────────────────┐  ┌──────────────────┐  ┌────────────────┐  │
+│  │ ChunkService    │  │ EmbeddingService │  │ RaptorService  │  │
+│  │                 │  │                  │  │                │  │
+│  │ build_chunks()  │  │ embed_chunks()   │  │ run_raptor_    │  │
+│  │ insert_chunks() │  │                  │  │   for_kb()     │  │
+│  └─────────────────┘  └──────────────────┘  └────────────────┘  │
+│                                                                 │
+│  ┌─────────────────┐  ┌──────────────────┐  ┌────────────────┐  │
+│  │DataflowService  │  │ PostProcessor    │  │ InsertService  │  │
+│  │                 │  │                  │  │                │  │
+│  │ run_dataflow()  │  │ process_table_   │  │ insert_chunks()│  │
+│  │                 │  │   parser_        │  │                │  │
+│  │                 │  │   metadata()     │  │                │  │
+│  └─────────────────┘  └──────────────────┘  └────────────────┘  │
+│                                                                 │
+│  ┌─────────────────┐  ┌──────────────────┐                      │
+│  │ ChunkBuilder    │  │ChunkPostProcessor│                      │
+│  │                 │  │                  │                      │
+│  │ Chunk building  │  │ Post-processing  │                      │
+│  │ logic           │  │ logic            │                      │
+│  └─────────────────┘  └──────────────────┘                      │
+├─────────────────────────────────────────────────────────────────┤
+│                     Infrastructure Layer                         │
+│                                                                 │
+│  ┌─────────────────┐  ┌──────────────────┐  ┌────────────────┐  │
+│  │ TaskContext     │  │ RecordingContext │  │ Comparator     │  │
+│  │                 │  │                  │  │                │  │
+│  │ Task property   │  │ Execution result │  │ Production vs  │  │
+│  │ accessors       │  │ recording        │  │ Dry-run        │  │
+│  │ Rate limiter    │  │ Function return  │  │ Difference     │  │
+│  │ encapsulation   │  │ value recording  │  │ report gen     │  │
+│  │ Interceptor     │  │ Timing decorator │  │                │  │
+│  │ references      │  │                  │  │                │  │
+│  └─────────────────┘  └──────────────────┘  └────────────────┘  │
+│                                                                 │
+│  ┌──────────────────────────────────┐  ┌────────────────────┐   │
+│  │ WriteOperationInterceptor        │  │ ReportGenerator    │   │
+│  │                                  │  │                    │   │
+│  │ Whitelist method interception    │  │ Difference report  │   │
+│  │ Pre-recorded return value replay │  │ Formatted output   │   │
+│  └──────────────────────────────────┘  └────────────────────┘   │
+│                                                                 │
+│  ┌──────────────────────────────────┐  ┌────────────────────┐   │
+│  │ TaskManager                      │  │ Constants & Utils  │   │
+│  │                                  │  │                    │   │
+│  │ Task lifecycle management        │  │ CANVAS_DEBUG_      │   │
+│  │ Task state tracking              │  │   DOC_ID           │   │
+│  └──────────────────────────────────┘  │ GRAPH_RAPTOR_      │   │
+│                                        │   FAKE_DOC_ID      │   │
+│                                        └────────────────────┘   │
+└─────────────────────────────────────────────────────────────────┘
+```
+
+---
+
+## 3. Core Design Patterns
+
+### 3.1 Dependency Injection
+
+All services receive `TaskContext` through constructors, rather than directly importing global state:
+
+```python
+class ChunkService:
+    def __init__(self, ctx: TaskContext):
+        self._task_context = ctx
+```
+
+### 3.2 Interceptor Pattern
+
+`WriteOperationInterceptor` is used to replay production execution return values in comparison mode:
+
+```python
+# Comparison mode: intercept write operations
+if ctx.write_interceptor:
+    update_result = ctx.write_interceptor.intercept("KnowledgebaseService.update_by_id")
+else:
+    update_result = KnowledgebaseService.update_by_id(kb.id, {"parser_config": kb_parser_config})
+```
+
+### 3.3 Recording Context Pattern
+
+`RecordingContext` captures intermediate results for comparison:
+
+```python
+# Record intermediate results
+get_recording_context().record("chunks", chunks)
+get_recording_context().record("token_count", token_count)
+```
+
+### 3.4 Factory Pattern
+
+Parser modules are registered through factory mapping:
+
+```python
+PARSER_FACTORY = {}
+
+def register_parser(parser_id: str, parser_module):
+    PARSER_FACTORY[parser_id] = parser_module
+```
+
+---
+
+## 4. Task Execution Flow
+
+### 4.1 Standard Task Flow
+
+```
+run_refactored_task()
+    │
+    ▼
+TaskContext Creation
+    │
+    ▼
+TaskHandler.handle_task()
+    │
+    ├── try: handle()
+    │       │
+    │       ├── Task type judgment
+    │       │   ├── "memory"    → handle_save_to_memory_task()
+    │       │   ├── "dataflow"  → DataflowService.run_dataflow()
+    │       │   ├── "raptor"    → _run_raptor()
+    │       │   ├── "graphrag"  → _run_graphrag()
+    │       │   ├── "mindmap"   → Placeholder
+    │       │   └── Others      → _run_standard_chunking()
+    │       │
+    │       └── _run_standard_chunking()
+    │           │
+    │           ├── Bind embedding model
+    │           ├── Retrieve storage binary
+    │           ├── ChunkService.build_chunks()
+    │           │   ├── File size validation
+    │           │   ├── Parser chunking
+    │           │   ├── Outline extraction
+    │           │   ├── MinIO upload
+    │           │   ├── Keyword extraction
+    │           │   ├── Question generation
+    │           │   ├── Metadata generation
+    │           │   └── Content tagging
+    │           ├── EmbeddingService.embed_chunks()
+    │           ├── TOC generation (async)
+    │           ├── ChunkService.insert_chunks()
+    │           ├── PostProcessor.process_table_parser_metadata()
+    │           ├── TOC insertion
+    │           └── DocumentService.increment_chunk_num()
+    │
+    └── finally: Cancel task cleanup
+```
+
+### 4.2 Comparison Mode Flow
+
+```
+dry_run_task()
+    │
+    ├── Create WriteOperationInterceptor (using pre-recorded values from recording_ctx1)
+    ├── Create new RecordingContext (recording_ctx2)
+    ├── Set recording_context to recording_ctx2
+    │
+    ▼
+TaskHandler.handle_task()  # Execute with interceptor replay
+    │
+    ▼
+ContextComparator.compare(task_id, recording_ctx1, recording_ctx2)
+    │
+    ├── Key-by-key comparison
+    ├── Generate difference report
+    └── Output mismatched_keys and remaining_values
+```
+
+---
+
+## 5. Testing Strategy
+
+### 5.1 Test Coverage Status
+
+| Module | Test File | Test Lines | Coverage Focus |
+|--------|-----------|------------|----------------|
+| `TaskContext` | `test_task_context.py` | ~410 | Property accessors, rate limiters, interceptors |
+| `RecordingContext` | `test_recording_context.py` | ~330 | Record/retrieve, function return values, timing |
+| `WriteOperationInterceptor` | `test_write_operation_interceptor.py` | ~450 | Whitelist validation, FIFO replay |
+| `ChunkService` | `test_chunk_service.py` | ~560 | Chunking logic, post-processing, insertion |
+| `ChunkBuilder` | `test_chunk_builder.py` | ~290 | Chunk building logic |
+| `ChunkPostProcessor` | `test_chunk_post_processor.py` | ~550 | Post-processing logic |
+| `EmbeddingService` | `test_embedding_service.py` | ~190 | Batch encoding, vector stacking |
+| `EmbeddingUtils` | `test_embedding_utils.py` | ~370 | Text preparation, truncation, stacking |
+| `RaptorService` | `test_raptor_service.py` | ~350 | RAPTOR execution |
+| `DataflowService` | `test_dataflow_service.py` | ~250 | Dataflow execution |
+| `PostProcessor` | `test_post_processor.py` | ~120 | Table metadata processing |
+| `Comparator` | `test_comparator.py` | ~570 | Various type comparison logic |
+| `TaskHandler` | `test_task_handler.py` | ~800 | Routing, model binding, task types |
+| `TaskHandler` | `test_task_handler_integration.py` | ~1400 | Full flow integration tests |
+| `constants.py` | `test_constants.py` | ~40 | Constant value validation |
+
+**Total Test Code**: Approximately 6,700+ lines
+
+### 5.2 Mock Strategy
+
+```python
+# conftest.py shared fixtures
+
+@pytest.fixture
+def mock_task():
+    """Standard test task"""
+    return {
+        "id": "task-001",
+        "task_type": "standard",
+        "tenant_id": "tenant-001",
+        "kb_id": "kb-001",
+        "doc_id": "doc-001",
+        "name": "test.pdf",
+        ...
+    }
+
+@pytest.fixture
+def mock_task_context(mock_task):
+    """TaskContext fixture"""
+    return TaskContext(
+        task=mock_task,
+        chat_limiter=asyncio.Semaphore(1),
+        minio_limiter=asyncio.Semaphore(1),
+        chunk_limiter=asyncio.Semaphore(1),
+        embed_limiter=asyncio.Semaphore(1),
+        kg_limiter=asyncio.Semaphore(1),
+        progress_callback=lambda **kwargs: None,
+        has_canceled_func=lambda task_id: False,
+    )
+```
+
+### 5.3 Test Coverage Targets
+
+| Module | Current Coverage | Target Coverage | Notes |
+|--------|-----------------|-----------------|-------|
+| `task_context.py` | ~90% | 95%+ | Good |
+| `recording_context.py` | ~85% | 90%+ | Good |
+| `write_operation_interceptor.py` | ~90% | 95%+ | Good |
+| `chunk_service.py` | ~80% | 90%+ | Good |
+| `chunk_builder.py` | ~75% | 85%+ | Needs more edge case tests |
+| `chunk_post_processor.py` | ~80% | 90%+ | Good |
+| `embedding_service.py` | ~85% | 90%+ | Good |
+| `raptor_service.py` | ~70% | 85%+ | Improved |
+| `dataflow_service.py` | ~75% | 85%+ | Good |
+| `post_processor.py` | ~75% | 85%+ | Good |
+| `comparator.py` | ~85% | 90%+ | Good |
+| `task_handler.py` | ~75% | 85%+ | Needs more integration tests |
+
+---
+
+## 6. Backward Compatibility Strategy
+
+### 6.1 Dual Code Path Coexistence
+
+Original `task_executor.py` is preserved, importing refactored modules:
+
+```python
+# rag/svr/task_executor.py (modified)
+from rag.svr.task_executor_refactor.task_handler import dry_run_task, run_refactored_task
+from rag.svr.task_executor_refactor.recording_context import timed_with_recording, get_recording_context, \
+    RecordingContext, set_recording_context
+```
+
+### 6.2 Migration Plan
+
+| Phase | Status | Description |
+|-------|--------|-------------|
+| Phase 1 | ✅ Completed | Dual code paths parallel, `run_refactored_task()` and `dry_run_task()` available |
+| Phase 2 | ⏳ Pending | Switch default execution to refactored code, keep old code as fallback |
+| Phase 3 | ⏳ Pending | Remove old code after validation period |
+
+---
+
+## 7. Equivalence Guarantee Strategy
+
+### 7.1 Comparison Mode
+
+The refactoring introduces a unique comparison mode to verify equivalence:
+
+1. **Production Execution**: Run original code path, record all intermediate results to `RecordingContext`
+2. **Dry Run**: Use `WriteOperationInterceptor` to replay production results, record new intermediate results
+3. **Comparison**: `ContextComparator` compares differences between two contexts
+
+### 7.2 Comparison Strategy
+
+| Data Type | Comparison Strategy |
+|-----------|---------------------|
+| Primitives (int, str, bool) | Direct equality |
+| Floating point | Tolerance range |
+| Lists | Length + ID set + sampled content |
+| Dictionaries | Key set + recursive value comparison |
+| None | Equal |
+
+---
+
+## 8. Risks and Mitigations
+
+| Risk | Mitigation | Status |
+|------|------------|--------|
+| Refactoring introduces bugs | Comparison mode verifies equivalence | ✅ Implemented |
+| Performance regression | Benchmark comparison | ⏳ Pending |
+| Memory increase | RecordingContext stores intermediate results | ⚠️ Needs monitoring |
+| Insufficient test coverage | Supplement RaptorService tests | ✅ Improved |
+| Large modules | Split chunk_service.py | ✅ Split |
+
+---
+
+## 9. Future Improvement Suggestions
+
+### 9.1 High Priority
+
+1. **Performance Benchmarking**: Compare performance before and after refactoring
+2. **Improve Integration Tests**: Add more end-to-end test scenarios
+3. **Fix Type Annotations**: Add `Any` type for `default_value` and similar parameters
+
+### 9.2 Medium Priority
+
+4. **Improve Exception Handling**: Preserve more context information when wrapping exceptions
+5. **Documentation Improvement**: Add usage examples to docstrings
+
+### 9.3 Low Priority
+
+6. **Memory Optimization**: Consider streaming recording for large tasks
+7. **Code Cleanup**: Remove unused imports and functions
+
+---
+
+## 10. Code Statistics
+
+### 10.1 Source Code
+
+| Module | Lines | Type |
+|--------|-------|------|
+| `task_context.py` | ~450 | Infrastructure |
+| `recording_context.py` | ~330 | Infrastructure |
+| `write_operation_interceptor.py` | ~130 | Infrastructure |
+| `comparator.py` | ~550 | Infrastructure |
+| `report_generator.py` | ~130 | Infrastructure |
+| `constants.py` | ~25 | Infrastructure |
+| `task_manager.py` | ~200 | Infrastructure |
+| `chunk_service.py` | ~430 | Service |
+| `chunk_builder.py` | ~130 | Service |
+| `chunk_post_processor.py` | ~350 | Service |
+| `embedding_service.py` | ~130 | Service |
+| `embedding_utils.py` | ~210 | Utility |
+| `raptor_service.py` | ~520 | Service |
+| `raptor_utils.py` | ~100 | Utility |
+| `dataflow_service.py` | ~430 | Service |
+| `post_processor.py` | ~150 | Service |
+| `insert_service.py` | ~150 | Service |
+| `task_handler.py` | ~630 | Business |
+| **Source Code Total** | **~4,900** | |
+
+### 10.2 Test Code
+
+| Test File | Lines |
+|-----------|-------|
+| `conftest.py` | ~260 |
+| `test_task_context.py` | ~410 |
+| `test_recording_context.py` | ~330 |
+| `test_write_operation_interceptor.py` | ~450 |
+| `test_chunk_service.py` | ~560 |
+| `test_chunk_builder.py` | ~290 |
+| `test_chunk_post_processor.py` | ~550 |
+| `test_embedding_service.py` | ~190 |
+| `test_embedding_utils.py` | ~370 |
+| `test_raptor_service.py` | ~350 |
+| `test_dataflow_service.py` | ~250 |
+| `test_post_processor.py` | ~120 |
+| `test_comparator.py` | ~570 |
+| `test_task_handler.py` | ~800 |
+| `test_task_handler_integration.py` | ~1400 |
+| `test_constants.py` | ~40 |
+| **Test Code Total** | **~6,700+** |
+
+### 10.3 Documentation
+
+| Document | Lines |
+|----------|-------|
+| `task_executor_refactoring_plan.md` | This document |
+
+---
+
+## 11. Time Estimation
+
+| Phase | Completed | Estimated Time |
+|-------|-----------|----------------|
+| Infrastructure Preparation | ✅ Completed | - |
+| Core Logic Decoupling | ✅ Completed | - |
+| Advanced Feature Decoupling | ✅ Completed | - |
+| Test Writing | ✅ Mostly Completed | - |
+| Performance Benchmarking | ⏳ Pending | 1-2 days |
+| Migration to Production | ⏳ Pending | 1-2 days |
+| **Remaining Total** | | **2-4 days** |
+
+---
+
+## 12. Summary
+
+This refactoring has successfully decomposed the monolithic `task_executor.py` into a layered, testable module architecture:
+
+- ✅ **Layered Architecture**: Infrastructure Layer → Service Layer → Business Layer
+- ✅ **Dependency Injection**: Execution resources injected via `TaskContext`
+- ✅ **Comparison Mode**: Innovative Production vs Dry-run comparison framework
+- ✅ **Test Coverage**: Approximately 6,700+ lines of test code
+- ✅ **Module Decomposition**: Large modules split into smaller responsibility units
+- ⚠️ **Pending Improvements**: Performance benchmarking, production migration validation
+
+**Overall Status**: Core refactoring completed, test coverage is good, ready for validation and migration phases.

rag/svr/task_executor_refactor/task_handler.py

@@ -0,0 +1,576 @@
+#
+#  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.
+
+"""
+Task Handler Module.
+
+Provides [`TaskHandler`](rag/svr/task_executor_refactor/task_handler.py:56) as the main entry point
+for handling document processing tasks with refactored, testable methods.
+"""
+
+import asyncio
+import logging
+import json
+import xxhash
+
+from timeit import default_timer as timer
+from typing import Callable, Dict, List, Optional
+
+from api.db.services.document_service import DocumentService
+from api.db.services.knowledgebase_service import KnowledgebaseService
+from api.db.joint_services.memory_message_service import handle_save_to_memory_task
+from api.db.joint_services.tenant_model_service import (
+    get_model_config_by_type_and_name,
+    get_tenant_default_model_by_type,
+)
+from api.db.services.llm_service import LLMBundle
+from api.db.services.task_service import GRAPH_RAPTOR_FAKE_DOC_ID
+from common.constants import LLMType
+from common.exceptions import TaskCanceledException
+from common.misc_utils import thread_pool_exec
+from rag.nlp import search
+from rag.svr.task_executor_refactor.constants import CANVAS_DEBUG_DOC_ID
+from rag.svr.task_executor_refactor.chunk_service import ChunkService
+from rag.svr.task_executor_refactor.dataflow_service import BillingHook, DataflowService
+from rag.svr.task_executor_refactor.embedding_service import EmbeddingService
+from rag.svr.task_executor_refactor.post_processor import PostProcessor
+from rag.svr.task_executor_refactor.raptor_service import RaptorService
+from rag.svr.task_executor_refactor.raptor_utils import delete_raptor_chunks
+from rag.svr.task_executor_refactor.recording_context import RecordingContext
+from rag.svr.task_executor_refactor.task_context import TaskContext
+from rag.graphrag.general.index import run_graphrag_for_kb
+from api.db.services.file2document_service import File2DocumentService
+from rag.prompts.generator import run_toc_from_text
+from common import settings
+
+
+class TaskHandler:
+    """Main task handler for document processing.
+
+    This class orchestrates the entire document processing pipeline:
+    1. Task type detection (memory, dataflow, raptor, graphrag, standard)
+    2. Model binding (embedding, chat)
+    3. Chunk building or RAPTOR/GraphRAG execution
+    4. Embedding
+    5. Indexing
+    6. Post-processing (TOC, table metadata)
+
+    All intermediate results are recorded via RecordingContext for comparison.
+    """
+
+    def __init__(
+        self,
+        ctx: TaskContext,
+        billing_hook: Optional[BillingHook] = None,
+    ):
+        """Initialize TaskHandler.
+
+        Args:
+            ctx: TaskContext containing task configuration and execution resources.
+            billing_hook: Optional billing hook for pipeline success/error callbacks.
+        """
+        self._task_context = ctx
+        self._billing_hook = billing_hook
+
+    async def handle_task(self) -> None:
+        try:
+            await self.handle()
+        finally:
+            task_id = self._task_context.id
+            task_tenant_id = self._task_context.tenant_id
+            task_dataset_id = self._task_context.kb_id
+            task_doc_id = self._task_context.doc_id
+            if self._task_context.has_canceled_func(task_id):
+                try:
+                    exists = await thread_pool_exec(
+                        settings.docStoreConn.index_exist,
+                        search.index_name(task_tenant_id),
+                        task_dataset_id,
+                    )
+                    if exists:
+                        ret = await thread_pool_exec(
+                            settings.docStoreConn.delete,
+                            {"doc_id": task_doc_id},
+                            search.index_name(task_tenant_id),
+                            task_dataset_id,
+                        )
+                        self._task_context.recording_context.save_func_return_value("docStoreConn.delete", ret)
+                except Exception as e:
+                    logging.exception(
+                        f"Remove doc({task_doc_id}) from docStore failed when task({task_id}) canceled, exception: {e}")
+
+    async def handle(self) -> None:
+        """Handle a document processing task."""
+        ctx = self._task_context
+        task_type = ctx.task_type
+        task_id = ctx.id
+
+        # Handle memory tasks
+        if task_type == "memory":
+            # ignore when it's dry run - no change on handle_save_to_memory_task when refactor
+            if isinstance(ctx.write_interceptor, RecordingContext):
+                logging.info(f"dry run, ignore handle_save_to_memory_task {task_id}")
+            else:
+                # actual run - not dry run
+                await handle_save_to_memory_task(ctx.raw_task)
+
+        # Handle dataflow debug mode
+        if task_type == "dataflow" and ctx.doc_id == CANVAS_DEBUG_DOC_ID:
+            await self._run_dataflow()
+            return
+
+        if task_type.startswith("dataflow"):
+            await self._run_dataflow()
+            return
+
+        # Check if task is canceled
+        if ctx.has_canceled_func(task_id):
+            ctx.progress_cb(-1, msg="Task has been canceled.")
+            return
+
+        # Bind embedding model
+        embedding_model = await self._bind_embedding_model()
+        if embedding_model is None:
+            return
+
+        with embedding_model:
+            vector_size = self._get_vector_size(embedding_model)
+            self._init_kb(vector_size)
+
+            # Route to appropriate handler
+            if task_type == "raptor":
+                await self._run_raptor(embedding_model, vector_size)
+            elif task_type == "graphrag":
+                await self._run_graphrag(embedding_model)
+            elif task_type == "mindmap":
+                ctx.progress_cb(1, "place holder")
+            elif task_type == "evaluation":
+                await self._run_evaluation()
+            elif task_type == "reembedding":
+                await self._run_reembedding()
+            elif task_type == "clone":
+                await self._run_clone()
+            else:
+                await self._run_standard_chunking(embedding_model)
+
+
+    @classmethod
+    def _get_vector_size(cls, embedding_model: LLMBundle) -> int:
+        """Get vector size from embedding model."""
+        vts, _ = embedding_model.encode(["ok"])
+        return len(vts[0])
+
+    def _init_kb(self, vector_size: int) -> None:
+        """Initialize knowledge base index."""
+        ctx = self._task_context
+        idxnm = search.index_name(ctx.tenant_id)
+        parser_id = ctx.parser_id
+        # Create index if not exists
+        settings.docStoreConn.create_idx(idxnm, ctx.kb_id, vector_size, parser_id)
+
+    async def _run_dataflow(self) -> None:
+        """Run dataflow pipeline."""
+        dataflow_service = DataflowService(
+            ctx=self._task_context,
+            billing_hook=self._billing_hook,
+        )
+        await dataflow_service.run_dataflow()
+
+    async def _run_evaluation(self) -> None:
+        """Run evaluation task."""
+        ctx = self._task_context
+        ctx.progress_cb(1, "Evaluation task placeholder")
+
+    async def _run_reembedding(self) -> None:
+        """Run reembedding task."""
+        ctx = self._task_context
+        ctx.progress_cb(1, "Reembedding task placeholder")
+
+    async def _run_clone(self) -> None:
+        """Run clone task."""
+        ctx = self._task_context
+        ctx.progress_cb(1, "Clone task placeholder")
+
+    async def _bind_embedding_model(self) -> Optional[LLMBundle]:
+        """Bind embedding model to task."""
+        ctx = self._task_context
+        task_tenant_id = ctx.tenant_id
+        task_embedding_id = ctx.embd_id
+        task_language = ctx.language
+
+        try:
+            if task_embedding_id:
+                embd_model_config = get_model_config_by_type_and_name(
+                    task_tenant_id, LLMType.EMBEDDING, task_embedding_id
+                )
+            else:
+                embd_model_config = get_tenant_default_model_by_type(
+                    task_tenant_id, LLMType.EMBEDDING
+                )
+            embedding_model = LLMBundle(task_tenant_id, embd_model_config, lang=task_language)
+            vts, _ = embedding_model.encode(["ok"])
+            return embedding_model
+        except Exception as e:
+            error_message = f'Fail to bind embedding model: {str(e)}'
+            ctx.progress_cb(-1, msg=error_message)
+            logging.exception(error_message)
+            raise
+
+    async def _run_raptor(
+        self,
+        embedding_model: LLMBundle,
+        vector_size: int,
+    ) -> None:
+        """Run RAPTOR summary generation."""
+        ctx = self._task_context
+        task_tenant_id = ctx.tenant_id
+        task_dataset_id = ctx.kb_id
+        kb_task_llm_id = ctx.kb_parser_config.get("llm_id") or ctx.llm_id
+
+        ok, kb = KnowledgebaseService.get_by_id(task_dataset_id)
+        if not ok:
+            ctx.progress_cb(prog=-1.0, msg="Cannot found valid dataset for RAPTOR task")
+            return
+
+        kb_parser_config = kb.parser_config
+        if not kb_parser_config.get("raptor", {}).get("use_raptor", False):
+            kb_parser_config.update({
+                "raptor": {
+                    "use_raptor": True,
+                    "prompt": "Please summarize the following paragraphs. Be careful with the numbers, do not make things up. Paragraphs as following:\n      {cluster_content}\nThe above is the content you need to summarize.",
+                    "max_token": 256,
+                    "threshold": 0.1,
+                    "max_cluster": 64,
+                    "random_seed": 0,
+                    "scope": "file",
+                    "clustering_method": "gmm",
+                    "tree_builder": "raptor",
+                },
+            })
+            if ctx.write_interceptor:
+                update_result = ctx.write_interceptor.intercept("KnowledgebaseService.update_by_id")
+            else:
+                update_result = KnowledgebaseService.update_by_id(kb.id, {"parser_config": kb_parser_config})
+
+            if not update_result:
+                ctx.progress_cb(prog=-1.0, msg="Internal error: Invalid RAPTOR configuration")
+                return
+
+        # Bind LLM for raptor
+        chat_model_config = get_model_config_by_type_and_name(
+            task_tenant_id, LLMType.CHAT, kb_task_llm_id
+        )
+        with LLMBundle(task_tenant_id, chat_model_config, lang=ctx.language) as chat_model:
+
+            # Run RAPTOR
+            raptor_service = RaptorService(ctx=ctx)
+
+            async with ctx.kg_limiter:
+                chunks, token_count, raptor_cleanup_chunks = await raptor_service.run_raptor_for_kb(
+                    kb_parser_config=kb_parser_config,
+                    chat_mdl=chat_model,
+                    embd_mdl=embedding_model,
+                    vector_size=vector_size,
+                    doc_ids=ctx.doc_ids,
+                )
+
+            ctx.recording_context.record("raptor_chunks", chunks)
+            ctx.recording_context.record("raptor_token_count", token_count)
+
+            # Insert RAPTOR chunks
+            if chunks:
+                task_doc_id = (ctx.doc_ids or [GRAPH_RAPTOR_FAKE_DOC_ID])[0]
+                chunk_service = ChunkService(ctx=ctx)
+                insert_result = await chunk_service.insert_chunks(ctx.id, task_tenant_id, task_dataset_id, chunks)
+                if insert_result:
+                    ctx.recording_context.record("insertion_result", "success")
+                else:
+                    ctx.recording_context.record("insertion_result", "failed")
+
+                # Cleanup stale RAPTOR chunks
+                cleaned_chunks = 0
+                for cleanup_doc_id, keep_method in raptor_cleanup_chunks:
+                    ret = await self._delete_raptor_chunks(
+                        cleanup_doc_id, task_tenant_id, task_dataset_id, keep_method
+                    )
+                    cleaned_chunks += ret
+
+                if cleaned_chunks:
+                    ctx.progress_cb(msg=f"Cleaned up {cleaned_chunks} stale RAPTOR chunks.")
+
+                # Update document stats
+                if ctx.write_interceptor:
+                    ctx.write_interceptor.intercept("DocumentService.increment_chunk_num")
+                else:
+                    DocumentService.increment_chunk_num(task_doc_id, task_dataset_id, token_count, len(chunks), 0)
+
+            ctx.recording_context.record("task_status", "completed")
+            ctx.progress_cb(prog=1.0, msg="RAPTOR done")
+
+    async def _run_graphrag(
+        self,
+        embedding_model: LLMBundle
+    ) -> None:
+        """Run GraphRAG."""
+        ctx = self._task_context
+        task_tenant_id = ctx.tenant_id
+        task_dataset_id = ctx.kb_id
+        kb_task_llm_id = ctx.kb_parser_config.get("llm_id") or ctx.llm_id
+        task_language = ctx.language
+
+        ok, kb = KnowledgebaseService.get_by_id(task_dataset_id)
+        if not ok:
+            ctx.progress_cb(prog=-1.0, msg="Cannot found valid dataset for GraphRAG task")
+            return
+
+        kb_parser_config = kb.parser_config
+        if not kb_parser_config.get("graphrag", {}).get("use_graphrag", False):
+            kb_parser_config.update({
+                "graphrag": {
+                    "use_graphrag": True,
+                    "entity_types": ["organization", "person", "geo", "event", "category"],
+                    "method": "light",
+                }
+            })
+            if ctx.write_interceptor:
+                update_result = ctx.write_interceptor.intercept("KnowledgebaseService.update_by_id")
+            else:
+                update_result = KnowledgebaseService.update_by_id(kb.id, {"parser_config": kb_parser_config})
+            if not update_result:
+                ctx.progress_cb(prog=-1.0, msg="Internal error: Invalid GraphRAG configuration")
+                return
+
+        graphrag_conf = kb_parser_config.get("graphrag", {})
+        start_ts = timer()
+        chat_model_config = get_model_config_by_type_and_name(
+            task_tenant_id, LLMType.CHAT, kb_task_llm_id
+        )
+        with LLMBundle(task_tenant_id, chat_model_config, lang=task_language) as chat_model:
+
+            with_resolution = graphrag_conf.get("resolution", False)
+            with_community = graphrag_conf.get("community", False)
+
+            async with ctx.kg_limiter:
+                result = await run_graphrag_for_kb(
+                    row=ctx.raw_task,
+                    doc_ids=ctx.doc_ids,
+                    language=task_language,
+                    kb_parser_config=kb_parser_config,
+                    chat_model=chat_model,
+                    embedding_model=embedding_model,
+                    callback=ctx.progress_cb,
+                    with_resolution=with_resolution,
+                    with_community=with_community,
+                )
+                logging.info(f"GraphRAG task result for task {ctx.raw_task}:\n{result}")
+
+            ctx.recording_context.record("graphrag_result", result)
+            ctx.progress_cb(prog=1.0, msg="Knowledge Graph done ({:.2f}s)".format(timer() - start_ts))
+
+    async def _run_standard_chunking(
+        self,
+        embedding_model: LLMBundle
+    ) -> None:
+        """Run standard chunking pipeline."""
+        ctx = self._task_context
+        task_id = ctx.id
+        task_tenant_id = ctx.tenant_id
+        task_dataset_id = ctx.kb_id
+        task_doc_id = ctx.doc_id
+        task_start_ts = timer()
+        doc_task_llm_id = ctx.parser_config.get("llm_id") or ctx.llm_id
+        ctx.raw_task['llm_id'] = doc_task_llm_id
+
+        # Build chunks
+        start_ts = timer()
+        chunk_service = ChunkService(ctx=ctx)
+
+        # Get storage binary
+        bucket, name = File2DocumentService.get_storage_address(doc_id=ctx.doc_id)
+        binary = await self._get_storage_binary(bucket, name)
+
+        chunks = await chunk_service.build_chunks(binary)
+        ctx.recording_context.record("chunks", chunks)
+        chunk_ids = [c.get("id") for c in chunks if isinstance(c, dict) and "id" in c]
+        ctx.recording_context.record("chunk_ids_count", len(chunk_ids))
+
+        logging.info("Build document {}: {:.2f}s".format(ctx.name, timer() - start_ts))
+
+        if not chunks:
+            ctx.progress_cb(1., msg=f"No chunk built from {ctx.name}")
+            return
+
+        ctx.progress_cb(msg="Generate {} chunks".format(len(chunks)))
+
+        # Embed chunks
+        start_ts = timer()
+        embedding_service = EmbeddingService(ctx=ctx)
+        try:
+            token_count, vector_size = embedding_service.embed_chunks(
+                chunks, embedding_model, ctx.parser_config
+            )
+        except TaskCanceledException:
+            raise
+        except Exception as e:
+            error_message = "Generate embedding error:{}".format(str(e))
+            ctx.progress_cb(-1, error_message)
+            logging.exception(error_message)
+            raise
+
+        ctx.recording_context.record("token_count", token_count)
+        ctx.recording_context.record("vector_size", vector_size)
+        progress_message = "Embedding chunks ({:.2f}s)".format(timer() - start_ts)
+        logging.info(progress_message)
+        ctx.progress_cb(msg=progress_message)
+
+        # Build TOC if needed
+        toc_thread = None
+        if ctx.parser_id.lower() == "naive" and ctx.parser_config.get("toc_extraction", False):
+            toc_thread = asyncio.create_task(asyncio.to_thread(self._build_toc, ctx, chunks, ctx.progress_cb))
+
+        # Insert chunks
+        chunk_count = len(set([chunk["id"] for chunk in chunks]))
+        start_ts = timer()
+
+        chunk_service = ChunkService(ctx=ctx)
+
+        if ctx.has_canceled_func(task_id):
+            ctx.progress_cb(-1, msg="Task has been canceled.")
+            return
+
+        insert_result = await chunk_service.insert_chunks(
+            task_id, task_tenant_id, task_dataset_id, chunks
+        )
+
+        if not insert_result:
+            ctx.recording_context.record("insertion_result", "failed")
+            return
+        ctx.recording_context.record("insertion_result", "success")
+
+        # Post-processing
+        post_processor = PostProcessor(ctx=ctx)
+        await post_processor.process_table_parser_metadata(task_doc_id, chunks)
+
+        ctx.progress_cb(msg="Indexing done ({:.2f}s).".format(timer() - start_ts))
+
+        toc_chunk = await self._process_toc_thread(toc_thread)
+        if toc_chunk:
+            ctx.recording_context.record("toc_chunk", [toc_chunk])
+            await post_processor.insert_toc_chunk(toc_chunk, chunk_service)
+
+        if ctx.has_canceled_func(task_id):
+            ctx.progress_cb(-1, msg="Task has been canceled.")
+            return
+
+        # Update document stats
+        if ctx.write_interceptor:
+            ctx.write_interceptor.intercept("DocumentService.increment_chunk_num")
+        else:
+            DocumentService.increment_chunk_num(task_doc_id, task_dataset_id, token_count, chunk_count, 0)
+
+        task_time_cost = timer() - task_start_ts
+        ctx.recording_context.record("task_status", "completed")
+        ctx.progress_cb(prog=1.0, msg="Task done ({:.2f}s)".format(task_time_cost))
+
+        logging.info(
+            "Chunk doc({}), page({}-{}), chunks({}), token({}), elapsed:{:.2f}".format(
+                ctx.name, ctx.from_page, ctx.to_page,
+                len(chunks), token_count, task_time_cost
+            )
+        )
+
+
+    async def _process_toc_thread(self, toc_thread):
+        try:
+            if toc_thread:
+                return await toc_thread
+            else:
+                return None
+        finally:
+            if toc_thread is not None and not toc_thread.done():
+                toc_thread.cancel()
+
+    @classmethod
+    async def _get_storage_binary(cls, bucket: str, name: str) -> bytes:
+        from common import settings
+        """Get binary from storage."""
+        return await thread_pool_exec(settings.STORAGE_IMPL.get, bucket, name)
+
+    @classmethod
+    def _build_toc(cls, ctx: TaskContext, docs: List[Dict], progress_cb: Callable) -> Optional[Dict]:
+        """Build table of contents."""
+        progress_cb(msg="Start to generate table of content ...")
+        chat_model_config = get_model_config_by_type_and_name(
+            ctx.tenant_id, LLMType.CHAT, ctx.llm_id
+        )
+        with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_mdl:
+
+            docs = sorted(docs, key=lambda d: (
+                d.get("page_num_int", 0)[0] if isinstance(d.get("page_num_int", 0), list) else d.get("page_num_int", 0),
+                d.get("top_int", 0)[0] if isinstance(d.get("top_int", 0), list) else d.get("top_int", 0)
+            ))
+
+            # NOTE: asyncio.run() creates a new event loop in the worker thread
+            # (this method is called via asyncio.to_thread), which is the
+            # intended pattern for bridging sync -> async in a thread context.
+            toc: list[dict] = asyncio.run(
+                run_toc_from_text([d["content_with_weight"] for d in docs], chat_mdl, progress_cb)
+            )
+            logging.info("------------ T O C -------------\n" + json.dumps(toc, ensure_ascii=False, indent='  '))
+
+            for ii, item in enumerate(toc):
+                try:
+                    chunk_val = item.pop("chunk_id", None)
+                    if chunk_val is None or str(chunk_val).strip() == "":
+                        logging.warning(f"Index {ii}: chunk_id is missing or empty. Skipping.")
+                        continue
+                    curr_idx = int(chunk_val or -1)
+                    if curr_idx >= len(docs):
+                        logging.error(f"Index {ii}: chunk_id {curr_idx} exceeds docs length {len(docs)}.")
+                        continue
+                    item["ids"] = [docs[curr_idx]["id"]]
+                    if ii + 1 < len(toc):
+                        next_chunk_val = toc[ii + 1].get("chunk_id", "")
+                        if str(next_chunk_val).strip() != "":
+                            next_idx = int(next_chunk_val)
+                            for jj in range(curr_idx + 1, min(next_idx + 1, len(docs))):
+                                item["ids"].append(docs[jj]["id"])
+                        else:
+                            logging.warning(f"Index {ii + 1}: next chunk_id is empty, range fill skipped.")
+                except (ValueError, TypeError) as e:
+                    logging.error(f"Index {ii}: Data conversion error - {e}")
+                except Exception as e:
+                    logging.exception(f"Index {ii}: Unexpected error - {e}")
+
+            if toc:
+                import copy
+                d = copy.deepcopy(docs[-1])
+                d["content_with_weight"] = json.dumps(toc, ensure_ascii=False)
+                d["toc_kwd"] = "toc"
+                d["available_int"] = 0
+                d["page_num_int"] = [100000000]
+                d["id"] = xxhash.xxh64(
+                    (d["content_with_weight"] + str(d["doc_id"])).encode("utf-8", "surrogatepass")).hexdigest()
+                return d
+            return None
+
+    async def _delete_raptor_chunks(
+        self, doc_id: str, tenant_id: str, kb_id: str, keep_method: Optional[str]
+    ) -> int:
+        """Delete RAPTOR chunks."""
+        if self._task_context.write_interceptor:
+            return self._task_context.write_interceptor.intercept("delete_raptor_chunks")
+        else:
+            return await delete_raptor_chunks(doc_id, tenant_id, kb_id, keep_method)

rag/svr/task_executor_refactor/task_manager.py

@@ -0,0 +1,177 @@
+#
+#  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.
+
+"""
+Task Manager Module.
+
+Provides [`TaskManager`](rag/svr/task_executor_refactor/task_manager.py:50) as the entry point
+for executing document processing tasks, supporting both production and dry-run (comparison) modes.
+"""
+
+import logging
+from typing import Any, Optional
+
+from rag.svr.task_executor_refactor.comparator import ContextComparator
+from rag.svr.task_executor_refactor.task_context import TaskCallbacks, TaskDict, TaskLimiters
+from rag.svr.task_executor_refactor.dataflow_service import BillingHook
+from rag.svr.task_executor_refactor.recording_context import (
+    BaseRecordingContext,
+    RecordingContext,
+    _NULL_RECORDING_CONTEXT,
+    set_recording_context, recording_context_manager,
+)
+from rag.svr.task_executor_refactor.task_context import TaskContext
+from rag.svr.task_executor_refactor.task_handler import TaskHandler
+from rag.svr.task_executor_refactor.write_operation_interceptor import (
+    WriteOperationInterceptor,
+)
+
+
+class TaskManager:
+    """Entry point for executing document processing tasks.
+
+    This class provides methods for:
+    - Production task execution (run_refactored_task)
+    - Dry-run task execution with comparison (dry_run_task)
+
+    Usage:
+        manager = TaskManager()
+        await manager.run_refactored_task(task, chat_limiter, ...)
+        # or
+        await manager.dry_run_task(task, recording_ctx1, ...)
+    """
+
+    @classmethod
+    async def run_refactored_task(
+        cls,
+        task: dict,
+        chat_limiter: Any,
+        minio_limiter: Any,
+        chunk_limiter: Any,
+        embed_limiter: Any,
+        kg_limiter: Any,
+        set_progress: Any,
+        has_canceled: Any,
+        billing_hook: Optional[BillingHook] = None,
+    ) -> None:
+        """Run a document processing task in production mode.
+
+        Args:
+            task: Task configuration dictionary.
+            chat_limiter: Rate limiter for chat operations.
+            minio_limiter: Rate limiter for MinIO operations.
+            chunk_limiter: Rate limiter for chunking operations.
+            embed_limiter: Rate limiter for embedding operations.
+            kg_limiter: Rate limiter for knowledge graph operations.
+            set_progress: Progress callback function.
+            has_canceled: Function to check if task is canceled.
+            billing_hook: Optional billing hook for pipeline success/error callbacks.
+        """
+        with recording_context_manager(_NULL_RECORDING_CONTEXT):
+            # Use NullRecordingContext in production to avoid memory allocation
+            set_recording_context(_NULL_RECORDING_CONTEXT)
+
+            # Create TaskContext with all execution resources
+            task_context = TaskContext(
+                task=task,
+                limiters=TaskLimiters(
+                    chat=chat_limiter,
+                    minio=minio_limiter,
+                    chunk=chunk_limiter,
+                    embed=embed_limiter,
+                    kg=kg_limiter,
+                ),
+                callbacks=TaskCallbacks(
+                    progress=set_progress,
+                    has_canceled=has_canceled,
+                ),
+                recording_context=_NULL_RECORDING_CONTEXT,
+            )
+
+            # Execute with TaskHandler
+            handler = TaskHandler(ctx=task_context, billing_hook=billing_hook)
+            await handler.handle_task()
+
+    @classmethod
+    async def dry_run_task(
+        cls,
+        task: TaskDict,
+        recording_ctx1: BaseRecordingContext,
+        chat_limiter: Any,
+        minio_limiter: Any,
+        chunk_limiter: Any,
+        embed_limiter: Any,
+        kg_limiter: Any,
+        set_progress: Any,
+        has_canceled: Any,
+    ) -> None:
+        """Run a document processing task in dry-run mode for comparison.
+
+        This executes the task with a write operation interceptor that records
+        all write operations, then compares the results with the production run.
+
+        Args:
+            task: Task configuration dictionary.
+            recording_ctx1: RecordingContext from production execution.
+            chat_limiter: Rate limiter for chat operations.
+            minio_limiter: Rate limiter for MinIO operations.
+            chunk_limiter: Rate limiter for chunking operations.
+            embed_limiter: Rate limiter for embedding operations.
+            kg_limiter: Rate limiter for knowledge graph operations.
+            set_progress: Progress callback function.
+            has_canceled: Function to check if task is canceled.
+        """
+        interceptor = WriteOperationInterceptor(recording_ctx1.get_all_func_return_values())
+        recording_ctx2 = RecordingContext()
+
+        with recording_context_manager(recording_ctx2):
+            set_recording_context(recording_ctx2)
+
+            # Create TaskContext with all execution resources
+            task_context = TaskContext(
+                task=task,
+                limiters=TaskLimiters(
+                    chat=chat_limiter,
+                    minio=minio_limiter,
+                    chunk=chunk_limiter,
+                    embed=embed_limiter,
+                    kg=kg_limiter,
+                ),
+                callbacks=TaskCallbacks(
+                    progress=set_progress,
+                    has_canceled=has_canceled,
+                ),
+                write_interceptor=interceptor,
+                recording_context=recording_ctx2,
+            )
+
+            # Execute with TaskHandler
+            handler = TaskHandler(ctx=task_context)
+            await handler.handle_task()
+
+            # Compare results
+            comp: ContextComparator = ContextComparator()
+            comp_result = comp.compare(task_context.id, recording_ctx1, recording_ctx2)
+            logging.info(f"-------{task_context.name}, compare result:{comp_result.to_markdown()}")
+            if interceptor.remaining_values_count() > 0 or comp_result.mismatched_keys > 0:
+                logging.info(f"------task:{task_context.id} {task_context.name} differs, "
+                             f"interceptor.remaining_values_count():{interceptor.remaining_values_count()}, "
+                             f"mismatched_keys:{comp_result.mismatched_keys}")
+                if interceptor.remaining_values_count() > 0:
+                    logging.info(f"------task:{task_context.id}, remaining values:{interceptor.remaining_values()}")
+                if comp_result.mismatched_keys > 0:
+                    logging.info(f"-------compare result:{comp_result.details}")
+            else:
+                logging.info(f"------task:{task_context.id} {task_context.name} same result for prod and dry run ")

rag/svr/task_executor_refactor/write_operation_interceptor.py

@@ -0,0 +1,138 @@
+#
+#  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.
+
+"""
+Write Operation Interceptor Module
+
+Provides a mechanism to intercept write operations during comparison mode.
+The interceptor consumes pre-recorded return values (from production execution)
+and returns them one by one when the corresponding methods are called.
+"""
+import logging
+from typing import Any, Dict, List
+
+# Set of allowed method names that can be intercepted
+ALLOWED_METHOD_NAMES = {
+    "KnowledgebaseService.update_by_id",
+    "TaskService.update_chunk_ids",
+    "DocumentService.increment_chunk_num",
+    "DocMetadataService.update_document_metadata",
+    "PipelineOperationLogService.record_pipeline_operation",
+    "PipelineOperationLogService.create",
+    "delete_raptor_chunks",
+    "handle_save_to_memory_task",
+    "docStoreConn.insert",
+    "docStoreConn.delete"
+}
+
+_NO_DEFAULT = object()
+
+
+class WriteOperationInterceptor:
+    """Intercepts write operations and returns pre-recorded values.
+
+    This interceptor is used in comparison mode to replay production execution
+    results. When a method is called, the interceptor pops the first recorded
+    return value from the corresponding list and returns it.
+
+    Usage:
+        # Create interceptor with pre-recorded values
+        interceptor = WriteOperationInterceptor({
+            "build_chunks": [chunks1, chunks2],
+            "embedding": [(token_count1, vector_size1)],
+            ...
+        })
+
+        # Intercept a method call
+        result = interceptor.intercept("build_chunks")  # Returns chunks1
+        result = interceptor.intercept("build_chunks")  # Returns chunks2
+    """
+
+    def __init__(self, recorded_values: Dict[str, List[Any]]):
+        """Initialize the interceptor with pre-recorded values.
+
+        Args:
+            recorded_values: A dictionary where keys are method names and
+                values are lists of pre-recorded return values. Each call
+                to intercept() will pop and return the first value from
+                the corresponding list.
+
+        Note:
+            If a key from ALLOWED_METHOD_NAMES is not in recorded_values,
+            it will be initialized with an empty list. This allows the
+            interceptor to be created even if not all methods have recorded
+            values, and it will fall through to original execution when
+            no recorded values are available.
+        """
+        self._recorded_values: Dict[str, List[Any]] = dict()
+        for key in ALLOWED_METHOD_NAMES:
+            self._recorded_values[key] = list(recorded_values.get(key, []))
+
+    def intercept(self, method_name: str, default_value = _NO_DEFAULT) -> Any:
+        """Intercept a method call and return the next pre-recorded value.
+
+        Args:
+            method_name: Name of the method being intercepted.
+            default_value: default value
+
+        Returns:
+            The next pre-recorded return value for this method.
+
+        Raises:
+            ValueError: If method_name is not in the allowed method names set.
+            KeyError: If method_name has no recorded values list.
+            IndexError: If the recorded values list for method_name is empty.
+        """
+        if method_name not in ALLOWED_METHOD_NAMES:
+            raise ValueError(
+                f"Cannot intercept method '{method_name}'. "
+                f"Allowed method names: {ALLOWED_METHOD_NAMES}"
+            )
+
+        if method_name not in self._recorded_values:
+            raise KeyError(f"No recorded values found for method '{method_name}'")
+
+        values_list = self._recorded_values[method_name]
+        if not values_list:
+            if default_value is not _NO_DEFAULT:
+                logging.info(f"return default value for {method_name}")
+                return default_value
+            raise IndexError(f"No more recorded values for method '{method_name}'")
+
+        return values_list.pop(0)
+
+
+    def remaining_count(self, method_name: str) -> int:
+        """Get the number of remaining recorded values for a method.
+
+        Args:
+            method_name: Name of the method to check.
+
+        Returns:
+            Number of remaining recorded values.
+        """
+        if method_name not in self._recorded_values:
+            return 0
+        return len(self._recorded_values[method_name])
+
+
+    def remaining_values(self):
+        return {k: list(v) for k, v in self._recorded_values.items()}
+
+    def remaining_values_count(self):
+        return sum(len(values) for values in self._recorded_values.values())
+
+    def __repr__(self) -> str:
+        return f"WriteOperationInterceptor(total_recorded={self._recorded_values})"

rag/utils/base64_image.py

@@ -33,7 +33,7 @@ test_image = base64.b64decode(test_image_base64)
 async def image2id(d: dict, storage_put_func: partial, objname: str, bucket: str = "imagetemps"):
     import logging
     from io import BytesIO
-    from rag.svr.task_executor import minio_limiter
+    from rag.svr.task_executor_limiter import minio_limiter
 
     if "image" not in d:
         return

rag/utils/raptor_utils.py

@@ -62,16 +62,28 @@ def _as_extra_dict(extra) -> dict:
     if isinstance(extra, dict):
         return extra
     if isinstance(extra, str) and extra:
+        # Try standard JSON first (double quotes)
         try:
             parsed = json.loads(extra)
+            return parsed if isinstance(parsed, dict) else {}
         except json.JSONDecodeError:
-            logging.warning(
-                "Ignoring malformed RAPTOR extra payload while collecting chunk metadata: %s",
-                extra[:200],
-                exc_info=True,
-            )
-            return {}
-        return parsed if isinstance(parsed, dict) else {}
+            last_exc = True
+
+        # Fallback: try parsing Python dict literal (single quotes)
+        try:
+            import ast
+            parsed = ast.literal_eval(extra)
+            if isinstance(parsed, dict):
+                return parsed
+        except (ValueError, SyntaxError):
+            last_exc = True
+
+        logging.warning(
+            "Ignoring malformed RAPTOR extra payload while collecting chunk metadata: %s",
+            extra[:200],
+            exc_info=last_exc,
+        )
+        return {}
     return {}