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>

api/db/joint_services/memory_message_service.py

@@ -161,19 +161,19 @@ async def extract_by_llm(tenant_id: str, tenant_llm_id: int, extract_conf: dict,
         )
     else:
         llm_config = get_model_config_by_type_and_name(tenant_id, LLMType.CHAT, llm_id)
-    llm = LLMBundle(tenant_id, llm_config)
-    if task_id:
-        TaskService.update_progress(task_id, {"progress": 0.15, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Prepared prompts and LLM."})
-    res = await llm.async_chat(system_prompt, user_prompts, extract_conf)
-    res_json = get_json_result_from_llm_response(res)
-    if task_id:
-        TaskService.update_progress(task_id, {"progress": 0.35, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Get extracted result from LLM."})
-    return [{
-        "content": extracted_content["content"],
-        "valid_at": format_iso_8601_to_ymd_hms(extracted_content["valid_at"]),
-        "invalid_at": format_iso_8601_to_ymd_hms(extracted_content["invalid_at"]) if extracted_content.get("invalid_at") else "",
-        "message_type": message_type
-    } for message_type, extracted_content_list in res_json.items() for extracted_content in extracted_content_list]
+    with LLMBundle(tenant_id, llm_config) as llm:
+        if task_id:
+            TaskService.update_progress(task_id, {"progress": 0.15, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Prepared prompts and LLM."})
+        res = await llm.async_chat(system_prompt, user_prompts, extract_conf)
+        res_json = get_json_result_from_llm_response(res)
+        if task_id:
+            TaskService.update_progress(task_id, {"progress": 0.35, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Get extracted result from LLM."})
+        return [{
+            "content": extracted_content["content"],
+            "valid_at": format_iso_8601_to_ymd_hms(extracted_content["valid_at"]),
+            "invalid_at": format_iso_8601_to_ymd_hms(extracted_content["invalid_at"]) if extracted_content.get("invalid_at") else "",
+            "message_type": message_type
+        } for message_type, extracted_content_list in res_json.items() for extracted_content in extracted_content_list]
 
 
 async def embed_and_save(memory, message_list: list[dict], task_id: str=None):
@@ -185,48 +185,48 @@ async def embed_and_save(memory, message_list: list[dict], task_id: str=None):
         )
     else:
         embd_model_config = get_model_config_by_type_and_name(memory.tenant_id, LLMType.EMBEDDING, memory.embd_id)
-    embedding_model = LLMBundle(memory.tenant_id, embd_model_config)
-    if task_id:
-        TaskService.update_progress(task_id, {"progress": 0.65, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Prepared embedding model."})
-    vector_list, _ = embedding_model.encode([msg["content"] for msg in message_list])
-    for idx, msg in enumerate(message_list):
-        msg["content_embed"] = vector_list[idx]
-    if task_id:
-        TaskService.update_progress(task_id, {"progress": 0.85, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Embedded extracted content."})
-    vector_dimension = len(vector_list[0])
-    if not MessageService.has_index(memory.tenant_id, memory.id):
-        created = MessageService.create_index(memory.tenant_id, memory.id, vector_size=vector_dimension)
-        if not created:
-            error_msg = "Failed to create message index."
-            if task_id:
-                TaskService.update_progress(task_id, {"progress": -1, "progress_msg": timestamp_to_date(current_timestamp())+ " " + error_msg})
-            return False, error_msg
-
-    new_msg_size = sum([MessageService.calculate_message_size(m) for m in message_list])
-    current_memory_size = get_memory_size_cache(memory.tenant_id, memory.id)
-    if new_msg_size + current_memory_size > memory.memory_size:
-        size_to_delete = current_memory_size + new_msg_size - memory.memory_size
-        if memory.forgetting_policy == "FIFO":
-            message_ids_to_delete, delete_size = MessageService.pick_messages_to_delete_by_fifo(memory.id, memory.tenant_id,
-                                                                                                size_to_delete)
-            MessageService.delete_message({"message_id": message_ids_to_delete}, memory.tenant_id, memory.id)
-            decrease_memory_size_cache(memory.id, delete_size)
-        else:
-            error_msg = "Failed to insert message into memory. Memory size reached limit and cannot decide which to delete."
-            if task_id:
-                TaskService.update_progress(task_id, {"progress": -1, "progress_msg": timestamp_to_date(current_timestamp())+ " " + error_msg})
-            return False, error_msg
-    fail_cases = MessageService.insert_message(message_list, memory.tenant_id, memory.id)
-    if fail_cases:
-        error_msg = "Failed to insert message into memory. Details: " + "; ".join(fail_cases)
+    with LLMBundle(memory.tenant_id, embd_model_config) as embedding_model:
         if task_id:
-            TaskService.update_progress(task_id, {"progress": -1, "progress_msg": timestamp_to_date(current_timestamp())+ " " + error_msg})
-        return False, error_msg
+            TaskService.update_progress(task_id, {"progress": 0.65, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Prepared embedding model."})
+        vector_list, _ = embedding_model.encode([msg["content"] for msg in message_list])
+        for idx, msg in enumerate(message_list):
+            msg["content_embed"] = vector_list[idx]
+        if task_id:
+            TaskService.update_progress(task_id, {"progress": 0.85, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Embedded extracted content."})
+        vector_dimension = len(vector_list[0])
+        if not MessageService.has_index(memory.tenant_id, memory.id):
+            created = MessageService.create_index(memory.tenant_id, memory.id, vector_size=vector_dimension)
+            if not created:
+                error_msg = "Failed to create message index."
+                if task_id:
+                    TaskService.update_progress(task_id, {"progress": -1, "progress_msg": timestamp_to_date(current_timestamp())+ " " + error_msg})
+                return False, error_msg
 
-    if task_id:
-        TaskService.update_progress(task_id, {"progress": 0.95, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Saved messages to storage."})
-    increase_memory_size_cache(memory.id, new_msg_size)
-    return True, "Message saved successfully."
+        new_msg_size = sum([MessageService.calculate_message_size(m) for m in message_list])
+        current_memory_size = get_memory_size_cache(memory.tenant_id, memory.id)
+        if new_msg_size + current_memory_size > memory.memory_size:
+            size_to_delete = current_memory_size + new_msg_size - memory.memory_size
+            if memory.forgetting_policy == "FIFO":
+                message_ids_to_delete, delete_size = MessageService.pick_messages_to_delete_by_fifo(memory.id, memory.tenant_id,
+                                                                                                    size_to_delete)
+                MessageService.delete_message({"message_id": message_ids_to_delete}, memory.tenant_id, memory.id)
+                decrease_memory_size_cache(memory.id, delete_size)
+            else:
+                error_msg = "Failed to insert message into memory. Memory size reached limit and cannot decide which to delete."
+                if task_id:
+                    TaskService.update_progress(task_id, {"progress": -1, "progress_msg": timestamp_to_date(current_timestamp())+ " " + error_msg})
+                return False, error_msg
+        fail_cases = MessageService.insert_message(message_list, memory.tenant_id, memory.id)
+        if fail_cases:
+            error_msg = "Failed to insert message into memory. Details: " + "; ".join(fail_cases)
+            if task_id:
+                TaskService.update_progress(task_id, {"progress": -1, "progress_msg": timestamp_to_date(current_timestamp())+ " " + error_msg})
+            return False, error_msg
+
+        if task_id:
+            TaskService.update_progress(task_id, {"progress": 0.95, "progress_msg": timestamp_to_date(current_timestamp())+ " " + "Saved messages to storage."})
+        increase_memory_size_cache(memory.id, new_msg_size)
+        return True, "Message saved successfully."
 
 
 def query_message(filter_dict: dict, params: dict):

api/db/services/document_service.py

@@ -1098,12 +1098,12 @@ def queue_raptor_o_graphrag_tasks(sample_doc, ty, priority, fake_doc_id="", doc_
 
     task["doc_ids"] = doc_ids
     DocumentService.begin2parse(task["doc_id"], keep_progress=True)
-    assert REDIS_CONN.queue_product(settings.get_svr_queue_name(priority), message=task), "Can't access Redis. Please check the Redis' status."
+    assert REDIS_CONN.queue_product(settings.get_svr_queue_name(priority, ty), message=task), "Can't access Redis. Please check the Redis' status."
     return task["id"]
 
 
-def get_queue_length(priority):
-    group_info = REDIS_CONN.queue_info(settings.get_svr_queue_name(priority), SVR_CONSUMER_GROUP_NAME)
+def get_queue_length(priority, suffix="common"):
+    group_info = REDIS_CONN.queue_info(settings.get_svr_queue_name(priority, suffix), SVR_CONSUMER_GROUP_NAME)
     if not group_info:
         return 0
     return int(group_info.get("lag", 0) or 0)

api/db/services/llm_service.py

@@ -86,6 +86,19 @@ class LLMBundle(LLM4Tenant):
     def __init__(self, tenant_id: str, model_config: dict, lang="Chinese", **kwargs):
         super().__init__(tenant_id, model_config, lang, **kwargs)
 
+    def close(self):
+        """Release resources held by this LLMBundle instance."""
+        super().close()
+
+    def __enter__(self):
+        """Enter context manager."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit context manager and release resources."""
+        self.close()
+        return False
+
     def bind_tools(self, toolcall_session, tools):
         if not self.is_tools:
             logging.warning(f"Model {self.model_config['llm_name']} does not support tool call, but you have assigned one or more tools to it!")
@@ -124,7 +137,7 @@ class LLMBundle(LLM4Tenant):
 
         embeddings, used_tokens = self.mdl.encode(safe_texts)
         if self.model_config["llm_factory"] == "Builtin":
-            logging.info("LLMBundle.encode_queries query: {}, emd len: {}, used_tokens: {}. Builtin model don't need to update token usage".format(texts, len(embeddings), used_tokens))
+            logging.debug("LLMBundle.encode_queries query: {}, emd len: {}, used_tokens: {}. Builtin model don't need to update token usage".format(texts, len(embeddings), used_tokens))
         elif not TenantLLMService.increase_usage_by_id(self.model_config["id"], used_tokens):
             logging.error("LLMBundle.encode can't update token usage for <tenant redacted>/EMBEDDING used_tokens: {}".format(used_tokens))
 

api/db/services/task_service.py

@@ -419,6 +419,9 @@ def queue_tasks(doc: dict, bucket: str, name: str, priority: int):
     else:
         parse_task_array.append(new_task())
 
+    # Determine suffix based on parser_id (consistent with SAAS version line 444)
+    suffix = "common" if doc["parser_id"] != "resume" else "resume"
+
     chunking_config = DocumentService.get_chunking_config(doc["id"])
     for task in parse_task_array:
         hasher = xxhash.xxh64()
@@ -456,7 +459,7 @@ def queue_tasks(doc: dict, bucket: str, name: str, priority: int):
     unfinished_task_array = [task for task in parse_task_array if task["progress"] < 1.0]
     for unfinished_task in unfinished_task_array:
         assert REDIS_CONN.queue_product(
-            settings.get_svr_queue_name(priority), message=unfinished_task
+            settings.get_svr_queue_name(priority, suffix), message=unfinished_task
         ), "Can't access Redis. Please check the Redis' status."
 
 
@@ -547,7 +550,7 @@ def queue_dataflow(tenant_id:str, flow_id:str, task_id:str, doc_id:str=CANVAS_DE
     task["file"] = file
 
     if not REDIS_CONN.queue_product(
-            settings.get_svr_queue_name(priority), message=task
+            settings.get_svr_queue_name(priority, "common"), message=task
     ):
         return False, "Can't access Redis. Please check the Redis' status."
 

api/db/services/tenant_llm_service.py

@@ -520,3 +520,31 @@ class LLM4Tenant:
             except Exception:
                 # Skip langfuse tracing if connection fails
                 pass
+
+    def close(self):
+        """Release resources held by this LLM4Tenant instance.
+
+        This method should be called when the instance is no longer needed
+        to properly release resources such as:
+        - Langfuse tracing client (flush and shutdown)
+        - Underlying model instance resources (HTTP sessions, etc.)
+        """
+        # Flush and shutdown Langfuse client if it was initialized
+        if self.langfuse:
+            try:
+                self.langfuse.flush()
+                if hasattr(self.langfuse, 'shutdown'):
+                    self.langfuse.shutdown()
+            except Exception:
+                # Ignore errors during cleanup
+                pass
+            finally:
+                self.langfuse = None
+
+        # Release underlying model instance if it has a close method
+        if self.mdl and hasattr(self.mdl, 'close') and callable(getattr(self.mdl, 'close')):
+            try:
+                self.mdl.close()
+            except Exception:
+                # Ignore errors during cleanup
+                pass

common/constants.py

@@ -246,7 +246,7 @@ class ForgettingPolicy(StrEnum):
 # ENV_TRACE_MALLOC_ENABLED = "TRACE_MALLOC_ENABLED"
 
 PAGERANK_FLD = "pagerank_fea"
-SVR_QUEUE_NAME = "rag_flow_svr_queue"
+SVR_QUEUE_NAME = "te"
 SVR_CONSUMER_GROUP_NAME = "rag_flow_svr_task_broker"
 TAG_FLD = "tag_feas"
 

common/decorator.py

@@ -13,7 +13,12 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import functools
+import inspect
+import logging
 import os
+import time
+
 
 def singleton(cls, *args, **kw):
     instances = {}
@@ -24,4 +29,58 @@ def singleton(cls, *args, **kw):
             instances[key] = cls(*args, **kw)
         return instances[key]
 
-    return _singleton
+    return _singleton
+
+
+def timing(func=None, *, name=None, context=None):
+    """Decorator that records function execution time.
+
+    Usage:
+        @timing
+        async def my_func(): ...
+
+        @timing(name="custom_name")
+        def my_func(): ...
+
+        @timing(context=recording_ctx)
+        async def my_func(): ...
+
+    Args:
+        func: The function to decorate (auto-passed when used as @timing)
+        name: Custom name for the timing record, defaults to function name
+        context: A RecordingContext-like object to record timing data into.
+                 If not provided, will try to use global recording_context from task_executor.
+                 Timing data will be recorded as "{name}_time".
+    """
+    if func is None:
+        return functools.partial(timing, name=name, context=context)
+
+    func_name = name or func.__name__
+    log = logging.getLogger(__name__)
+
+    if inspect.iscoroutinefunction(func):
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            start = time.perf_counter()
+            try:
+                result = await func(*args, **kwargs)
+                return result
+            finally:
+                elapsed = time.perf_counter() - start
+                log.debug(f"[TIMING] {func_name} took {elapsed:.3f}s")
+                if context is not None:
+                    context.record(f"{func_name}_time", elapsed)
+        return async_wrapper
+    else:
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            start = time.perf_counter()
+            try:
+                result = func(*args, **kwargs)
+                return result
+            finally:
+                elapsed = time.perf_counter() - start
+                log.debug(f"[TIMING] {func_name} took {elapsed:.3f}s")
+                if context is not None:
+                    context.record(f"{func_name}_time", elapsed)
+        return sync_wrapper

common/settings.py

@@ -133,13 +133,30 @@ PARALLEL_DEVICES: int = 0
 STORAGE_IMPL_TYPE = os.getenv('STORAGE_IMPL', 'MINIO')
 STORAGE_IMPL = None
 
-def get_svr_queue_name(priority: int) -> str:
-    if priority == 0:
-        return SVR_QUEUE_NAME
-    return f"{SVR_QUEUE_NAME}_{priority}"
+def get_svr_queue_name(priority: int, suffix: str = "common") -> str:
+    """
+    Generate queue name with two dimensions: priority and suffix.
+    
+    Args:
+        priority: Task priority (0=low, 1=high)
+        suffix: Task type suffix (common/resume/graphrag/raptor/mindmap)
+               Currently only "common" is used, other suffixes are reserved.
+    
+    Returns:
+        Queue name string
+    
+    Examples:
+        get_svr_queue_name(0, "common") -> "te.0.common"
+        get_svr_queue_name(1, "common") -> "te.1.common"
+        get_svr_queue_name(0) -> "te.0.common"  # default suffix="common"
 
-def get_svr_queue_names():
-    return [get_svr_queue_name(priority) for priority in [1, 0]]
+    """
+    return f"{SVR_QUEUE_NAME}.{priority}.common"
+
+
+def get_svr_queue_names(suffix:str):
+    """Return queue names sorted by priority (high to low)."""
+    return [get_svr_queue_name(priority, suffix) for priority in [1, 0]]
 
 def init_secret_key():
     secret_key = os.environ.get("RAGFLOW_SECRET_KEY")

docker/entrypoint.sh

@@ -210,7 +210,7 @@ function task_exe() {
     JEMALLOC_PATH="$(pkg-config --variable=libdir jemalloc)/libjemalloc.so"
     while true; do
         LD_PRELOAD="$JEMALLOC_PATH" \
-        "$PY" rag/svr/task_executor.py "${host_id}_${consumer_id}"  &
+        "$PY" rag/svr/task_executor.py -i "${host_id}_${consumer_id}" -t "common" &
         wait;
         sleep 1;
     done

docker/launch_backend_service.sh

@@ -73,7 +73,7 @@ task_exe(){
     local retry_count=0
     while ! $STOP && [ $retry_count -lt $MAX_RETRIES ]; do
         echo "Starting task_executor.py for task $task_id (Attempt $((retry_count+1)))"
-        LD_PRELOAD=$JEMALLOC_PATH $PY rag/svr/task_executor.py "$task_id"
+        LD_PRELOAD=$JEMALLOC_PATH $PY rag/svr/task_executor.py -i "$task_id"
         EXIT_CODE=$?
         if [ $EXIT_CODE -eq 0 ]; then
             echo "task_executor.py for task $task_id exited successfully."

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 {}
 
 

test/unit_test/common/test_settings_queue.py

@@ -0,0 +1,208 @@
+#
+#  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.
+#
+"""Test cases for get_svr_queue_name and get_svr_queue_names functions in common.settings."""
+
+from common.settings import get_svr_queue_name, get_svr_queue_names
+
+
+class TestGetSvrQueueName:
+    """Test cases for get_svr_queue_name function."""
+
+    def test_default_suffix(self):
+        """Test that default suffix is 'common'."""
+
+        result = get_svr_queue_name(0)
+        assert result == "te.0.common"
+
+    def test_priority_zero(self):
+        """Test queue name with priority 0 (low)."""
+
+        result = get_svr_queue_name(0)
+        assert result == "te.0.common"
+
+    def test_priority_one(self):
+        """Test queue name with priority 1 (high)."""
+
+        result = get_svr_queue_name(1)
+        assert result == "te.1.common"
+
+    def test_explicit_suffix_common(self):
+        """Test with explicit 'common' suffix."""
+
+        result = get_svr_queue_name(0, "common")
+        assert result == "te.0.common"
+
+    def test_suffix_parameter_ignored(self):
+        """Test that suffix parameter is currently ignored (hardcoded to 'common').
+        
+        Note: The function signature accepts a suffix parameter but currently
+        hardcodes 'common' in the return value. This test documents this behavior.
+        """
+
+        # Even with different suffix values, result should be the same
+        result_default = get_svr_queue_name(0, "common")
+        result_resume = get_svr_queue_name(0, "resume")
+        result_graphrag = get_svr_queue_name(0, "graphrag")
+
+        # All should return the same value since suffix is hardcoded
+        assert result_default == result_resume == result_graphrag == "te.0.common"
+
+    def test_format_structure(self):
+        """Test that queue name follows expected format: {SVR_QUEUE_NAME}.{priority}.common."""
+
+        for priority in [0, 1]:
+            result = get_svr_queue_name(priority)
+            parts = result.split(".")
+            assert len(parts) == 3
+            assert parts[0] == "te"  # SVR_QUEUE_NAME
+            assert parts[1] == str(priority)
+            assert parts[2] == "common"
+
+    def test_different_priorities_produce_different_results(self):
+        """Test that different priorities produce different queue names."""
+
+        result_0 = get_svr_queue_name(0)
+        result_1 = get_svr_queue_name(1)
+
+        assert result_0 != result_1
+        assert result_0 == "te.0.common"
+        assert result_1 == "te.1.common"
+
+    def test_with_various_priority_values(self):
+        """Test with various priority values beyond 0 and 1."""
+
+        # Test with other priority values to ensure format is correct
+        for priority in [2, 5, 10, 100]:
+            result = get_svr_queue_name(priority)
+            expected = f"te.{priority}.common"
+            assert result == expected
+
+    def test_returns_string_type(self):
+        """Test that function returns a string."""
+
+        result = get_svr_queue_name(0)
+        assert isinstance(result, str)
+
+    def test_no_whitespace_issues(self):
+        """Test that queue name has no unexpected whitespace."""
+
+        for priority in [0, 1]:
+            result = get_svr_queue_name(priority)
+            assert " " not in result
+            assert "\t" not in result
+            assert "\n" not in result
+
+
+class TestGetSvrQueueNames:
+    """Test cases for get_svr_queue_names function."""
+
+    def test_returns_list(self):
+        """Test that function returns a list."""
+
+        result = get_svr_queue_names("common")
+        assert isinstance(result, list)
+
+    def test_returns_two_queues(self):
+        """Test that function returns exactly two queue names."""
+
+        result = get_svr_queue_names("common")
+        assert len(result) == 2
+
+    def test_sorted_high_to_low(self):
+        """Test that queue names are sorted from high priority to low priority."""
+
+        result = get_svr_queue_names("common")
+        assert result[0] == "te.1.common"  # High priority first
+        assert result[1] == "te.0.common"  # Low priority second
+
+    def test_expected_values(self):
+        """Test that returned values match expected queue names."""
+
+        result = get_svr_queue_names("common")
+        expected = ["te.1.common", "te.0.common"]
+        assert result == expected
+
+    def test_suffix_parameter_passed_through(self):
+        """Test that suffix parameter is passed to get_svr_queue_name.
+        
+        Note: Since get_svr_queue_name currently hardcodes 'common' as the suffix,
+        different suffix values will still produce the same result.
+        """
+
+        # All suffixes should produce same result due to hardcoded suffix in get_svr_queue_name
+        result_common = get_svr_queue_names("common")
+        result_resume = get_svr_queue_names("resume")
+        result_graphrag = get_svr_queue_names("graphrag")
+
+        expected = ["te.1.common", "te.0.common"]
+        assert result_common == expected
+        assert result_resume == expected  # suffix is currently ignored
+        assert result_graphrag == expected  # suffix is currently ignored
+
+    def test_all_elements_are_strings(self):
+        """Test that all elements in the returned list are strings."""
+
+        result = get_svr_queue_names("common")
+        for item in result:
+            assert isinstance(item, str)
+
+    def test_consistent_results(self):
+        """Test that multiple calls return consistent results."""
+
+        result1 = get_svr_queue_names("common")
+        result2 = get_svr_queue_names("common")
+        result3 = get_svr_queue_names("common")
+
+        assert result1 == result2 == result3
+
+    def test_with_empty_suffix(self):
+        """Test with empty string suffix."""
+
+        result = get_svr_queue_names("")
+        # Should still work since suffix is ignored
+        assert result == ["te.1.common", "te.0.common"]
+
+
+class TestGetSvrQueueNameWithMockedConstant:
+    """Test cases with mocked SVR_QUEUE_NAME constant."""
+
+    def test_with_custom_queue_name(self):
+        """Test with a custom SVR_QUEUE_NAME constant."""
+        # Need to patch where the constant is imported in settings module
+        import common.settings as settings_mod
+
+        original_value = settings_mod.SVR_QUEUE_NAME
+        try:
+            settings_mod.SVR_QUEUE_NAME = "custom_queue"
+            result = settings_mod.get_svr_queue_name(0)
+            assert result == "custom_queue.0.common"
+
+            result = settings_mod.get_svr_queue_name(1)
+            assert result == "custom_queue.1.common"
+        finally:
+            settings_mod.SVR_QUEUE_NAME = original_value
+
+    def test_with_custom_queue_names(self):
+        """Test get_svr_queue_names with a custom SVR_QUEUE_NAME constant."""
+        import common.settings as settings_mod
+
+        original_value = settings_mod.SVR_QUEUE_NAME
+        try:
+            settings_mod.SVR_QUEUE_NAME = "custom_queue"
+            result = settings_mod.get_svr_queue_names("common")
+            assert result == ["custom_queue.1.common", "custom_queue.0.common"]
+        finally:
+            settings_mod.SVR_QUEUE_NAME = original_value

test/unit_test/rag/svr/task_executor_refactor/conftest.py

@@ -0,0 +1,494 @@
+#
+#  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 pytest fixtures for task_executor_refactor integration tests.
+
+This module provides reusable fixtures for integration tests that verify
+the complete orchestration flow of TaskHandler and its collaborating services.
+
+Design principles:
+- Mock external system boundaries (LLM, ES, MinIO, MySQL)
+- Use real TaskContext, TaskHandler, and service instances
+- Verify RecordingContext for data flow assertions
+"""
+# =============================================================================
+# TensorFlow/UMAP Import Workaround
+# =============================================================================
+# Mock umap.parametric_umap before any other imports to prevent TensorFlow
+# dependency errors during test collection. This allows tests to run without
+# requiring TensorFlow to be installed.
+import sys
+from unittest.mock import MagicMock
+
+# Create a mock module for parametric_umap to satisfy umap's import check
+_mock_parametric_umap = MagicMock()
+sys.modules.setdefault("umap.parametric_umap", _mock_parametric_umap)
+sys.modules.setdefault("umap", MagicMock())
+
+import asyncio
+import uuid
+from typing import Any, Dict, List
+from unittest.mock import MagicMock, AsyncMock, patch
+
+import numpy as np
+import pytest
+
+from rag.svr.task_executor_refactor.task_context import TaskContext, TaskLimiters, TaskCallbacks
+from rag.svr.task_executor_refactor.recording_context import (
+    RecordingContext,
+    set_recording_context,
+)
+
+
+# =============================================================================
+# Async Limiter Fixtures
+# =============================================================================
+
+class AsyncMockLimiter:
+    """Mock asyncio semaphore that does not actually limit."""
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *args):
+        pass
+
+
+@pytest.fixture
+def mock_limiter():
+    """Provide a no-op async limiter."""
+    return asyncio.Semaphore(5)
+
+
+# =============================================================================
+# Task Dictionary Fixtures
+# =============================================================================
+
+@pytest.fixture
+def standard_task_dict() -> Dict[str, Any]:
+    """Provide a minimal but complete task dict for standard chunking."""
+    return {
+        "id": f"task_{uuid.uuid4().hex[:8]}",
+        "tenant_id": "tenant_test",
+        "kb_id": "kb_test",
+        "doc_id": "doc_test",
+        "name": "test_document.pdf",
+        "location": "/path/to/test_document.pdf",
+        "size": 1024,
+        "parser_id": "naive",
+        "parser_config": {
+            "auto_keywords": 0,
+            "auto_questions": 0,
+            "enable_metadata": False,
+        },
+        "kb_parser_config": {},
+        "language": "en",
+        "llm_id": "llm_test",
+        "embd_id": "embd_test",
+        "from_page": 0,
+        "to_page": -1,
+        "task_type": "standard",
+        "pagerank": 0,
+    }
+
+
+@pytest.fixture
+def dataflow_task_dict() -> Dict[str, Any]:
+    """Provide a task dict for dataflow tasks."""
+    task = standard_task_dict()
+    task["task_type"] = "dataflow"
+    task["dataflow_id"] = "dataflow_test"
+    return task
+
+
+@pytest.fixture
+def raptor_task_dict() -> Dict[str, Any]:
+    """Provide a task dict for RAPTOR tasks."""
+    task = standard_task_dict()
+    task["task_type"] = "raptor"
+    task["doc_ids"] = ["doc_1", "doc_2"]
+    return task
+
+
+@pytest.fixture
+def graphrag_task_dict() -> Dict[str, Any]:
+    """Provide a task dict for GraphRAG tasks."""
+    task = standard_task_dict()
+    task["task_type"] = "graphrag"
+    task["doc_ids"] = ["doc_1"]
+    return task
+
+
+@pytest.fixture
+def memory_task_dict() -> Dict[str, Any]:
+    """Provide a task dict for memory tasks."""
+    return {
+        "id": f"task_{uuid.uuid4().hex[:8]}",
+        "task_type": "memory",
+        "memory_id": "mem_test",
+        "source_id": "src_test",
+        "message_dict": {"role": "user", "content": "test"},
+    }
+
+
+# =============================================================================
+# TaskContext Fixtures
+# =============================================================================
+
+@pytest.fixture
+def task_context(standard_task_dict, mock_limiter, recording_context):
+    """Provide a real TaskContext instance with mocked limiters."""
+    ctx = TaskContext(
+        task=standard_task_dict,
+        limiters=TaskLimiters(
+            chat=mock_limiter,
+            minio=mock_limiter,
+            chunk=mock_limiter,
+            embed=mock_limiter,
+            kg=mock_limiter,
+        ),
+        callbacks=TaskCallbacks(
+            progress=MagicMock(),
+            has_canceled=MagicMock(return_value=False),
+        ),
+        recording_context=recording_context,
+    )
+    return ctx
+
+
+@pytest.fixture
+def canceled_task_context(standard_task_dict, mock_limiter, recording_context):
+    """Provide a TaskContext where the task is already canceled."""
+    ctx = TaskContext(
+        task=standard_task_dict,
+        limiters=TaskLimiters(
+            chat=mock_limiter,
+            minio=mock_limiter,
+            chunk=mock_limiter,
+            embed=mock_limiter,
+            kg=mock_limiter,
+        ),
+        callbacks=TaskCallbacks(
+            progress=MagicMock(),
+            has_canceled=MagicMock(return_value=True),
+        ),
+        recording_context=recording_context,
+    )
+    return ctx
+
+
+# =============================================================================
+# RecordingContext Fixtures
+# =============================================================================
+
+@pytest.fixture(autouse=True)
+def recording_context():
+    """Provide a fresh RecordingContext for each test.
+
+    This fixture is autouse=True to ensure every test has a clean
+    recording context for assertions.
+    """
+    ctx = RecordingContext()
+    set_recording_context(ctx)
+    yield ctx
+    # Cleanup: reset the global context after test
+    set_recording_context(RecordingContext())
+
+
+@pytest.fixture(autouse=True)
+def cleanup_resources(request):
+    """Global resource cleanup fixture.
+
+    Runs after each test to clean up:
+    - Unclosed event loops
+    - Unclosed sockets (via garbage collection)
+    - Unawaited coroutines
+    - MagicMock objects that may hold unclosed resources
+
+    This prevents ResourceWarning and RuntimeWarning from failing
+    tests when filterwarnings is set to "error".
+
+    Optimization: Uses minimal gc cycles and generation-2 collection
+    for faster teardown.
+    """
+    yield
+    import warnings
+
+    # Suppress warnings during cleanup to avoid recursive warning issues
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+
+        # Close any unclosed event loops
+        try:
+            policy = asyncio.get_event_loop_policy()
+            loop = policy.get_event_loop()
+            if not loop.is_closed():
+                loop.close()
+        except RuntimeError:
+            # No event loop exists, which is fine
+            pass
+
+
+# =============================================================================
+# External System Mocks (Boundary Mocks)
+# =============================================================================
+
+class MockEmbeddingModel:
+    """Mock embedding model that returns deterministic vectors."""
+
+    def __init__(self, vector_size: int = 128):
+        self.vector_size = vector_size
+        self.max_length = 512
+        self.llm_name = "mock_embedding"
+
+    def encode(self, texts: List[str]):
+        """Return random vectors for the given texts."""
+        vectors = np.random.rand(len(texts), self.vector_size).astype(np.float32)
+        token_count = sum(len(t.split()) for t in texts)
+        return vectors, token_count
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+
+class MockChatModel:
+    """Mock chat model that returns canned responses."""
+
+    def __init__(self):
+        self.llm_name = "mock_chat"
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+
+@pytest.fixture
+def mock_embedding_model():
+    """Provide a mock embedding model."""
+    return MockEmbeddingModel(vector_size=128)
+
+
+@pytest.fixture
+def mock_chat_model():
+    """Provide a mock chat model."""
+    return MockChatModel()
+
+
+# =============================================================================
+# Patching Helpers
+# =============================================================================
+
+def create_patch_embedding_model(vectors=None, vector_size=128):
+    """Create a patcher for the embedding model binding.
+
+    This patches the entire _bind_embedding_model flow to return a mock model.
+    """
+    if vectors is None:
+        vectors = np.random.rand(1, vector_size).astype(np.float32)
+
+    mock_model = MagicMock()
+    mock_model.encode.return_value = (vectors, 10)
+    mock_model.max_length = 512
+    mock_model.llm_name = "mock_embedding"
+    mock_model.__enter__ = MagicMock(return_value=mock_model)
+    mock_model.__exit__ = MagicMock(return_value=False)
+
+    return patch(
+        "rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name",
+        return_value=MagicMock(),
+    ), patch(
+        "rag.svr.task_executor_refactor.task_handler.LLMBundle",
+        return_value=mock_model,
+    ), patch(
+        "rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type",
+        return_value=MagicMock(),
+    )
+
+
+def create_patch_docstore_insert():
+    """Create a patcher for docStoreConn.insert that always succeeds."""
+    return patch(
+        "common.settings.docStoreConn",
+        new_callable=MagicMock,
+    )
+
+
+def create_patch_storage_binary(binary_data=b"fake pdf content"):
+    """Create a patcher for storage retrieval."""
+    mock_async = AsyncMock(return_value=binary_data)
+    return patch(
+        "rag.svr.task_executor_refactor.task_handler.File2DocumentService.get_storage_address",
+        return_value=("bucket_test", "name_test"),
+    ), patch(
+        "rag.svr.task_executor_refactor.task_handler.thread_pool_exec",
+        new_callable=MagicMock,
+        return_value=mock_async,
+    )
+
+
+def create_patch_parser_chunking(chunks=None):
+    """Create a patcher for the parser chunking to return predefined chunks.
+
+    Args:
+        chunks: List of chunk dicts to return from the parser.
+                If None, returns a default single chunk.
+    """
+    if chunks is None:
+        chunks = [{
+            "content_with_weight": "This is a test chunk content.",
+            "page_num_int": [0],
+            "top_int": [0],
+            "position_int": [0, 0, 0, 0],
+        }]
+
+    mock_async = AsyncMock(return_value=chunks)
+    return patch(
+        "rag.svr.task_executor_refactor.chunk_service.thread_pool_exec",
+        new_callable=MagicMock,
+        return_value=mock_async,
+    )
+
+
+# =============================================================================
+# Shared Helper Functions for Integration Tests
+# =============================================================================
+
+
+def create_mock_embedding_model(vector_size: int = 128):
+    """Create a mock embedding model that returns deterministic vectors matching input size."""
+    mock_model = MagicMock()
+
+    def mock_encode(texts):
+        n = len(texts) if isinstance(texts, list) else 1
+        return (
+            np.random.rand(n, vector_size).astype(np.float32),
+            10 * n,
+        )
+
+    mock_model.encode = mock_encode
+    mock_model.max_length = 512
+    mock_model.llm_name = "mock_embedding"
+    mock_model.__enter__ = MagicMock(return_value=mock_model)
+    mock_model.__exit__ = MagicMock(return_value=False)
+    return mock_model
+
+
+def create_mock_chat_model():
+    """Create a mock chat model."""
+    mock_model = MagicMock()
+    mock_model.llm_name = "mock_chat"
+    mock_model.__enter__ = MagicMock(return_value=mock_model)
+    mock_model.__exit__ = MagicMock(return_value=False)
+    return mock_model
+
+
+def create_mock_settings():
+    """Create a mock settings object with STORAGE_IMPL and docStoreConn."""
+    mock_settings = MagicMock()
+    mock_settings.STORAGE_IMPL = MagicMock()
+    mock_settings.STORAGE_IMPL.get = MagicMock(return_value=b"fake binary content")
+    mock_settings.docStoreConn = MagicMock()
+    mock_settings.docStoreConn.create_idx = MagicMock(return_value=None)
+    mock_settings.docStoreConn.insert = MagicMock(return_value=None)
+    mock_settings.docStoreConn.delete = MagicMock(return_value=None)
+    mock_settings.docStoreConn.index_exist = MagicMock(return_value=True)
+    mock_settings.docStoreConn.search = MagicMock(return_value={"hits": []})
+    mock_settings.DOC_MAXIMUM_SIZE = 100 * 1024 * 1024  # 100MB
+    mock_settings.DOC_BULK_SIZE = 100
+    mock_settings.retriever = MagicMock()
+    return mock_settings
+
+
+def create_default_chunks(count: int = 2) -> List[Dict[str, Any]]:
+    """Create default chunk dictionaries for testing."""
+    chunks = []
+    for i in range(count):
+        chunks.append({
+            "id": f"chunk_{i}_{uuid.uuid4().hex[:6]}",
+            "content_with_weight": f"This is test chunk content number {i}.",
+            "page_num_int": [i],
+            "top_int": [i * 100],
+            "position_int": [i, 0, i + 1, 0],
+            "doc_id": "doc_test",
+            "kb_id": "kb_test",
+            "docnm_kwd": "test_document.pdf",
+        })
+    return chunks
+
+
+def create_mock_chunk_service(chunks=None):
+    """Create a mock ChunkService instance."""
+    if chunks is None:
+        chunks = create_default_chunks(count=3)
+    mock_service = MagicMock()
+    mock_service.build_chunks = AsyncMock(return_value=chunks)
+    mock_service.insert_chunks = AsyncMock(return_value=True)
+    return mock_service
+
+
+@pytest.fixture
+def mock_embedding_model_factory():
+    """Provide a factory for mock embedding models."""
+    return create_mock_embedding_model
+
+
+@pytest.fixture
+def mock_chat_model_factory():
+    """Provide a factory for mock chat models."""
+    return create_mock_chat_model
+
+
+@pytest.fixture
+def mock_settings_factory():
+    """Provide a factory for mock settings."""
+    return create_mock_settings
+
+
+@pytest.fixture
+def mock_chunk_service_factory():
+    """Provide a factory for mock chunk services."""
+    return create_mock_chunk_service
+
+
+# =============================================================================
+# RaptorService Fixtures
+# =============================================================================
+
+def create_mock_raptor_context():
+    """Create a mock TaskContext suitable for RaptorService tests."""
+    ctx = MagicMock()
+    ctx.tenant_id = "tenant_1"
+    ctx.kb_id = "kb_1"
+    ctx.write_interceptor = None
+    ctx.progress_cb = MagicMock()
+    ctx.raw_task = {"type": ""}
+    ctx.parser_id = "naive"
+    ctx.parser_config = {}
+    ctx.name = "test.pdf"
+    ctx.pagerank = 0
+    ctx.id = "task_1"
+    return ctx
+
+
+@pytest.fixture
+def mock_raptor_context():
+    """Provide a mock TaskContext for RaptorService tests."""
+    return create_mock_raptor_context()

test/unit_test/rag/svr/task_executor_refactor/test_chunk_builder.py

@@ -0,0 +1,219 @@
+#
+#  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.
+
+"""
+Unit tests for ChunkBuilder module.
+"""
+
+import pytest
+from unittest.mock import MagicMock, patch, AsyncMock
+from rag.svr.task_executor_refactor.chunk_builder import (
+    get_parser,
+    run_chunking,
+    extract_outline,
+)
+
+
+class TestGetParser:
+    """Tests for get_parser function."""
+
+    @pytest.mark.parametrize("parser_id", [
+        "naive", "general", "table", "paper", "book",
+        "picture", "audio", "email", "presentation", "manual",
+        "laws", "qa", "resume", "one", "tag",
+    ])
+    def test_get_parser_returns_non_none(self, parser_id):
+        """Test that get_parser returns non-None for all parser types."""
+        parser = get_parser(parser_id)
+        assert parser is not None
+
+    def test_get_parser_kg(self):
+        """Test getting kg parser (maps to naive)."""
+        from common.constants import ParserType
+        parser = get_parser(ParserType.KG.value)
+        assert parser is not None
+
+
+class TestRunChunking:
+    """Tests for run_chunking function."""
+
+    def _create_mock_context(self):
+        """Helper to create a mock TaskContext."""
+        ctx = MagicMock()
+        ctx.name = "test.pdf"
+        ctx.location = "/path/to/test.pdf"
+        ctx.from_page = 0
+        ctx.to_page = -1
+        ctx.language = "en"
+        ctx.kb_id = "kb_1"
+        ctx.parser_config = {}
+        ctx.tenant_id = "tenant_1"
+        ctx.progress_cb = MagicMock()
+        ctx.raw_task = {}
+        ctx.chunk_limiter = MagicMock()
+        ctx.chunk_limiter.__aenter__ = AsyncMock()
+        ctx.chunk_limiter.__aexit__ = AsyncMock()
+        return ctx
+
+    @pytest.mark.asyncio
+    async def test_run_chunking_success(self):
+        """Test successful chunking."""
+        ctx = self._create_mock_context()
+
+        mock_chunker = MagicMock()
+        mock_chunker.chunk = MagicMock(return_value=[{"content_with_weight": "chunk1"}])
+
+        with patch("rag.svr.task_executor_refactor.chunk_builder.thread_pool_exec") as mock_thread:
+            # thread_pool_exec returns an awaitable that returns the list
+            mock_thread.return_value = [{"content_with_weight": "chunk1"}]
+
+            result = await run_chunking(mock_chunker, b"binary", ctx)
+
+            assert result is not None
+            assert len(result) == 1
+
+    @pytest.mark.asyncio
+    async def test_run_chunking_with_parser_config(self):
+        """Test chunking merges table parser config."""
+        ctx = self._create_mock_context()
+        ctx.raw_task = {"parser_config": {"chunk_token_num": 128}}
+
+        mock_chunker = MagicMock()
+        mock_chunker.chunk = MagicMock(return_value=[])
+
+        with patch("rag.svr.task_executor_refactor.chunk_builder.thread_pool_exec") as mock_thread:
+            mock_thread.return_value = []
+
+            with patch("rag.svr.task_executor_refactor.chunk_builder.merge_table_parser_config_from_kb") as mock_merge:
+                mock_merge.return_value = {"chunk_token_num": 128}
+
+                await run_chunking(mock_chunker, b"binary", ctx)
+
+                mock_merge.assert_called_once_with(ctx.raw_task)
+
+    @pytest.mark.asyncio
+    async def test_run_chunking_exception(self):
+        """Test chunking handles exception."""
+        ctx = self._create_mock_context()
+
+        mock_chunker = MagicMock()
+        mock_chunker.chunk = MagicMock(side_effect=Exception("Test error"))
+
+        with patch("rag.svr.task_executor_refactor.chunk_builder.thread_pool_exec") as mock_thread:
+            mock_thread.side_effect = Exception("Test error")
+
+            with pytest.raises(Exception):
+                await run_chunking(mock_chunker, b"binary", ctx)
+
+            # Verify progress_cb was called with error message
+            ctx.progress_cb.assert_called()
+
+
+class TestExtractOutline:
+    """Tests for extract_outline function."""
+
+    def _create_mock_context(self):
+        """Helper to create a mock TaskContext."""
+        ctx = MagicMock()
+        ctx.doc_id = "doc_1"
+        ctx.write_interceptor = None
+        ctx.progress_cb = MagicMock()
+        return ctx
+
+    @pytest.mark.asyncio
+    async def test_extract_outline_with_data(self):
+        """Test outline extraction when outline data is present."""
+        ctx = self._create_mock_context()
+
+        outline_data = [{"title": "Chapter 1", "page": 1}]
+        cks = [{"__outline__": outline_data}]
+
+        mock_rec_ctx = MagicMock()
+        ctx.recording_context = mock_rec_ctx
+
+        with patch("rag.svr.task_executor_refactor.chunk_builder.DocMetadataService") as mock_meta:
+            mock_meta.get_document_metadata.return_value = {}
+            mock_meta.update_document_metadata = MagicMock()
+
+            await extract_outline(cks, ctx)
+
+            mock_rec_ctx.record.assert_called_with("outline_data", outline_data)
+            # Outline should be popped from first chunk
+            assert "__outline__" not in cks[0]
+            mock_meta.update_document_metadata.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_extract_outline_without_data(self):
+        """Test outline extraction when no outline data."""
+        ctx = self._create_mock_context()
+
+        cks = [{"content_with_weight": "test"}]
+
+        mock_rec_ctx = MagicMock()
+        ctx.recording_context = mock_rec_ctx
+
+        await extract_outline(cks, ctx)
+
+        mock_rec_ctx.record.assert_called_with("outline_data", None)
+
+    @pytest.mark.asyncio
+    async def test_extract_outline_empty_chunks(self):
+        """Test outline extraction with empty chunks list."""
+        ctx = self._create_mock_context()
+
+        mock_rec_ctx = MagicMock()
+        ctx.recording_context = mock_rec_ctx
+
+        await extract_outline([], ctx)
+
+        mock_rec_ctx.record.assert_called_with("outline_data", None)
+
+    @pytest.mark.asyncio
+    async def test_extract_outline_with_write_interceptor(self):
+        """Test outline extraction with write interceptor."""
+        ctx = self._create_mock_context()
+        ctx.write_interceptor = MagicMock()
+
+        outline_data = [{"title": "Chapter 1", "page": 1}]
+        cks = [{"__outline__": outline_data}]
+
+        mock_rec_ctx = MagicMock()
+        ctx.recording_context = mock_rec_ctx
+
+        await extract_outline(cks, ctx)
+
+        ctx.write_interceptor.intercept.assert_called_once_with(
+            "DocMetadataService.update_document_metadata"
+        )
+
+    @pytest.mark.asyncio
+    async def test_extract_outline_persistence_exception(self):
+        """Test outline extraction handles persistence exception."""
+        ctx = self._create_mock_context()
+
+        outline_data = [{"title": "Chapter 1", "page": 1}]
+        cks = [{"__outline__": outline_data}]
+
+        mock_rec_ctx = MagicMock()
+        ctx.recording_context = mock_rec_ctx
+
+        with patch("rag.svr.task_executor_refactor.chunk_builder.DocMetadataService") as mock_meta:
+            mock_meta.get_document_metadata.return_value = {}
+            mock_meta.update_document_metadata.side_effect = Exception("DB error")
+
+            # Should not raise exception, just log warning
+            await extract_outline(cks, ctx)
+
+            mock_rec_ctx.record.assert_called_with("outline_data", outline_data)

test/unit_test/rag/svr/task_executor_refactor/test_chunk_post_processor.py

@@ -0,0 +1,460 @@
+#
+#  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.
+
+"""
+Unit tests for ChunkPostProcessor module.
+"""
+
+import pytest
+from unittest.mock import MagicMock, patch, AsyncMock
+from rag.svr.task_executor_refactor.chunk_post_processor import (
+    extract_keywords,
+    generate_questions,
+    generate_metadata,
+    apply_tags,
+    count_with_key,
+    build_metadata_config,
+)
+
+
+class TestExtractKeywords:
+    """Tests for extract_keywords function."""
+
+    def _create_mock_context(self):
+        """Helper to create a mock TaskContext."""
+        ctx = MagicMock()
+        ctx.tenant_id = "tenant_1"
+        ctx.llm_id = "llm_1"
+        ctx.language = "en"
+        ctx.parser_config = {"auto_keywords": 5}
+        ctx.id = "task_1"
+        ctx.progress_cb = MagicMock()
+        ctx.has_canceled_func = MagicMock(return_value=False)
+        ctx.chat_limiter = MagicMock()
+        ctx.chat_limiter.__aenter__ = AsyncMock()
+        ctx.chat_limiter.__aexit__ = AsyncMock()
+        return ctx
+
+    @pytest.mark.asyncio
+    async def test_extract_keywords_success(self):
+        """Test successful keyword extraction."""
+        ctx = self._create_mock_context()
+        docs = [
+            {"content_with_weight": "This is test content one"},
+            {"content_with_weight": "This is test content two"},
+        ]
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") as mock_config:
+            mock_config.return_value = MagicMock()
+
+            with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
+                mock_llm_instance = MagicMock()
+                mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
+                mock_llm.return_value.__exit__ = MagicMock(return_value=False)
+
+                with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache") as mock_cache:
+                    mock_cache.return_value = "keyword1, keyword2"
+
+                    with patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache"):
+                        with patch("rag.svr.task_executor_refactor.chunk_post_processor.rag_tokenizer") as mock_tokenizer:
+                            mock_tokenizer.tokenize.return_value = "keyword1 keyword2"
+
+                            await extract_keywords(docs, ctx)
+
+                            # Verify keywords were set
+                            assert "important_kwd" in docs[0]
+                            assert "important_tks" in docs[0]
+
+    @pytest.mark.asyncio
+    async def test_extract_keywords_canceled(self):
+        """Test keyword extraction when task is canceled."""
+        ctx = self._create_mock_context()
+        ctx.has_canceled_func = MagicMock(return_value=True)
+        docs = [{"content_with_weight": "This is test content"}]
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") as mock_config:
+            mock_config.return_value = MagicMock()
+
+            with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
+                mock_llm_instance = MagicMock()
+                mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
+                mock_llm.return_value.__exit__ = MagicMock(return_value=False)
+
+                with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache") as mock_cache:
+                    mock_cache.return_value = None  # No cache
+
+                    await extract_keywords(docs, ctx)
+
+                    # Should return early due to cancellation
+                    assert "important_kwd" not in docs[0]
+
+    @pytest.mark.asyncio
+    async def test_extract_keywords_empty_docs(self):
+        """Test keyword extraction with empty docs list."""
+        ctx = self._create_mock_context()
+        docs = []
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") as mock_config:
+            mock_config.return_value = MagicMock()
+
+            with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
+                mock_llm_instance = MagicMock()
+                mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
+                mock_llm.return_value.__exit__ = MagicMock(return_value=False)
+
+                await extract_keywords(docs, ctx)
+
+                # Should complete without error
+                ctx.progress_cb.assert_called()
+
+
+class TestGenerateQuestions:
+    """Tests for generate_questions function."""
+
+    def _create_mock_context(self):
+        """Helper to create a mock TaskContext."""
+        ctx = MagicMock()
+        ctx.tenant_id = "tenant_1"
+        ctx.llm_id = "llm_1"
+        ctx.language = "en"
+        ctx.parser_config = {"auto_questions": 3}
+        ctx.id = "task_1"
+        ctx.progress_cb = MagicMock()
+        ctx.has_canceled_func = MagicMock(return_value=False)
+        ctx.chat_limiter = MagicMock()
+        ctx.chat_limiter.__aenter__ = AsyncMock()
+        ctx.chat_limiter.__aexit__ = AsyncMock()
+        return ctx
+
+    @pytest.mark.asyncio
+    async def test_generate_questions_success(self):
+        """Test successful question generation."""
+        ctx = self._create_mock_context()
+        docs = [
+            {"content_with_weight": "This is test content one"},
+        ]
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") as mock_config:
+            mock_config.return_value = MagicMock()
+
+            with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
+                mock_llm_instance = MagicMock()
+                mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
+                mock_llm.return_value.__exit__ = MagicMock(return_value=False)
+
+                with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache") as mock_cache:
+                    mock_cache.return_value = "Question 1\nQuestion 2"
+
+                    with patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache"):
+                        with patch("rag.svr.task_executor_refactor.chunk_post_processor.rag_tokenizer") as mock_tokenizer:
+                            mock_tokenizer.tokenize.return_value = "Question 1 Question 2"
+
+                            await generate_questions(docs, ctx)
+
+                            # Verify questions were set
+                            assert "question_kwd" in docs[0]
+                            assert "question_tks" in docs[0]
+
+    @pytest.mark.asyncio
+    async def test_generate_questions_canceled(self):
+        """Test question generation when task is canceled."""
+        ctx = self._create_mock_context()
+        ctx.has_canceled_func = MagicMock(return_value=True)
+        docs = [{"content_with_weight": "This is test content"}]
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") as mock_config:
+            mock_config.return_value = MagicMock()
+
+            with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
+                mock_llm_instance = MagicMock()
+                mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
+                mock_llm.return_value.__exit__ = MagicMock(return_value=False)
+
+                with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache") as mock_cache:
+                    mock_cache.return_value = None  # No cache
+
+                    await generate_questions(docs, ctx)
+
+                    # Should return early due to cancellation
+                    assert "question_kwd" not in docs[0]
+
+
+class TestGenerateMetadata:
+    """Tests for generate_metadata function."""
+
+    def _create_mock_context(self):
+        """Helper to create a mock TaskContext."""
+        ctx = MagicMock()
+        ctx.tenant_id = "tenant_1"
+        ctx.llm_id = "llm_1"
+        ctx.language = "en"
+        ctx.parser_config = {
+            "enable_metadata": True,
+            "metadata": [{"name": "category", "type": "string"}],
+            "built_in_metadata": ["author", "date"],
+        }
+        ctx.doc_id = "doc_1"
+        ctx.id = "task_1"
+        ctx.progress_cb = MagicMock()
+        ctx.has_canceled_func = MagicMock(return_value=False)
+        ctx.write_interceptor = None
+        ctx.chat_limiter = MagicMock()
+        ctx.chat_limiter.__aenter__ = AsyncMock()
+        ctx.chat_limiter.__aexit__ = AsyncMock()
+        return ctx
+
+    @pytest.mark.asyncio
+    async def test_generate_metadata_success(self):
+        """Test successful metadata generation."""
+        ctx = self._create_mock_context()
+        docs = [
+            {"content_with_weight": "This is test content", "metadata_obj": {"category": "test"}},
+        ]
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") as mock_config:
+            mock_config.return_value = MagicMock()
+
+            with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
+                mock_llm_instance = MagicMock()
+                mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
+                mock_llm.return_value.__exit__ = MagicMock(return_value=False)
+
+                with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache") as mock_cache:
+                    mock_cache.return_value = {"category": "test"}
+
+                    with patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache"):
+                        with patch("rag.svr.task_executor_refactor.chunk_post_processor.update_metadata_to") as mock_update:
+                            mock_update.return_value = {"category": "test"}
+
+                            with patch("rag.svr.task_executor_refactor.chunk_post_processor.DocMetadataService") as mock_meta:
+                                mock_meta.get_document_metadata.return_value = {}
+                                mock_meta.update_document_metadata = MagicMock()
+
+                                await generate_metadata(docs, ctx)
+
+                                # Verify metadata_obj was processed
+                                mock_meta.update_document_metadata.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_generate_metadata_with_write_interceptor(self):
+        """Test metadata generation with write interceptor."""
+        ctx = self._create_mock_context()
+        ctx.write_interceptor = MagicMock()
+        docs = [
+            {"content_with_weight": "This is test content", "metadata_obj": {"category": "test"}},
+        ]
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") as mock_config:
+            mock_config.return_value = MagicMock()
+
+            with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
+                mock_llm_instance = MagicMock()
+                mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
+                mock_llm.return_value.__exit__ = MagicMock(return_value=False)
+
+                with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache") as mock_cache:
+                    mock_cache.return_value = {"category": "test"}
+
+                    with patch("rag.svr.task_executor_refactor.chunk_post_processor.update_metadata_to") as mock_update:
+                        mock_update.return_value = {"category": "test"}
+
+                        with patch("rag.svr.task_executor_refactor.chunk_post_processor.DocMetadataService") as mock_meta:
+                            mock_meta.get_document_metadata.return_value = {}
+                            mock_meta.update_document_metadata = MagicMock()
+
+                            await generate_metadata(docs, ctx)
+
+                            ctx.write_interceptor.intercept.assert_called_once_with(
+                                "DocMetadataService.update_document_metadata"
+                            )
+
+
+class TestApplyTags:
+    """Tests for apply_tags function."""
+
+    def _create_mock_context(self):
+        """Helper to create a mock TaskContext."""
+        ctx = MagicMock()
+        ctx.tenant_id = "tenant_1"
+        ctx.llm_id = "llm_1"
+        ctx.language = "en"
+        ctx.kb_parser_config = {"tag_kb_ids": ["kb_1"], "topn_tags": 3}
+        ctx.id = "task_1"
+        ctx.progress_cb = MagicMock()
+        ctx.has_canceled_func = MagicMock(return_value=False)
+        ctx.chat_limiter = MagicMock()
+        ctx.chat_limiter.__aenter__ = AsyncMock()
+        ctx.chat_limiter.__aexit__ = AsyncMock()
+        return ctx
+
+    @pytest.mark.asyncio
+    async def test_apply_tags_success(self):
+        """Test successful tag application."""
+        ctx = self._create_mock_context()
+        docs = [
+            {"content_with_weight": "This is test content"},
+        ]
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") as mock_config:
+            mock_config.return_value = MagicMock()
+
+            with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
+                mock_llm_instance = MagicMock()
+                mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
+                mock_llm.return_value.__exit__ = MagicMock(return_value=False)
+
+                with patch("rag.svr.task_executor_refactor.chunk_post_processor.settings") as mock_settings:
+                    mock_settings.retriever.all_tags_in_portion.return_value = {"tag1": 10, "tag2": 5}
+                    mock_settings.retriever.tag_content.return_value = True
+
+                    with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache") as mock_cache:
+                        mock_cache.return_value = '{"tag1": 1}'
+
+                        with patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache"):
+                            await apply_tags(docs, ctx)
+
+                            # Verify tags were applied
+                            assert len(docs) == 1
+
+    @pytest.mark.asyncio
+    async def test_apply_tags_canceled(self):
+        """Test tag application when task is canceled."""
+        ctx = self._create_mock_context()
+        ctx.has_canceled_func = MagicMock(return_value=True)
+        docs = [
+            {"content_with_weight": "This is test content"},
+        ]
+
+        with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_by_type_and_name") as mock_config:
+            mock_config.return_value = MagicMock()
+
+            with patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle") as mock_llm:
+                mock_llm_instance = MagicMock()
+                mock_llm.return_value.__enter__ = MagicMock(return_value=mock_llm_instance)
+                mock_llm.return_value.__exit__ = MagicMock(return_value=False)
+
+                with patch("rag.svr.task_executor_refactor.chunk_post_processor.settings") as mock_settings:
+                    mock_settings.retriever.all_tags_in_portion.return_value = {"tag1": 10}
+
+                    await apply_tags(docs, ctx)
+
+                    # Should return early due to cancellation
+
+
+class TestCountWithKey:
+    """Tests for count_with_key function."""
+
+    def test_count_with_key_all_have_key(self):
+        """Test counting when all docs have the key."""
+        docs = [{"tag": 1}, {"tag": 2}, {"tag": 3}]
+        result = count_with_key(docs, "tag")
+        assert result == 3
+
+    def test_count_with_key_some_have_key(self):
+        """Test counting when some docs have the key."""
+        docs = [{"tag": 1}, {"other": 2}, {"tag": 3}]
+        result = count_with_key(docs, "tag")
+        assert result == 2
+
+    def test_count_with_key_none_have_key(self):
+        """Test counting when no docs have the key."""
+        docs = [{"other": 1}, {"other": 2}]
+        result = count_with_key(docs, "tag")
+        assert result == 0
+
+    def test_count_with_key_empty_docs(self):
+        """Test counting with empty docs list."""
+        result = count_with_key([], "tag")
+        assert result == 0
+
+    def test_count_with_key_falsy_value(self):
+        """Test counting when key exists but has falsy value."""
+        docs = [{"tag": 0}, {"tag": ""}, {"tag": None}]
+        result = count_with_key(docs, "tag")
+        # Falsy values should not be counted (since d.get(key) returns falsy)
+        assert result == 0
+
+    def test_count_with_key_truthy_value(self):
+        """Test counting when key has truthy value."""
+        docs = [{"tag": 1}, {"tag": "value"}, {"tag": [1, 2]}]
+        result = count_with_key(docs, "tag")
+        assert result == 3
+
+
+class TestBuildMetadataConfig:
+    """Tests for build_metadata_config function."""
+
+    def test_dict_without_properties_returns_schema(self):
+        """When metadata is a dict without properties, return {type: object, properties: {}}."""
+        parser_config = {"metadata": {"type": "object"}, "built_in_metadata": []}
+        result = build_metadata_config(parser_config)
+        assert result == {"type": "object", "properties": {}}
+
+    def test_dict_with_properties_and_built_in(self):
+        """When metadata is a dict with properties AND built_in_metadata, merge them."""
+        parser_config = {
+            "metadata": {"type": "object", "properties": {"a": {"type": "string"}}},
+            "built_in_metadata": [{"key": "author", "description": "Author name", "enum": ["alice", "bob"]}],
+        }
+        result = build_metadata_config(parser_config)
+        assert result["type"] == "object"
+        assert "a" in result["properties"]
+        assert "author" in result["properties"]
+
+    def test_dict_with_properties_no_built_in(self):
+        """When metadata is a dict with properties and no built_in, return as-is."""
+        parser_config = {
+            "metadata": {"type": "object", "properties": {"a": {"type": "string"}}},
+            "built_in_metadata": [],
+        }
+        result = build_metadata_config(parser_config)
+        assert result == {"type": "object", "properties": {"a": {"type": "string"}}}
+
+    def test_list_with_built_in(self):
+        """When metadata is a list and built_in_metadata is present, concatenate."""
+        parser_config = {
+            "metadata": [{"key": "category"}],
+            "built_in_metadata": [{"key": "author"}],
+        }
+        result = build_metadata_config(parser_config)
+        assert result == [{"key": "category"}, {"key": "author"}]
+
+    def test_list_without_built_in(self):
+        """When metadata is a list and built_in_metadata is empty, return metadata as-is."""
+        parser_config = {"metadata": [{"key": "category"}], "built_in_metadata": []}
+        result = build_metadata_config(parser_config)
+        assert result == [{"key": "category"}]
+
+    def test_other_type_with_built_in(self):
+        """When metadata is not dict or list (empty list), return built_in_metadata only."""
+        parser_config = {"metadata": [], "built_in_metadata": [{"key": "author"}]}
+        result = build_metadata_config(parser_config)
+        assert result == [{"key": "author"}]
+
+    def test_idempotent_same_input(self):
+        """Same input produces structurally equal results."""
+        parser_config = {
+            "metadata": [{"key": "category"}],
+            "built_in_metadata": [{"key": "author"}],
+        }
+        result1 = build_metadata_config(parser_config)
+        result2 = build_metadata_config(parser_config)
+        assert result1 == result2
+
+    def test_missing_metadata_key(self):
+        """When parser_config has no 'metadata' key, built_in_metadata alone is returned."""
+        parser_config = {"built_in_metadata": []}
+        result = build_metadata_config(parser_config)
+        assert result == []

test/unit_test/rag/svr/task_executor_refactor/test_chunk_service.py

@@ -0,0 +1,453 @@
+#
+#  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.
+
+"""
+Unit tests for ChunkService module.
+
+Note: After refactoring, some functionality has been moved to:
+- chunk_builder.py: Parser factory, run_chunking, extract_outline
+- chunk_post_processor.py: Keyword extraction, question generation, metadata, tagging
+
+This test file now focuses on ChunkService-specific functionality:
+- build_chunks orchestration
+- _prepare_docs_and_upload
+- insert_chunks and related methods
+"""
+
+import pytest
+from unittest.mock import MagicMock, patch, AsyncMock
+from rag.svr.task_executor_refactor.chunk_service import ChunkService
+
+
+class TestChunkServiceInit:
+    """Tests for ChunkService initialization."""
+
+    def test_init_stores_task_context(self):
+        """Test that task context is stored."""
+        ctx = MagicMock()
+        service = ChunkService(ctx=ctx)
+        assert service._task_context is ctx
+
+
+class TestChunkServiceBuildChunks:
+    """Tests for build_chunks method."""
+
+    def _create_mock_context(self, parser_id="naive", size=1000, parser_config=None, kb_parser_config=None):
+        """Helper to create a mock TaskContext."""
+        ctx = MagicMock()
+        ctx.parser_id = parser_id
+        ctx.name = "test.pdf"
+        ctx.size = size
+        ctx.from_page = 0
+        ctx.to_page = -1
+        ctx.parser_config = parser_config or {}
+        ctx.kb_parser_config = kb_parser_config or {}
+        ctx.language = "en"
+        ctx.id = "task_1"
+        ctx.tenant_id = "tenant_1"
+        ctx.kb_id = "kb_1"
+        ctx.doc_id = "doc_1"
+        ctx.progress_cb = MagicMock()
+        ctx.has_canceled_func = MagicMock(return_value=False)
+        ctx.write_interceptor = None
+        ctx.raw_task = {}
+        ctx.llm_id = "llm_1"
+        ctx.pagerank = 0
+        ctx.location = "/path/to/test.pdf"
+        ctx.chunk_limiter = MagicMock()
+        ctx.chunk_limiter.__aenter__ = AsyncMock()
+        ctx.chunk_limiter.__aexit__ = AsyncMock()
+        ctx.chat_limiter = MagicMock()
+        ctx.chat_limiter.__aenter__ = AsyncMock()
+        ctx.chat_limiter.__aexit__ = AsyncMock()
+        return ctx
+
+    @pytest.mark.asyncio
+    async def test_build_chunks_file_size_exceeded(self):
+        """Test build_chunks returns empty list when file size exceeds limit."""
+        ctx = self._create_mock_context(size=1000000000)  # Very large size
+
+        service = ChunkService(ctx=ctx)
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.DOC_MAXIMUM_SIZE = 1000  # Small limit
+
+            mock_rec_ctx = MagicMock()
+            ctx.recording_context = mock_rec_ctx
+
+            result = await service.build_chunks(b"test binary")
+
+            assert result == []
+            mock_rec_ctx.record.assert_any_call("file_size_exceeded", True)
+
+    @pytest.mark.asyncio
+    async def test_build_chunks_file_size_ok(self):
+        """Test build_chunks proceeds when file size is within limit."""
+        ctx = self._create_mock_context(size=1000)
+
+        service = ChunkService(ctx=ctx)
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.DOC_MAXIMUM_SIZE = 10000000  # Large limit
+
+            mock_rec_ctx = MagicMock()
+            ctx.recording_context = mock_rec_ctx
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.get_parser") as mock_get_parser:
+                mock_parser = MagicMock()
+                mock_get_parser.return_value = mock_parser
+
+                with patch("rag.svr.task_executor_refactor.chunk_service.run_chunking", new_callable=AsyncMock) as mock_run_chunking:
+                    mock_run_chunking.return_value = [{"content_with_weight": "test"}]
+
+                    with patch("rag.svr.task_executor_refactor.chunk_service.extract_outline", new_callable=AsyncMock):
+                        with patch.object(service, '_prepare_docs_and_upload', new_callable=AsyncMock) as mock_prepare:
+                            mock_prepare.return_value = [{"id": "chunk_1", "content_with_weight": "test"}]
+
+                            await service.build_chunks(b"test binary")
+
+                            mock_rec_ctx.record.assert_any_call("file_size_exceeded", False)
+                            mock_rec_ctx.record.assert_any_call("parser_id", "naive")
+                            mock_get_parser.assert_called_once_with("naive")
+
+    @pytest.mark.asyncio
+    async def test_build_chunks_with_auto_keywords(self):
+        """Test build_chunks triggers keyword extraction when configured."""
+        ctx = self._create_mock_context(parser_config={"auto_keywords": 5})
+
+        service = ChunkService(ctx=ctx)
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.DOC_MAXIMUM_SIZE = 10000000
+
+            mock_rec_ctx = MagicMock()
+            ctx.recording_context = mock_rec_ctx
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.get_parser") as mock_get_parser:
+                mock_get_parser.return_value = MagicMock()
+
+                with patch("rag.svr.task_executor_refactor.chunk_service.run_chunking", new_callable=AsyncMock) as mock_run_chunking:
+                    mock_run_chunking.return_value = []
+
+                    with patch("rag.svr.task_executor_refactor.chunk_service.extract_outline", new_callable=AsyncMock):
+                        with patch.object(service, '_prepare_docs_and_upload', new_callable=AsyncMock) as mock_prepare:
+                            mock_prepare.return_value = [{"id": "chunk_1", "content_with_weight": "test"}]
+
+                            with patch("rag.svr.task_executor_refactor.chunk_service.extract_keywords", new_callable=AsyncMock) as mock_extract:
+                                await service.build_chunks(b"test binary")
+                                mock_extract.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_build_chunks_with_auto_questions(self):
+        """Test build_chunks triggers question generation when configured."""
+        ctx = self._create_mock_context(parser_config={"auto_questions": 3})
+
+        service = ChunkService(ctx=ctx)
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.DOC_MAXIMUM_SIZE = 10000000
+
+            mock_rec_ctx = MagicMock()
+            ctx.recording_context = mock_rec_ctx
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.get_parser") as mock_get_parser:
+                mock_get_parser.return_value = MagicMock()
+
+                with patch("rag.svr.task_executor_refactor.chunk_service.run_chunking", new_callable=AsyncMock) as mock_run_chunking:
+                    mock_run_chunking.return_value = []
+
+                    with patch("rag.svr.task_executor_refactor.chunk_service.extract_outline", new_callable=AsyncMock):
+                        with patch.object(service, '_prepare_docs_and_upload', new_callable=AsyncMock) as mock_prepare:
+                            mock_prepare.return_value = [{"id": "chunk_1", "content_with_weight": "test"}]
+
+                            with patch("rag.svr.task_executor_refactor.chunk_service.generate_questions", new_callable=AsyncMock) as mock_gen:
+                                await service.build_chunks(b"test binary")
+                                mock_gen.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_build_chunks_with_tag_kb_ids(self):
+        """Test build_chunks triggers tag application when tag_kb_ids configured."""
+        ctx = self._create_mock_context(kb_parser_config={"tag_kb_ids": ["kb_1"]})
+
+        service = ChunkService(ctx=ctx)
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.DOC_MAXIMUM_SIZE = 10000000
+
+            mock_rec_ctx = MagicMock()
+            ctx.recording_context = mock_rec_ctx
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.get_parser") as mock_get_parser:
+                mock_get_parser.return_value = MagicMock()
+
+                with patch("rag.svr.task_executor_refactor.chunk_service.run_chunking", new_callable=AsyncMock) as mock_run_chunking:
+                    mock_run_chunking.return_value = []
+
+                    with patch("rag.svr.task_executor_refactor.chunk_service.extract_outline", new_callable=AsyncMock):
+                        with patch.object(service, '_prepare_docs_and_upload', new_callable=AsyncMock) as mock_prepare:
+                            mock_prepare.return_value = [{"id": "chunk_1", "content_with_weight": "test"}]
+
+                            with patch("rag.svr.task_executor_refactor.chunk_service.apply_tags", new_callable=AsyncMock) as mock_apply:
+                                await service.build_chunks(b"test binary")
+                                mock_apply.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_build_chunks_with_metadata(self):
+        """Test build_chunks triggers metadata generation when configured."""
+        ctx = self._create_mock_context(
+            parser_config={
+                "enable_metadata": True,
+                "metadata": [{"name": "category", "type": "string"}]
+            }
+        )
+
+        service = ChunkService(ctx=ctx)
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.DOC_MAXIMUM_SIZE = 10000000
+
+            mock_rec_ctx = MagicMock()
+            ctx.recording_context = mock_rec_ctx
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.get_parser") as mock_get_parser:
+                mock_get_parser.return_value = MagicMock()
+
+                with patch("rag.svr.task_executor_refactor.chunk_service.run_chunking", new_callable=AsyncMock) as mock_run_chunking:
+                    mock_run_chunking.return_value = []
+
+                    with patch("rag.svr.task_executor_refactor.chunk_service.extract_outline", new_callable=AsyncMock):
+                        with patch.object(service, '_prepare_docs_and_upload', new_callable=AsyncMock) as mock_prepare:
+                            mock_prepare.return_value = [{"id": "chunk_1", "content_with_weight": "test"}]
+
+                            with patch("rag.svr.task_executor_refactor.chunk_service.generate_metadata", new_callable=AsyncMock) as mock_meta:
+                                await service.build_chunks(b"test binary")
+                                mock_meta.assert_called_once()
+
+
+class TestChunkServicePrepareDocsAndUpload:
+    """Tests for _prepare_docs_and_upload method."""
+
+    def _create_mock_context(self):
+        """Helper to create a mock TaskContext."""
+        ctx = MagicMock()
+        ctx.doc_id = "doc_1"
+        ctx.kb_id = "kb_1"
+        ctx.tenant_id = "tenant_1"
+        ctx.name = "test.pdf"
+        ctx.location = "/path/to/test.pdf"
+        ctx.pagerank = 0
+        ctx.progress_cb = MagicMock()
+        return ctx
+
+    @pytest.mark.asyncio
+    async def test_prepare_docs_and_upload_basic(self):
+        """Test basic document preparation."""
+        ctx = self._create_mock_context()
+        service = ChunkService(ctx=ctx)
+
+        cks = [{"content_with_weight": "test chunk"}]
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.STORAGE_IMPL = MagicMock()
+            mock_settings.STORAGE_IMPL.put = MagicMock()
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.image2id", new_callable=AsyncMock):
+
+                docs = await service._prepare_docs_and_upload(cks)
+
+                assert len(docs) == 1
+                assert docs[0]["doc_id"] == "doc_1"
+                assert docs[0]["kb_id"] == "kb_1"
+
+    @pytest.mark.asyncio
+    async def test_prepare_docs_and_upload_with_pagerank(self):
+        """Test document preparation with pagerank."""
+        ctx = self._create_mock_context()
+        ctx.pagerank = 5
+        service = ChunkService(ctx=ctx)
+
+        cks = [{"content_with_weight": "test chunk"}]
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.STORAGE_IMPL = MagicMock()
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.image2id", new_callable=AsyncMock):
+
+                docs = await service._prepare_docs_and_upload(cks)
+
+                assert docs[0].get("pagerank_fea") == 5
+
+
+class TestChunkServiceInsertChunks:
+    """Tests for insert_chunks method."""
+
+    def _create_mock_context(self):
+        """Helper to create a mock TaskContext."""
+        ctx = MagicMock()
+        ctx.id = "task_1"
+        ctx.tenant_id = "tenant_1"
+        ctx.kb_id = "kb_1"
+        ctx.doc_id = "doc_1"
+        ctx.parser_id = "naive"
+        ctx.progress_cb = MagicMock()
+        ctx.has_canceled_func = MagicMock(return_value=False)
+        ctx.write_interceptor = None
+        return ctx
+
+    @pytest.mark.asyncio
+    async def test_insert_chunks_success(self):
+        """Test successful chunk insertion."""
+        ctx = self._create_mock_context()
+        service = ChunkService(ctx=ctx)
+
+        chunks = [
+            {"id": "chunk_1", "content_with_weight": "test1"},
+            {"id": "chunk_2", "content_with_weight": "test2"},
+        ]
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.DOC_BULK_SIZE = 100
+            mock_settings.docStoreConn = MagicMock()
+            mock_settings.docStoreConn.insert = MagicMock(return_value=None)
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.search.index_name") as mock_index:
+                mock_index.return_value = "test_index"
+
+                with patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_thread:
+                    mock_thread.return_value = None
+
+                    with patch("rag.svr.task_executor_refactor.chunk_service.TaskService") as mock_task:
+                        mock_task.update_chunk_ids = MagicMock()
+
+                        result = await service.insert_chunks("task_1", "tenant_1", "kb_1", chunks)
+
+                        assert result is True
+
+    @pytest.mark.asyncio
+    async def test_insert_chunks_canceled(self):
+        """Test chunk insertion when task is canceled."""
+        ctx = self._create_mock_context()
+        ctx.has_canceled_func = MagicMock(return_value=True)
+        service = ChunkService(ctx=ctx)
+
+        chunks = [{"id": "chunk_1", "content_with_weight": "test1"}]
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.DOC_BULK_SIZE = 100
+            mock_settings.docStoreConn = MagicMock()
+            mock_settings.docStoreConn.insert = MagicMock(return_value=None)
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.search.index_name") as mock_index:
+                mock_index.return_value = "test_index"
+
+                with patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_thread:
+                    mock_thread.return_value = None
+
+                    result = await service.insert_chunks("task_1", "tenant_1", "kb_1", chunks)
+
+                    assert result is False
+                    ctx.progress_cb.assert_called_with(-1, msg="Task has been canceled.")
+
+    @pytest.mark.asyncio
+    async def test_insert_chunks_doc_store_error(self):
+        """Test chunk insertion when doc store returns error."""
+        ctx = self._create_mock_context()
+        service = ChunkService(ctx=ctx)
+
+        chunks = [{"id": "chunk_1", "content_with_weight": "test1"}]
+
+        with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings:
+            mock_settings.DOC_BULK_SIZE = 100
+            mock_settings.docStoreConn = MagicMock()
+            mock_settings.docStoreConn.insert = MagicMock(return_value="Error message")
+
+            with patch("rag.svr.task_executor_refactor.chunk_service.search.index_name") as mock_index:
+                mock_index.return_value = "test_index"
+
+                with patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_thread:
+                    mock_thread.return_value = "Error"
+
+                    with pytest.raises(Exception, match="Insert chunk error"):
+                        await service.insert_chunks("task_1", "tenant_1", "kb_1", chunks)
+
+
+class TestChunkServiceCreateMotherChunks:
+    """Tests for _create_mother_chunks class method."""
+
+    def test_create_mother_chunks_with_mom_field(self):
+        """Test creating mother chunks from mom field."""
+        chunks = [
+            {"id": "chunk_1", "mom": "Summary text 1", "content_with_weight": "test1"},
+        ]
+
+        mothers = ChunkService._create_mother_chunks(chunks)
+
+        assert len(mothers) == 1
+        assert mothers[0]["content_with_weight"] == "Summary text 1"
+        assert mothers[0]["available_int"] == 0
+
+    def test_create_mother_chunks_with_mom_with_weight_field(self):
+        """Test creating mother chunks from mom_with_weight field."""
+        chunks = [
+            {"id": "chunk_1", "mom_with_weight": "Summary text 2", "content_with_weight": "test1"},
+        ]
+
+        mothers = ChunkService._create_mother_chunks(chunks)
+
+        assert len(mothers) == 1
+        assert mothers[0]["content_with_weight"] == "Summary text 2"
+
+    def test_create_mother_chunks_no_mom_field(self):
+        """Test creating mother chunks when no mom field present."""
+        chunks = [
+            {"id": "chunk_1", "content_with_weight": "test1"},
+        ]
+
+        mothers = ChunkService._create_mother_chunks(chunks)
+
+        assert len(mothers) == 0
+
+    def test_create_mother_chunks_empty_mom(self):
+        """Test creating mother chunks with empty mom field."""
+        chunks = [
+            {"id": "chunk_1", "mom": "", "content_with_weight": "test1"},
+        ]
+
+        mothers = ChunkService._create_mother_chunks(chunks)
+
+        assert len(mothers) == 0
+
+    def test_create_mother_chunks_deduplicates_ids(self):
+        """Test that mother chunks deduplicate by ID."""
+        chunks = [
+            {"id": "chunk_1", "mom": "Same summary", "content_with_weight": "test1"},
+            {"id": "chunk_2", "mom": "Same summary", "content_with_weight": "test2"},
+        ]
+
+        mothers = ChunkService._create_mother_chunks(chunks)
+
+        assert len(mothers) == 1
+
+    def test_create_mother_chunks_filters_fields(self):
+        """Test that mother chunks only keep allowed fields."""
+        chunks = [
+            {"id": "chunk_1", "mom": "Summary", "extra_field": "should be removed", "content_with_weight": "test1"},
+        ]
+
+        mothers = ChunkService._create_mother_chunks(chunks)
+
+        assert "extra_field" not in mothers[0]
+        assert "id" in mothers[0]
+        assert "content_with_weight" in mothers[0]

test/unit_test/rag/svr/task_executor_refactor/test_comparator.py

@@ -0,0 +1,598 @@
+#
+#  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.
+
+"""
+Unit tests for Comparator module.
+"""
+
+from rag.svr.task_executor_refactor.report_generator import (
+    ComparisonResult,
+    ComparisonReport,
+)
+from rag.svr.task_executor_refactor.comparator import (
+    ContextComparator,
+)
+from rag.svr.task_executor_refactor.recording_context import RecordingContext
+
+
+class TestComparisonResult:
+    """Tests for ComparisonResult dataclass."""
+
+    def test_init_with_required_fields(self):
+        """Test initialization with required fields."""
+        result = ComparisonResult(key="test_key", match=True)
+        assert result.key == "test_key"
+        assert result.match is True
+        assert result.production_value is None
+        assert result.dry_run_value is None
+        assert result.diff_details is None
+
+    def test_init_with_all_fields(self):
+        """Test initialization with all fields."""
+        result = ComparisonResult(
+            key="test_key",
+            match=False,
+            production_value=100,
+            dry_run_value=200,
+            diff_details="Values differ"
+        )
+        assert result.key == "test_key"
+        assert result.match is False
+        assert result.production_value == 100
+        assert result.dry_run_value == 200
+        assert result.diff_details == "Values differ"
+
+    def test_to_dict_match(self):
+        """Test to_dict for matching result."""
+        result = ComparisonResult(key="key", match=True)
+        d = result.to_dict()
+        assert d == {"key": "key", "match": True, "diff_details": None}
+
+    def test_to_dict_mismatch(self):
+        """Test to_dict for mismatching result."""
+        result = ComparisonResult(
+            key="key",
+            match=False,
+            diff_details="Difference"
+        )
+        d = result.to_dict()
+        assert d == {"key": "key", "match": False, "diff_details": "Difference"}
+
+
+class TestComparisonReport:
+    """Tests for ComparisonReport dataclass."""
+
+    def test_init_with_required_fields(self):
+        """Test initialization with required fields."""
+        report = ComparisonReport(task_id="task_123")
+        assert report.task_id == "task_123"
+        assert report.total_keys == 0
+        assert report.matched_keys == 0
+        assert report.mismatched_keys == 0
+        assert report.missing_in_production == []
+        assert report.missing_in_dry_run == []
+        assert report.details == []
+
+    def test_summary_no_keys(self):
+        """Test summary when no keys to compare."""
+        report = ComparisonReport(task_id="task_123")
+        assert "No keys to compare" in report.summary()
+
+    def test_summary_with_keys(self):
+        """Test summary with keys."""
+        report = ComparisonReport(
+            task_id="task_123",
+            total_keys=10,
+            matched_keys=8,
+            mismatched_keys=2
+        )
+        summary = report.summary()
+        assert "8/10" in summary
+        assert "80.0%" in summary
+
+    def test_to_dict(self):
+        """Test to_dict serialization."""
+        report = ComparisonReport(
+            task_id="task_123",
+            total_keys=1,
+            matched_keys=1,
+            details=[ComparisonResult(key="k", match=True)]
+        )
+        d = report.to_dict()
+        assert d["task_id"] == "task_123"
+        assert d["total_keys"] == 1
+        assert len(d["details"]) == 1
+
+    def test_to_markdown(self):
+        """Test to_markdown serialization."""
+        report = ComparisonReport(
+            task_id="task_123",
+            total_keys=1,
+            matched_keys=1,
+            mismatched_keys=0,
+            missing_in_production=[],
+            missing_in_dry_run=[],
+            details=[ComparisonResult(key="k", match=True)]
+        )
+        md = report.to_markdown()
+        assert "# Comparison Report: task_123" in md
+        assert "## Summary" in md
+        assert "## Details" in md
+
+    def test_to_markdown_empty_details(self):
+        """Test to_markdown with no details."""
+        report = ComparisonReport(task_id="task_123")
+        md = report.to_markdown()
+        assert "No comparison details" in md
+
+
+class TestContextComparatorInit:
+    """Tests for ContextComparator initialization."""
+
+    def test_init_default_tolerance(self):
+        """Test initialization with default tolerance."""
+        comparator = ContextComparator()
+        assert comparator.float_tolerance == 1e-6
+
+    def test_init_custom_tolerance(self):
+        """Test initialization with custom tolerance."""
+        comparator = ContextComparator(float_tolerance=0.01)
+        assert comparator.float_tolerance == 0.01
+
+
+class TestContextComparatorCompareValue:
+    """Tests for ContextComparator.compare_value method."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        self.comparator = ContextComparator()
+
+    def test_compare_none_values(self):
+        """Test comparing None values."""
+        result = self.comparator.compare_value("key", None, None)
+        assert result.match is True
+
+    def test_compare_one_none(self):
+        """Test comparing when one value is None."""
+        result = self.comparator.compare_value("key", 1, None)
+        assert result.match is False
+        assert "None" in result.diff_details
+
+    def test_compare_equal_strings(self):
+        """Test comparing equal strings."""
+        result = self.comparator.compare_value("key", "hello", "hello")
+        assert result.match is True
+
+    def test_compare_different_strings(self):
+        """Test comparing different strings."""
+        result = self.comparator.compare_value("key", "hello", "world")
+        assert result.match is False
+
+    def test_compare_equal_booleans(self):
+        """Test comparing equal booleans."""
+        result = self.comparator.compare_value("key", True, True)
+        assert result.match is True
+
+    def test_compare_different_booleans(self):
+        """Test comparing different booleans."""
+        result = self.comparator.compare_value("key", True, False)
+        assert result.match is False
+
+    def test_compare_equal_integers(self):
+        """Test comparing equal integers."""
+        result = self.comparator.compare_value("key", 42, 42)
+        assert result.match is True
+
+    def test_compare_equal_floats_within_tolerance(self):
+        """Test comparing equal floats within tolerance."""
+        result = self.comparator.compare_value("key", 1.0000001, 1.0000002)
+        assert result.match is True
+
+    def test_compare_different_floats_exceeding_tolerance(self):
+        """Test comparing floats exceeding tolerance."""
+        result = self.comparator.compare_value("key", 1.0, 2.0)
+        assert result.match is False
+        assert "exceeds tolerance" in result.diff_details
+
+    def test_compare_equal_lists(self):
+        """Test comparing equal lists."""
+        result = self.comparator.compare_value("key", [1, 2, 3], [1, 2, 3])
+        assert result.match is True
+
+    def test_compare_different_length_lists(self):
+        """Test comparing lists with different lengths."""
+        result = self.comparator.compare_value("key", [1, 2], [1, 2, 3])
+        assert result.match is False
+        assert "Length differs" in result.diff_details
+
+    def test_compare_equal_dicts(self):
+        """Test comparing equal dicts."""
+        result = self.comparator.compare_value("key", {"a": 1}, {"a": 1})
+        assert result.match is True
+
+    def test_compare_different_dicts(self):
+        """Test comparing different dicts."""
+        result = self.comparator.compare_value("key", {"a": 1}, {"a": 2})
+        assert result.match is False
+
+    def test_compare_chunks_key_uses_chunk_comparison(self):
+        """Test that chunk keys use chunk comparison strategy."""
+        result = self.comparator.compare_value(
+            "raw_chunks",
+            [{"id": "1", "content_with_weight": "a"}],
+            [{"id": "1", "content_with_weight": "a"}]
+        )
+        assert result.match is True
+
+
+class TestContextComparatorCompareLists:
+    """Tests for _compare_lists method."""
+
+    def test_equal_lists(self):
+        """Test comparing equal lists."""
+        result = ContextComparator._compare_lists("key", [1, 2], [1, 2])
+        assert result.match is True
+
+    def test_different_length_lists(self):
+        """Test comparing lists with different lengths."""
+        result = ContextComparator._compare_lists("key", [1], [1, 2])
+        assert result.match is False
+
+    def test_different_elements(self):
+        """Test comparing lists with different elements."""
+        result = ContextComparator._compare_lists("key", [1, 2], [1, 3])
+        assert result.match is False
+
+
+class TestContextComparatorCompareDicts:
+    """Tests for _compare_dicts method."""
+
+    def test_equal_dicts(self):
+        """Test comparing equal dicts."""
+        result = ContextComparator._compare_dicts("key", {"a": 1}, {"a": 1})
+        assert result.match is True
+
+    def test_dicts_different_keys(self):
+        """Test comparing dicts with different keys."""
+        result = ContextComparator._compare_dicts("key", {"a": 1}, {"b": 1})
+        assert result.match is False
+        assert "Keys differ" in result.diff_details
+
+    def test_dicts_same_keys_different_values(self):
+        """Test comparing dicts with same keys but different values."""
+        result = ContextComparator._compare_dicts("key", {"a": 1}, {"a": 2})
+        assert result.match is False
+
+
+class TestContextComparatorCompareNumbers:
+    """Tests for _compare_numbers method."""
+
+    def test_equal_numbers(self):
+        """Test comparing equal numbers."""
+        comparator = ContextComparator()
+        result = comparator._compare_numbers("key", 1.0, 1.0)
+        assert result.match is True
+
+    def test_numbers_within_tolerance(self):
+        """Test comparing numbers within tolerance."""
+        comparator = ContextComparator(float_tolerance=0.1)
+        result = comparator._compare_numbers("key", 1.0, 1.05)
+        assert result.match is True
+
+    def test_numbers_exceeding_tolerance(self):
+        """Test comparing numbers exceeding tolerance."""
+        comparator = ContextComparator(float_tolerance=0.01)
+        result = comparator._compare_numbers("key", 1.0, 1.1)
+        assert result.match is False
+
+
+class TestContextComparatorCompareChunks:
+    """Tests for _compare_chunks method."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        self.comparator = ContextComparator()
+
+    def test_equal_chunks(self):
+        """Test comparing equal chunk lists."""
+        prod = [{"id": "1", "content_with_weight": "a"}]
+        dry = [{"id": "1", "content_with_weight": "a"}]
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is True
+
+    def test_different_count_chunks(self):
+        """Test comparing chunks with different counts."""
+        prod = [{"id": "1"}]
+        dry = [{"id": "1"}, {"id": "2"}]
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is False
+        assert "Chunk count differs" in result.diff_details
+
+    def test_different_ids_chunks(self):
+        """Test comparing chunks with different IDs."""
+        prod = [{"id": "1"}]
+        dry = [{"id": "2"}]
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is False
+        assert "Chunk IDs differ" in result.diff_details
+
+    def test_empty_chunks_lists(self):
+        """Test comparing empty chunk lists."""
+        result = self.comparator._compare_chunks("raw_chunks", [], [])
+        assert result.match is True
+
+    def test_all_chunks_compared_not_sampled(self):
+        """Test that ALL chunks are compared, not just samples.
+        
+        This test creates 10 chunks where only the middle one (index 5) differs.
+        With the old sampling strategy, this difference might be missed.
+        With full comparison, the difference should always be detected.
+        """
+        prod = [{"id": str(i), "content_with_weight": f"content_{i}"} for i in range(10)]
+        dry = [{"id": str(i), "content_with_weight": f"content_{i}"} for i in range(10)]
+        # Only modify chunk at index 5 (which might not be sampled in old strategy)
+        dry[5]["content_with_weight"] = "different_content"
+        
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is False
+        assert "Content differs" in result.diff_details
+
+    def test_all_chunks_detect_first_difference(self):
+        """Test that first chunk difference is detected."""
+        prod = [{"id": "1", "content_with_weight": "a"}, {"id": "2", "content_with_weight": "b"}]
+        dry = [{"id": "1", "content_with_weight": "different"}, {"id": "2", "content_with_weight": "b"}]
+        
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is False
+
+    def test_all_chunks_detect_last_difference(self):
+        """Test that last chunk difference is detected."""
+        prod = [{"id": "1", "content_with_weight": "a"}, {"id": "2", "content_with_weight": "b"}]
+        dry = [{"id": "1", "content_with_weight": "a"}, {"id": "2", "content_with_weight": "different"}]
+        
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is False
+
+    def test_all_chunks_large_list_all_match(self):
+        """Test that large list of chunks all match."""
+        prod = [{"id": str(i), "content_with_weight": f"content_{i}"} for i in range(100)]
+        dry = [{"id": str(i), "content_with_weight": f"content_{i}"} for i in range(100)]
+        
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is True
+
+    def test_all_chunks_large_list_one_mismatch(self):
+        """Test that a single mismatch in a large list is detected."""
+        prod = [{"id": str(i), "content_with_weight": f"content_{i}"} for i in range(100)]
+        dry = [{"id": str(i), "content_with_weight": f"content_{i}"} for i in range(100)]
+        # Modify only the last chunk
+        dry[99]["content_with_weight"] = "different"
+        
+        result = self.comparator._compare_chunks("raw_chunks", prod, dry)
+        assert result.match is False
+
+
+class TestContextComparatorExtractChunkIds:
+    """Tests for _extract_chunk_ids method."""
+
+    def test_extract_ids_from_valid_chunks(self):
+        """Test extracting IDs from valid chunks."""
+        chunks = [{"id": "1"}, {"id": "2"}, {"id": "3"}]
+        ids = ContextComparator._extract_chunk_ids(chunks)
+        assert ids == {"1", "2", "3"}
+
+    def test_extract_ids_from_empty_chunks(self):
+        """Test extracting IDs from empty list."""
+        ids = ContextComparator._extract_chunk_ids([])
+        assert ids == set()
+
+    def test_extract_ids_from_chunks_without_id(self):
+        """Test extracting IDs from chunks without id field."""
+        chunks = [{"content": "a"}, {"id": "1"}]
+        ids = ContextComparator._extract_chunk_ids(chunks)
+        assert ids == {"1"}
+
+
+class TestContextComparatorGetChunkId:
+    """Tests for _get_chunk_id method."""
+
+    def test_get_id_from_valid_chunk(self):
+        """Test getting ID from valid chunk."""
+        chunk = {"id": "123"}
+        assert ContextComparator._get_chunk_id(chunk) == "123"
+
+    def test_get_id_from_chunk_without_id(self):
+        """Test getting ID from chunk without id."""
+        chunk = {"content": "a"}
+        assert ContextComparator._get_chunk_id(chunk) == ""
+
+    def test_get_id_from_non_dict(self):
+        """Test getting ID from non-dict."""
+        assert ContextComparator._get_chunk_id("not a dict") == ""
+
+
+class TestContextComparatorCompare:
+    """Tests for compare method."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        self.comparator = ContextComparator()
+
+    def test_compare_empty_contexts(self):
+        """Test comparing empty contexts."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        report = self.comparator.compare("task_1", ctx1, ctx2)
+        assert report.total_keys == 0
+
+    def test_compare_matching_values(self):
+        """Test comparing contexts with matching values."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        ctx1.record("key", "value")
+        ctx2.record("key", "value")
+        report = self.comparator.compare("task_1", ctx1, ctx2)
+        assert report.matched_keys == 1
+        assert report.mismatched_keys == 0
+
+    def test_compare_mismatching_values(self):
+        """Test comparing contexts with mismatching values."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        ctx1.record("key1", "value1")
+        ctx2.record("key1", "value2")
+        report = self.comparator.compare("task_1", ctx1, ctx2)
+        assert report.mismatched_keys == 1
+
+    def test_compare_missing_key_in_one_context(self):
+        """Test comparing when key is missing in one context."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        ctx1.record("key1", "value1")
+        report = self.comparator.compare("task_1", ctx1, ctx2)
+        assert "key1" in report.missing_in_dry_run
+
+    def test_compare_with_specific_keys(self):
+        """Test comparing with specific keys list."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        ctx1.record("key1", "value1")
+        ctx1.record("key2", "value2")
+        ctx2.record("key1", "value1")
+        ctx2.record("key2", "value2")
+        report = self.comparator.compare("task_1", ctx1, ctx2, comparison_keys=["key1"])
+        assert report.total_keys == 1
+
+    def test_compare_filters_out_time_keys(self):
+        """Test that _time keys are filtered out."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        ctx1.record("operation_time", 1.0)
+        ctx2.record("operation_time", 1.0)
+        report = self.comparator.compare("task_1", ctx1, ctx2)
+        assert report.total_keys == 0
+
+
+class TestContextComparatorStripNonDeterministicFields:
+    """Tests for _strip_non_deterministic_fields method."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        self.comparator = ContextComparator()
+
+    def test_strip_seconds_from_dict_value(self):
+        """Test that 'seconds' key is removed from dict values."""
+        data = {
+            "graphrag_result": {"seconds": 45.48, "status": "done"},
+            "other_key": "value"
+        }
+        result = self.comparator._strip_non_deterministic_fields(data)
+        assert "seconds" not in result["graphrag_result"]
+        assert result["graphrag_result"] == {"status": "done"}
+        assert result["other_key"] == "value"
+
+    def test_strip_seconds_from_multiple_dict_values(self):
+        """Test that 'seconds' is removed from multiple dict values."""
+        data = {
+            "result1": {"seconds": 10.0, "count": 5},
+            "result2": {"seconds": 20.0, "name": "test"},
+            "simple_key": 123
+        }
+        result = self.comparator._strip_non_deterministic_fields(data)
+        assert result["result1"] == {"count": 5}
+        assert result["result2"] == {"name": "test"}
+        assert result["simple_key"] == 123
+
+    def test_strip_does_not_modify_original_dict(self):
+        """Test that the original dict is not modified in place."""
+        data = {
+            "result": {"seconds": 1.0, "value": "test"}
+        }
+        _ = data["result"].copy()
+        self.comparator._strip_non_deterministic_fields(data)
+        # The original nested dict should still have seconds since we only do shallow copy
+        assert "seconds" in data["result"]
+
+    def test_strip_with_empty_dict_values(self):
+        """Test handling of empty dict values."""
+        data = {
+            "empty_dict": {},
+            "normal_key": "value"
+        }
+        result = self.comparator._strip_non_deterministic_fields(data)
+        assert result["empty_dict"] == {}
+        assert result["normal_key"] == "value"
+
+    def test_strip_with_non_dict_values(self):
+        """Test that non-dict values are not affected."""
+        data = {
+            "string_val": "test",
+            "int_val": 42,
+            "list_val": [1, 2, 3],
+            "dict_val": {"seconds": 1.0, "name": "test"}
+        }
+        result = self.comparator._strip_non_deterministic_fields(data)
+        assert result["string_val"] == "test"
+        assert result["int_val"] == 42
+        assert result["list_val"] == [1, 2, 3]
+        assert result["dict_val"] == {"name": "test"}
+
+    def test_strip_seconds_from_graphrag_result(self):
+        """Test the specific case from the bug report: graphrag_result with seconds."""
+        prod_data = {
+            "graphrag_result": {
+                "seconds": 45.48,
+                "status": "success",
+                "entity_count": 100
+            }
+        }
+        dry_run_data = {
+            "graphrag_result": {
+                "seconds": 0.99,
+                "status": "success",
+                "entity_count": 100
+            }
+        }
+        prod_stripped = self.comparator._strip_non_deterministic_fields(prod_data)
+        dry_run_stripped = self.comparator._strip_non_deterministic_fields(dry_run_data)
+        
+        # After stripping, both should be equal (except for seconds)
+        assert prod_stripped["graphrag_result"] == {"status": "success", "entity_count": 100}
+        assert dry_run_stripped["graphrag_result"] == {"status": "success", "entity_count": 100}
+        assert prod_stripped["graphrag_result"] == dry_run_stripped["graphrag_result"]
+
+    def test_compare_with_seconds_in_dict_values(self):
+        """Test that compare correctly handles dict values with 'seconds' field."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        ctx1.record("graphrag_result", {"seconds": 45.48, "status": "success"})
+        ctx2.record("graphrag_result", {"seconds": 0.99, "status": "success"})
+        
+        report = self.comparator.compare("task_1", ctx1, ctx2)
+        # Should match because seconds is stripped
+        assert report.matched_keys == 1
+        assert report.mismatched_keys == 0
+
+    def test_compare_with_different_dict_values_excluding_seconds(self):
+        """Test that compare correctly detects differences in dict values (excluding seconds)."""
+        ctx1 = RecordingContext()
+        ctx2 = RecordingContext()
+        ctx1.record("graphrag_result", {"seconds": 45.48, "status": "success", "count": 100})
+        ctx2.record("graphrag_result", {"seconds": 0.99, "status": "failed", "count": 50})
+        
+        report = self.comparator.compare("task_1", ctx1, ctx2)
+        # Should mismatch because status and count differ
+        assert report.mismatched_keys == 1
+        assert report.matched_keys == 0

test/unit_test/rag/svr/task_executor_refactor/test_constants.py

@@ -0,0 +1,43 @@
+#
+#  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.
+
+"""
+Unit tests for constants module.
+"""
+
+import pytest
+from rag.svr.task_executor_refactor.constants import CANVAS_DEBUG_DOC_ID
+
+
+class TestConstants:
+    """Tests for constants module."""
+
+    def test_canvas_debug_doc_id_exists(self):
+        """Test that CANVAS_DEBUG_DOC_ID constant exists."""
+        assert CANVAS_DEBUG_DOC_ID is not None
+
+    @pytest.mark.parametrize("expected_type", [str])
+    def test_canvas_debug_doc_id_type(self, expected_type):
+        """Test that CANVAS_DEBUG_DOC_ID is a string."""
+        assert isinstance(CANVAS_DEBUG_DOC_ID, expected_type)
+
+    @pytest.mark.parametrize("expected_value", ["dataflow_x"])
+    def test_canvas_debug_doc_id_value(self, expected_value):
+        """Test that CANVAS_DEBUG_DOC_ID has expected value."""
+        assert CANVAS_DEBUG_DOC_ID == expected_value
+
+    def test_canvas_debug_doc_id_not_empty(self):
+        """Test that CANVAS_DEBUG_DOC_ID is not empty."""
+        assert len(CANVAS_DEBUG_DOC_ID) > 0

test/unit_test/rag/svr/task_executor_refactor/test_dataflow_service.py

@@ -0,0 +1,381 @@
+#
+#  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.
+
+"""
+Unit tests for DataflowService module.
+
+Tests validate behavior through the public run_dataflow() entry point.
+Private orchestration helpers (_process_chunks, _encode_batch, _normalize_chunks,
+_get_output_type, _embed_chunks, _load_dsl, etc.) are exercised implicitly; no test
+reaches directly into those internals.
+"""
+
+import pytest
+from unittest.mock import MagicMock, AsyncMock, patch
+
+from rag.svr.task_executor_refactor.dataflow_service import DataflowService
+
+
+class TestDataflowServiceRunDataflow:
+    """Tests for the public run_dataflow() method.
+
+    Internal helpers (_load_dsl, _normalize_chunks, _get_output_type, _process_chunks,
+    _embed_chunks, _encode_batch) are exercised through this single entry point so
+    the suite stays resilient when internal method boundaries change.
+    """
+
+    @pytest.mark.asyncio
+    @patch("rag.svr.task_executor_refactor.dataflow_service.UserCanvasService")
+    @patch("rag.svr.task_executor_refactor.dataflow_service.PipelineOperationLogService")
+    async def test_run_dataflow_dsl_not_found(self, mock_pipeline_log, mock_canvas, task_context):
+        """Test run_dataflow returns early when DSL is not found."""
+        task_context._task["task_type"] = "dataflow"
+        task_context._task["dataflow_id"] = "dataflow_test"
+        mock_canvas.get_by_id.return_value = (False, None)
+
+        service = DataflowService(ctx=task_context)
+        with pytest.raises(AssertionError, match="User pipeline not found"):
+            await service.run_dataflow()
+
+    @pytest.mark.asyncio
+    @patch("rag.svr.task_executor_refactor.dataflow_service.Pipeline")
+    @patch("rag.svr.task_executor_refactor.dataflow_service.UserCanvasService")
+    async def test_run_dataflow_empty_chunks(self, mock_canvas, mock_pipeline_class, task_context):
+        """Test run_dataflow handles empty pipeline output."""
+        task_context._task["task_type"] = "dataflow"
+        task_context._task["dataflow_id"] = "dataflow_test"
+        mock_canvas.get_by_id.return_value = (True, MagicMock(dsl='{"id": "test"}'))
+        mock_pipeline = MagicMock()
+        mock_pipeline.run = AsyncMock(return_value={})
+        mock_pipeline_class.return_value = mock_pipeline
+
+        with patch.object(DataflowService, '_record_pipeline_log'):
+            service = DataflowService(ctx=task_context)
+            await service.run_dataflow()
+
+    @pytest.mark.asyncio
+    @patch("rag.svr.task_executor_refactor.dataflow_service.Pipeline")
+    @patch("rag.svr.task_executor_refactor.dataflow_service.UserCanvasService")
+    async def test_run_dataflow_with_chunks_output(self, mock_canvas, mock_pipeline_class, task_context):
+        """Test run_dataflow processes 'chunks' output type end-to-end."""
+        task_context._task["task_type"] = "dataflow"
+        task_context._task["dataflow_id"] = "dataflow_test"
+        task_context._task["tenant_id"] = "tenant_test"
+        task_context._task["kb_id"] = "kb_test"
+        task_context._task["doc_id"] = "doc_test"
+        task_context._task["name"] = "test.pdf"
+        task_context._write_interceptor = None
+
+        mock_canvas.get_by_id.return_value = (True, MagicMock(dsl='{"id": "test"}'))
+        chunks = {
+            "chunks": [
+                {"text": "Hello world", "content_with_weight": "Hello world"},
+            ],
+            "embedding_token_consumption": 5,
+        }
+        mock_pipeline = MagicMock()
+        mock_pipeline.run = AsyncMock(return_value=chunks)
+        mock_pipeline_class.return_value = mock_pipeline
+
+        # Patch internal heavy dependencies so run_dataflow completes
+        with patch.object(DataflowService, '_embed_chunks', new_callable=AsyncMock, return_value=(chunks["chunks"], 5)):
+            with patch.object(DataflowService, '_insert_chunks', new_callable=AsyncMock, return_value=True):
+                with patch.object(DataflowService, '_update_document_metadata'):
+                    with patch.object(DataflowService, '_record_pipeline_log'):
+                        with patch("api.db.services.document_service.DocumentService.increment_chunk_num"):
+                            service = DataflowService(ctx=task_context)
+                            await service.run_dataflow()
+
+                            # Verify chunks were inserted
+                            DataflowService._insert_chunks.assert_called_once()
+
+    @pytest.mark.asyncio
+    @patch("rag.svr.task_executor_refactor.dataflow_service.Pipeline")
+    @patch("rag.svr.task_executor_refactor.dataflow_service.UserCanvasService")
+    async def test_run_dataflow_with_json_output(self, mock_canvas, mock_pipeline_class, task_context):
+        """Test run_dataflow processes 'json' output type."""
+        task_context._task["task_type"] = "dataflow"
+        task_context._task["dataflow_id"] = "dataflow_test"
+        task_context._task["tenant_id"] = "tenant_test"
+        task_context._task["kb_id"] = "kb_test"
+        task_context._task["doc_id"] = "doc_test"
+        task_context._task["name"] = "test.pdf"
+        task_context._write_interceptor = None
+
+        mock_canvas.get_by_id.return_value = (True, MagicMock(dsl='{"id": "test"}'))
+        chunks = {
+            "json": [
+                {"text": "JSON content"},
+            ],
+            "embedding_token_consumption": 2,
+        }
+        mock_pipeline = MagicMock()
+        mock_pipeline.run = AsyncMock(return_value=chunks)
+        mock_pipeline_class.return_value = mock_pipeline
+
+        with patch.object(DataflowService, '_embed_chunks', new_callable=AsyncMock, return_value=(chunks["json"], 2)):
+            with patch.object(DataflowService, '_insert_chunks', new_callable=AsyncMock, return_value=True):
+                with patch.object(DataflowService, '_update_document_metadata'):
+                    with patch.object(DataflowService, '_record_pipeline_log'):
+                        with patch("api.db.services.document_service.DocumentService.increment_chunk_num"):
+                            service = DataflowService(ctx=task_context)
+                            await service.run_dataflow()
+
+    @pytest.mark.asyncio
+    @patch("rag.svr.task_executor_refactor.dataflow_service.Pipeline")
+    @patch("rag.svr.task_executor_refactor.dataflow_service.UserCanvasService")
+    async def test_run_dataflow_embedding_failure(self, mock_canvas, mock_pipeline_class, task_context):
+        """Test run_dataflow handles embedding failure gracefully."""
+        task_context._task["task_type"] = "dataflow"
+        task_context._task["dataflow_id"] = "dataflow_test"
+        task_context._task["name"] = "test.pdf"
+        task_context._write_interceptor = None
+
+        mock_canvas.get_by_id.return_value = (True, MagicMock(dsl='{"id": "test"}'))
+        chunks = {
+            "chunks": [
+                {"text": "Hello"},
+            ],
+            "embedding_token_consumption": 1,
+        }
+        mock_pipeline = MagicMock()
+        mock_pipeline.run = AsyncMock(return_value=chunks)
+        mock_pipeline_class.return_value = mock_pipeline
+
+        with patch.object(DataflowService, '_embed_chunks', new_callable=AsyncMock, return_value=(None, 0)):
+            with patch.object(DataflowService, '_record_pipeline_log'):
+                service = DataflowService(ctx=task_context)
+                await service.run_dataflow()
+
+                # Should not insert chunks when embedding fails
+                service._record_pipeline_log.assert_called()
+
+    @pytest.mark.asyncio
+    @patch("rag.svr.task_executor_refactor.dataflow_service.Pipeline")
+    @patch("rag.svr.task_executor_refactor.dataflow_service.UserCanvasService")
+    async def test_run_dataflow_with_billing_hook_success(self, mock_canvas, mock_pipeline_class, task_context):
+        """Test run_dataflow calls billing hook on success."""
+        task_context._task["task_type"] = "dataflow"
+        task_context._task["dataflow_id"] = "dataflow_test"
+        task_context._task["tenant_id"] = "tenant_test"
+        task_context._task["kb_id"] = "kb_test"
+        task_context._task["doc_id"] = "doc_test"
+        task_context._task["name"] = "test.pdf"
+        task_context._write_interceptor = None
+
+        mock_canvas.get_by_id.return_value = (True, MagicMock(dsl='{"id": "test"}'))
+        chunks = {
+            "chunks": [
+                {"text": "Hello"},
+            ],
+            "embedding_token_consumption": 1,
+        }
+        mock_pipeline = MagicMock()
+        mock_pipeline.run = AsyncMock(return_value=chunks)
+        mock_pipeline_class.return_value = mock_pipeline
+
+        billing_hook = MagicMock()
+        billing_hook.on_pipeline_success = AsyncMock()
+        billing_hook.on_pipeline_error = AsyncMock()
+
+        with patch.object(DataflowService, '_embed_chunks', new_callable=AsyncMock, return_value=(chunks["chunks"], 1)):
+            with patch.object(DataflowService, '_insert_chunks', new_callable=AsyncMock, return_value=True):
+                with patch.object(DataflowService, '_update_document_metadata'):
+                    with patch.object(DataflowService, '_record_pipeline_log'):
+                        with patch("api.db.services.document_service.DocumentService.increment_chunk_num"):
+                            service = DataflowService(ctx=task_context, billing_hook=billing_hook)
+                            await service.run_dataflow()
+
+                            billing_hook.on_pipeline_success.assert_called_once()
+                            billing_hook.on_pipeline_error.assert_not_called()
+
+    @pytest.mark.asyncio
+    @patch("rag.svr.task_executor_refactor.dataflow_service.Pipeline")
+    @patch("rag.svr.task_executor_refactor.dataflow_service.UserCanvasService")
+    async def test_run_dataflow_with_billing_hook_error(self, mock_canvas, mock_pipeline_class, task_context):
+        """Test run_dataflow calls billing hook on error."""
+        task_context._task["task_type"] = "dataflow"
+        task_context._task["dataflow_id"] = "dataflow_test"
+        task_context._task["name"] = "test.pdf"
+        task_context._write_interceptor = None
+
+        mock_canvas.get_by_id.return_value = (True, MagicMock(dsl='{"id": "test"}'))
+        mock_pipeline = MagicMock()
+        mock_pipeline.run = AsyncMock(side_effect=Exception("Pipeline failure"))
+        mock_pipeline_class.return_value = mock_pipeline
+
+        billing_hook = MagicMock()
+        billing_hook.on_pipeline_success = AsyncMock()
+        billing_hook.on_pipeline_error = AsyncMock()
+
+        service = DataflowService(ctx=task_context, billing_hook=billing_hook)
+        with pytest.raises(Exception, match="Pipeline failure"):
+            await service.run_dataflow()
+
+        billing_hook.on_pipeline_error.assert_called_once()
+        billing_hook.on_pipeline_success.assert_not_called()
+
+
+class TestDataflowServiceNormalizeChunks:
+    """Tests for _normalize_chunks — stable pure helper for output-format normalization."""
+
+    def test_normalize_chunks_from_chunks_key(self):
+        """Test normalization from 'chunks' key."""
+        result = DataflowService._normalize_chunks({"chunks": [{"a": 1}]})
+        assert result == [{"a": 1}]
+
+    def test_normalize_chunks_from_json_key(self):
+        """Test normalization from 'json' key."""
+        result = DataflowService._normalize_chunks({"json": [{"a": 1}]})
+        assert result == [{"a": 1}]
+
+    def test_normalize_chunks_from_markdown_key(self):
+        """Test normalization from 'markdown' key."""
+        result = DataflowService._normalize_chunks({"markdown": "# Title"})
+        assert result == [{"text": ["# Title"]}]
+
+    def test_normalize_chunks_from_text_key(self):
+        """Test normalization from 'text' key."""
+        result = DataflowService._normalize_chunks({"text": "plain text"})
+        assert result == [{"text": ["plain text"]}]
+
+    def test_normalize_chunks_from_html_key(self):
+        """Test normalization from 'html' key."""
+        result = DataflowService._normalize_chunks({"html": "<p>content</p>"})
+        assert result == [{"text": ["<p>content</p>"]}]
+
+    def test_normalize_chunks_unknown_key(self):
+        """Test normalization with unknown key returns empty."""
+        result = DataflowService._normalize_chunks({"unknown": "data"})
+        assert result == []
+
+    def test_normalize_chunks_empty_markdown(self):
+        """Test normalization with empty markdown value returns empty."""
+        result = DataflowService._normalize_chunks({"markdown": ""})
+        assert result == []
+
+    def test_normalize_chunks_preserves_deepcopy(self):
+        """Test normalization returns a deepcopy so mutations don't leak."""
+        input_data = {"chunks": [{"key": "value"}]}
+        result = DataflowService._normalize_chunks(input_data)
+        result[0]["key"] = "modified"
+        assert input_data["chunks"][0]["key"] == "value"
+
+
+class TestDataflowServiceGetOutputType:
+    """Tests for _get_output_type — stable pure helper for output-type detection."""
+
+    def test_get_output_type_chunks(self):
+        assert DataflowService._get_output_type({"chunks": []}) == "chunks"
+
+    def test_get_output_type_json(self):
+        assert DataflowService._get_output_type({"json": []}) == "json"
+
+    def test_get_output_type_markdown(self):
+        assert DataflowService._get_output_type({"markdown": ""}) == "markdown"
+
+    def test_get_output_type_text(self):
+        assert DataflowService._get_output_type({"text": ""}) == "text"
+
+    def test_get_output_type_html(self):
+        assert DataflowService._get_output_type({"html": ""}) == "html"
+
+    def test_get_output_type_empty(self):
+        assert DataflowService._get_output_type({}) == "empty"
+
+
+class TestDataflowServiceProcessChunks:
+    """Tests for _process_chunks — stable pure helper for chunk metadata processing."""
+
+    def test_process_chunks_adds_doc_id_and_kb_id(self, task_context):
+        """Test _process_chunks adds doc_id, kb_id, and metadata."""
+        task_context._task["doc_id"] = "doc_123"
+        task_context._task["kb_id"] = "kb_456"
+        task_context._task["name"] = "test.pdf"
+        chunks = [{"text": "content"}]
+        DataflowService._process_chunks(DataflowService(ctx=task_context), chunks)
+        assert chunks[0]["doc_id"] == "doc_123"
+        assert "kb_id" in chunks[0]
+        assert "content_with_weight" in chunks[0]
+        assert "text" not in chunks[0]
+
+    def test_process_chunks_generates_id(self, task_context):
+        """Test _process_chunks auto-generates id."""
+        task_context._task["doc_id"] = "doc_123"
+        task_context._task["kb_id"] = "kb_456"
+        task_context._task["name"] = "test.pdf"
+        chunks = [{"text": "content"}]
+        DataflowService._process_chunks(DataflowService(ctx=task_context), chunks)
+        assert "id" in chunks[0]
+
+    def test_process_chunks_questions_field(self, task_context):
+        """Test _process_chunks processes questions field."""
+        task_context._task["doc_id"] = "doc_123"
+        task_context._task["kb_id"] = "kb_456"
+        task_context._task["name"] = "test.pdf"
+        chunks = [{"text": "content", "questions": "Q1\nQ2"}]
+        DataflowService._process_chunks(DataflowService(ctx=task_context), chunks)
+        assert "questions" not in chunks[0]
+        assert "question_kwd" in chunks[0]
+
+    def test_process_chunks_summary_field(self, task_context):
+        """Test _process_chunks processes summary field."""
+        task_context._task["doc_id"] = "doc_123"
+        task_context._task["kb_id"] = "kb_456"
+        task_context._task["name"] = "test.pdf"
+        chunks = [{"text": "content", "summary": "summary text"}]
+        DataflowService._process_chunks(DataflowService(ctx=task_context), chunks)
+        assert "summary" not in chunks[0]
+        assert "content_ltks" in chunks[0]
+
+    def test_process_chunks_metadata_field(self, task_context):
+        """Test _process_chunks extracts metadata."""
+        task_context._task["doc_id"] = "doc_123"
+        task_context._task["kb_id"] = "kb_456"
+        task_context._task["name"] = "test.pdf"
+        chunks = [{"text": "content", "metadata": {"key": "val"}}]
+        metadata = DataflowService._process_chunks(DataflowService(ctx=task_context), chunks)
+        assert "metadata" not in chunks[0]
+        assert "key" in metadata
+
+
+class TestDataflowServiceInit:
+    """Tests for DataflowService initialization."""
+
+    @patch("rag.svr.task_executor_refactor.dataflow_service.settings")
+    def test_init_with_custom_batch_sizes(self, mock_settings):
+        """Test initialization with custom batch sizes."""
+        ctx = MagicMock()
+        service = DataflowService(ctx=ctx, embedding_batch_size=64, doc_bulk_size=50)
+        assert service._embedding_batch_size == 64
+        assert service._doc_bulk_size == 50
+
+    @patch("rag.svr.task_executor_refactor.dataflow_service.settings")
+    def test_init_with_default_sizes(self, mock_settings):
+        """Test initialization with default batch sizes."""
+        mock_settings.EMBEDDING_BATCH_SIZE = 32
+        mock_settings.DOC_BULK_SIZE = 100
+        ctx = MagicMock()
+        service = DataflowService(ctx=ctx)
+        assert service._embedding_batch_size == 32
+        assert service._doc_bulk_size == 100
+
+    def test_init_stores_context_and_hook(self):
+        """Test initialization stores context and billing hook."""
+        ctx = MagicMock()
+        hook = MagicMock()
+        service = DataflowService(ctx=ctx, billing_hook=hook)
+        assert service._task_context is ctx
+        assert service._billing_hook is hook

test/unit_test/rag/svr/task_executor_refactor/test_embedding_service.py

@@ -0,0 +1,145 @@
+#
+#  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.
+
+"""
+Unit tests for EmbeddingService module.
+
+All tests validate behavior through the public API (embed_chunks) rather than
+reaching into private orchestration methods like _encode_single, _encode_batch,
+or _run_encode.  Those internal boundaries may be reshaped during a refactor
+without changing the external behavior; the suite should not break in that case.
+"""
+
+import numpy as np
+from unittest.mock import MagicMock, patch
+
+from rag.svr.task_executor_refactor.embedding_service import EmbeddingService
+
+
+class TestEmbeddingServiceInit:
+    """Tests for EmbeddingService initialization."""
+
+    @patch("rag.svr.task_executor_refactor.embedding_service.settings")
+    def test_init_with_default_batch_size(self, mock_settings):
+        """Test initialization with default batch size."""
+        mock_settings.EMBEDDING_BATCH_SIZE = 32
+        ctx = MagicMock()
+        service = EmbeddingService(ctx=ctx)
+        assert service._embedding_batch_size == 32
+
+    @patch("rag.svr.task_executor_refactor.embedding_service.settings")
+    def test_init_with_custom_batch_size(self, mock_settings):
+        """Test initialization with custom batch size."""
+        ctx = MagicMock()
+        service = EmbeddingService(ctx=ctx, embedding_batch_size=64)
+        assert service._embedding_batch_size == 64
+
+    def test_init_stores_task_context(self):
+        """Test that task context is stored."""
+        ctx = MagicMock()
+        service = EmbeddingService(ctx=ctx)
+        assert service._task_context is ctx
+
+
+class TestEmbeddingServiceEmbedChunks:
+    """Tests for the public embed_chunks method.
+
+    Internal helpers _encode_single, _encode_batch, and _run_encode are
+    exercised through this public entry point so the suite stays resilient to
+    method-boundary reshuffles.
+    """
+
+    @patch.object(EmbeddingService, '_run_encode')
+    def test_embed_chunks_basic(self, mock_run_encode):
+        """Test basic chunk embedding."""
+        mock_run_encode.return_value = (np.array([[1.0, 2.0]]), 10)
+        ctx = MagicMock()
+        ctx.progress_cb = None
+        service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
+        model = MagicMock()
+        model.max_length = 100
+
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "Content1"},
+        ]
+        tk_count, vector_size = service.embed_chunks(docs, model)
+
+        assert tk_count > 0
+        assert vector_size == 2
+        assert "q_2_vec" in docs[0]
+
+    @patch.object(EmbeddingService, '_run_encode')
+    def test_embed_chunks_uses_embedding_utils(self, mock_run_encode):
+        """Test that embed_chunks uses EmbeddingUtils internally.
+
+        The internal path runs _encode_batch -> EmbeddingUtils.truncate_texts
+        -> _run_encode.  We verify via the public embed_chunks that the chain
+        is wired correctly without asserting on individual private method calls.
+        """
+        mock_run_encode.return_value = (np.array([[1.0, 2.0]]), 10)
+        ctx = MagicMock()
+        ctx.progress_cb = None
+        service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
+        model = MagicMock()
+        model.max_length = 100
+
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "Content1"},
+        ]
+        service.embed_chunks(docs, model)
+
+        mock_run_encode.assert_called()
+
+    @patch.object(EmbeddingService, '_run_encode')
+    def test_embed_chunks_with_title_content_combination(self, mock_run_encode):
+        """Test that title and content vectors are combined."""
+        mock_run_encode.return_value = (np.array([[1.0, 2.0]]), 10)
+        ctx = MagicMock()
+        ctx.progress_cb = None
+        service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
+        model = MagicMock()
+        model.max_length = 100
+
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "Content1"},
+        ]
+        _, vector_size = service.embed_chunks(docs, model, parser_config={"filename_embd_weight": 0.5})
+
+        assert vector_size == 2
+
+    @patch.object(EmbeddingService, '_run_encode')
+    def test_embed_chunks_handles_long_text(self, mock_run_encode):
+        """Test that long texts are handled by embedding pipeline.
+
+        Even with content exceeding model.max_length, embed_chunks produces
+        valid vectors, meaning truncation (via EmbeddingUtils) is wired
+        correctly in the encode path.
+        """
+        mock_run_encode.return_value = (np.array([[1.0, 2.0]]), 10)
+        ctx = MagicMock()
+        ctx.progress_cb = None
+        service = EmbeddingService(ctx=ctx, embedding_batch_size=10)
+        model = MagicMock()
+        model.max_length = 100
+
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "a" * 200},
+        ]
+        tk_count, vector_size = service.embed_chunks(docs, model)
+
+        # Public contract: embed_chunks returns valid token counts and vectors
+        assert tk_count > 0
+        assert vector_size == 2
+        assert "q_2_vec" in docs[0]

test/unit_test/rag/svr/task_executor_refactor/test_embedding_utils.py

@@ -0,0 +1,331 @@
+#
+#  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.
+
+"""
+Unit tests for EmbeddingUtils module.
+"""
+
+import numpy as np
+from unittest.mock import patch
+from rag.svr.task_executor_refactor.embedding_utils import EmbeddingUtils
+
+
+class TestEmbeddingUtilsPrepareTexts:
+    """Tests for prepare_texts_for_embedding class method."""
+
+    def test_prepare_texts_basic(self):
+        """Test basic text preparation."""
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "Content1"},
+            {"docnm_kwd": "Title2", "content_with_weight": "Content2"},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs)
+        assert titles == ["Title1", "Title2"]
+        assert contents == ["Content1", "Content2"]
+
+    def test_prepare_texts_with_question_kwd(self):
+        """Test text preparation with question_kwd."""
+        docs = [
+            {"docnm_kwd": "Title1", "question_kwd": ["Q1", "Q2"], "content_with_weight": "Content1"},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs)
+        assert titles == ["Title1"]
+        assert contents == ["Q1\nQ2"]
+
+    def test_prepare_texts_with_empty_question_kwd(self):
+        """Test text preparation with empty question_kwd falls back to content."""
+        docs = [
+            {"docnm_kwd": "Title1", "question_kwd": [], "content_with_weight": "Content1"},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs)
+        assert contents == ["Content1"]
+
+    def test_prepare_texts_with_missing_question_kwd(self):
+        """Test text preparation without question_kwd uses content."""
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "Content1"},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs)
+        assert contents == ["Content1"]
+
+    def test_prepare_texts_normalizes_table_html(self):
+        """Test that table HTML tags are normalized."""
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "<table><tr><td>Cell</td></tr></table>"},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs)
+        # Table tags should be replaced with spaces
+        assert "<table>" not in contents[0]
+
+    def test_prepare_texts_whitespace_only_becomes_none(self):
+        """Test that whitespace-only content becomes 'None'."""
+        docs = [
+            {"docnm_kwd": "Title1", "content_with_weight": "   \n\n  "},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs)
+        assert contents == ["None"]
+
+    def test_prepare_texts_default_title(self):
+        """Test that missing docnm_kwd uses 'Title' as default."""
+        docs = [
+            {"content_with_weight": "Content1"},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs)
+        assert titles == ["Title"]
+
+    def test_prepare_texts_without_question_kwd(self):
+        """Test text preparation with use_question_kwd=False."""
+        docs = [
+            {"docnm_kwd": "Title1", "question_kwd": ["Q1"], "content_with_weight": "Content1"},
+        ]
+        titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs, use_question_kwd=False)
+        assert contents == ["Content1"]
+
+
+class TestEmbeddingUtilsPrepareDataflowTexts:
+    """Tests for prepare_texts_for_dataflow_embedding class method."""
+
+    def test_prepare_dataflow_texts_with_questions(self):
+        """Test dataflow text preparation with questions field."""
+        chunks = [
+            {"questions": "Q1\nQ2"},
+            {"questions": "Q3"},
+        ]
+        texts = EmbeddingUtils.prepare_texts_for_dataflow_embedding(chunks)
+        assert texts == ["Q1\nQ2", "Q3"]
+
+    def test_prepare_dataflow_texts_with_summary(self):
+        """Test dataflow text preparation with summary field (no questions)."""
+        chunks = [
+            {"summary": "Summary1"},
+        ]
+        texts = EmbeddingUtils.prepare_texts_for_dataflow_embedding(chunks)
+        assert texts == ["Summary1"]
+
+    def test_prepare_dataflow_texts_with_text(self):
+        """Test dataflow text preparation with text field (no questions/summary)."""
+        chunks = [
+            {"text": "Text content"},
+        ]
+        texts = EmbeddingUtils.prepare_texts_for_dataflow_embedding(chunks)
+        assert texts == ["Text content"]
+
+    def test_prepare_dataflow_texts_priority(self):
+        """Test field priority: questions > summary > text."""
+        chunks = [
+            {"questions": "Q", "summary": "S", "text": "T"},
+        ]
+        texts = EmbeddingUtils.prepare_texts_for_dataflow_embedding(chunks)
+        assert texts == ["Q"]
+
+        chunks = [
+            {"summary": "S", "text": "T"},
+        ]
+        texts = EmbeddingUtils.prepare_texts_for_dataflow_embedding(chunks)
+        assert texts == ["S"]
+
+
+class TestEmbeddingUtilsTruncateTexts:
+    """Tests for truncate_texts class method."""
+
+    @patch("rag.svr.task_executor_refactor.embedding_utils.truncate")
+    def test_truncate_texts_calls_truncate(self, mock_truncate):
+        """Test truncate_texts calls truncate with correct max_length."""
+        mock_truncate.return_value = "truncated"
+        texts = ["long text 1", "long text 2"]
+        max_length = 100
+
+        _ = EmbeddingUtils.truncate_texts(texts, max_length)
+
+        assert mock_truncate.call_count == 2
+        # Should subtract 10 for safety margin
+        mock_truncate.assert_called_with("long text 2", 90)
+
+    @patch("rag.svr.task_executor_refactor.embedding_utils.truncate")
+    def test_truncate_texts_returns_list(self, mock_truncate):
+        """Test truncate_texts returns a list of same length."""
+        mock_truncate.return_value = "truncated"
+        texts = ["text1", "text2", "text3"]
+        result = EmbeddingUtils.truncate_texts(texts, 50)
+        assert len(result) == 3
+
+
+class TestEmbeddingUtilsStackVectors:
+    """Tests for stack_vectors class method."""
+
+    def test_stack_vectors_with_multiple_batches(self):
+        """Test stacking multiple vector batches."""
+        batch1 = np.array([[1.0, 2.0], [3.0, 4.0]])
+        batch2 = np.array([[5.0, 6.0]])
+        result = EmbeddingUtils.stack_vectors([batch1, batch2])
+        expected = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+        np.testing.assert_array_equal(result, expected)
+
+    def test_stack_vectors_with_empty_batches(self):
+        """Test stacking empty batches returns empty array."""
+        result = EmbeddingUtils.stack_vectors([])
+        assert result.size == 0
+
+    def test_stack_vectors_with_single_batch(self):
+        """Test stacking a single batch."""
+        batch = np.array([[1.0, 2.0]])
+        result = EmbeddingUtils.stack_vectors([batch])
+        np.testing.assert_array_equal(result, batch)
+
+
+class TestEmbeddingUtilsAttachVectors:
+    """Tests for attach_vectors class method."""
+
+    def test_attach_vectors_basic(self):
+        """Test attaching vectors to docs."""
+        docs = [{"id": 1}, {"id": 2}]
+        vectors = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+
+        vector_size = EmbeddingUtils.attach_vectors(docs, vectors)
+
+        assert vector_size == 3
+        assert "q_3_vec" in docs[0]
+        assert "q_3_vec" in docs[1]
+        assert docs[0]["q_3_vec"] == [1.0, 2.0, 3.0]
+        assert docs[1]["q_3_vec"] == [4.0, 5.0, 6.0]
+
+    def test_attach_vectors_custom_key_template(self):
+        """Test attaching vectors with custom key template."""
+        docs = [{"id": 1}]
+        vectors = np.array([[1.0, 2.0]])
+
+        EmbeddingUtils.attach_vectors(docs, vectors, vector_key_template="vec_%d")
+
+        assert "vec_2" in docs[0]
+
+    def test_attach_vectors_modifies_in_place(self):
+        """Test that attach_vectors modifies docs in place."""
+        docs = [{"id": 1}]
+        vectors = np.array([[1.0, 2.0]])
+        original_id = id(docs)
+
+        EmbeddingUtils.attach_vectors(docs, vectors)
+
+        assert id(docs) == original_id
+
+
+class TestEmbeddingUtilsCombineVectors:
+    """Tests for combine_title_content_vectors class method."""
+
+    def test_combine_vectors_with_title_and_content(self):
+        """Test combining title and content vectors with weight."""
+        title_vecs = np.array([[1.0, 2.0], [3.0, 4.0]])
+        content_vecs = np.array([[5.0, 6.0], [7.0, 8.0]])
+
+        result = EmbeddingUtils.combine_title_content_vectors(title_vecs, content_vecs, title_weight=0.3)
+
+        # Expected: 0.3 * title + 0.7 * content
+        expected = 0.3 * title_vecs + 0.7 * content_vecs
+        np.testing.assert_array_almost_equal(result, expected)
+
+    def test_combine_vectors_with_default_weight(self):
+        """Test combining with default weight when not specified."""
+        title_vecs = np.array([[1.0, 2.0]])
+        content_vecs = np.array([[5.0, 6.0]])
+
+        result = EmbeddingUtils.combine_title_content_vectors(title_vecs, content_vecs)
+
+        # Expected: 0.1 * title + 0.9 * content (default weight is 0.1)
+        expected = 0.1 * title_vecs + 0.9 * content_vecs
+        np.testing.assert_array_almost_equal(result, expected)
+
+    def test_combine_vectors_with_none_title(self):
+        """Test combining when title vectors is None returns content."""
+        content_vecs = np.array([[5.0, 6.0]])
+
+        result = EmbeddingUtils.combine_title_content_vectors(None, content_vecs, title_weight=0.3)
+
+        np.testing.assert_array_equal(result, content_vecs)
+
+    def test_combine_vectors_with_mismatched_shapes(self):
+        """Test combining when shapes don't match returns content."""
+        title_vecs = np.array([[1.0, 2.0]])
+        content_vecs = np.array([[5.0, 6.0], [7.0, 8.0]])
+
+        result = EmbeddingUtils.combine_title_content_vectors(title_vecs, content_vecs, title_weight=0.3)
+
+        # Should return content_vecs when shapes don't match
+        np.testing.assert_array_equal(result, content_vecs)
+
+    def test_combine_vectors_with_zero_weight(self):
+        """Test combining when weight is 0 uses default 0.1."""
+        title_vecs = np.array([[1.0, 2.0]])
+        content_vecs = np.array([[5.0, 6.0]])
+
+        result = EmbeddingUtils.combine_title_content_vectors(title_vecs, content_vecs, title_weight=0)
+
+        # Should use default weight of 0.1
+        expected = 0.1 * title_vecs + 0.9 * content_vecs
+        np.testing.assert_array_almost_equal(result, expected)
+
+
+class TestEmbeddingUtilsInternals:
+    """Tests for internal helper methods."""
+
+    def test_extract_content_with_question_kwd(self):
+        """Test _extract_content with question_kwd."""
+        doc = {"question_kwd": ["Q1", "Q2"], "content_with_weight": "Content"}
+        result = EmbeddingUtils._extract_content(doc, use_question_kwd=True)
+        assert result == "Q1\nQ2"
+
+    def test_extract_content_without_question_kwd(self):
+        """Test _extract_content without question_kwd."""
+        doc = {"content_with_weight": "Content"}
+        result = EmbeddingUtils._extract_content(doc, use_question_kwd=True)
+        assert result == "Content"
+
+    def test_extract_content_with_use_question_false(self):
+        """Test _extract_content with use_question_kwd=False."""
+        doc = {"question_kwd": ["Q1"], "content_with_weight": "Content"}
+        result = EmbeddingUtils._extract_content(doc, use_question_kwd=False)
+        assert result == "Content"
+
+    def test_normalize_table_html(self):
+        """Test _normalize_table_html removes table tags."""
+        html = "<table><tr><td>Cell</td></tr></table>"
+        result = EmbeddingUtils._normalize_table_html(html)
+        assert "<table>" not in result
+        assert "<tr>" not in result
+        assert "<td>" not in result
+
+    def test_handle_whitespace(self):
+        """Test _handle_whitespace replaces whitespace-only with placeholder."""
+        assert EmbeddingUtils._handle_whitespace("   \n  ") == "None"
+        assert EmbeddingUtils._handle_whitespace("  text  ") == "  text  "
+
+    def test_handle_whitespace_with_empty_string(self):
+        """Test _handle_whitespace with empty string."""
+        assert EmbeddingUtils._handle_whitespace("") == "None"
+
+
+class TestEmbeddingUtilsConstants:
+    """Tests for class constants."""
+
+    def test_default_title_weight(self):
+        """Test DEFAULT_TITLE_WEIGHT value."""
+        assert EmbeddingUtils.DEFAULT_TITLE_WEIGHT == 0.1
+
+    def test_default_title_placeholder(self):
+        """Test DEFAULT_TITLE_PLACEHOLDER value."""
+        assert EmbeddingUtils.DEFAULT_TITLE_PLACEHOLDER == "Title"
+
+    def test_content_placeholder_for_whitespace(self):
+        """Test CONTENT_PLACEHOLDER_FOR_WHITESPACE value."""
+        assert EmbeddingUtils.CONTENT_PLACEHOLDER_FOR_WHITESPACE == "None"

test/unit_test/rag/svr/task_executor_refactor/test_post_processor.py

@@ -0,0 +1,130 @@
+#
+#  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.
+
+"""
+Unit tests for PostProcessor module.
+"""
+
+import pytest
+from unittest.mock import MagicMock, patch, AsyncMock
+from rag.svr.task_executor_refactor.post_processor import PostProcessor
+
+
+class TestPostProcessorInit:
+    """Tests for PostProcessor initialization."""
+
+    def test_init_stores_task_context(self):
+        """Test that task context is stored."""
+        ctx = MagicMock()
+        service = PostProcessor(ctx=ctx)
+        assert service._task_context is ctx
+
+
+class TestPostProcessorProcessTableParserMetadata:
+    """Tests for process_table_parser_metadata method."""
+
+    @pytest.mark.asyncio
+    async def test_skips_non_table_parser(self):
+        """Test that processing is skipped for non-table parser."""
+        ctx = MagicMock()
+        ctx.parser_id = "naive"
+        service = PostProcessor(ctx=ctx)
+
+        await service.process_table_parser_metadata("doc_1", [])
+
+        # Should return early without any further processing
+
+    @pytest.mark.asyncio
+    async def test_skips_when_not_manual_column_mode(self):
+        """Test that processing is skipped when not in manual column mode."""
+        ctx = MagicMock()
+        ctx.parser_id = "table"
+        ctx.raw_task = {}
+        service = PostProcessor(ctx=ctx)
+
+        with patch("rag.svr.task_executor_refactor.post_processor.merge_table_parser_config_from_kb") as mock_merge:
+            mock_merge.return_value = {"table_column_mode": "auto"}
+            await service.process_table_parser_metadata("doc_1", [])
+
+            mock_merge.assert_called_once()
+
+
+class TestPostProcessorInsertTocChunk:
+    """Tests for insert_toc_chunk method."""
+
+    @pytest.mark.asyncio
+    async def test_returns_false_for_none_chunk(self):
+        """Test that method returns False when chunk is None."""
+        ctx = MagicMock()
+        service = PostProcessor(ctx=ctx)
+        chunk_service = MagicMock()
+
+        result = await service.insert_toc_chunk(None, chunk_service)
+
+        assert result is False
+        chunk_service.insert_chunks.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_checks_cancellation(self):
+        """Test that cancellation is checked."""
+        ctx = MagicMock()
+        ctx.id = "task_1"
+        ctx.has_canceled_func = MagicMock(return_value=True)
+        ctx.progress_cb = MagicMock()
+        service = PostProcessor(ctx=ctx)
+        chunk_service = MagicMock()
+        toc_chunk = {"id": "toc_1"}
+
+        result = await service.insert_toc_chunk(toc_chunk, chunk_service)
+
+        assert result is False
+        ctx.progress_cb.assert_called_with(-1, msg="Task has been canceled.")
+
+    @pytest.mark.asyncio
+    async def test_inserts_toc_chunk_successfully(self):
+        """Test successful TOC chunk insertion."""
+        ctx = MagicMock()
+        ctx.id = "task_1"
+        ctx.tenant_id = "tenant_1"
+        ctx.kb_id = "kb_1"
+        ctx.has_canceled_func = MagicMock(return_value=False)
+        service = PostProcessor(ctx=ctx)
+        chunk_service = AsyncMock()
+        chunk_service.insert_chunks = AsyncMock(return_value=True)
+        toc_chunk = {"id": "toc_1"}
+
+        result = await service.insert_toc_chunk(toc_chunk, chunk_service)
+
+        assert result is True
+        chunk_service.insert_chunks.assert_called_once_with(
+            "task_1", "tenant_1", "kb_1", [toc_chunk]
+        )
+
+    @pytest.mark.asyncio
+    async def test_handles_insert_failure(self):
+        """Test handling of insert failure."""
+        ctx = MagicMock()
+        ctx.id = "task_1"
+        ctx.tenant_id = "tenant_1"
+        ctx.kb_id = "kb_1"
+        ctx.has_canceled_func = MagicMock(return_value=False)
+        service = PostProcessor(ctx=ctx)
+        chunk_service = AsyncMock()
+        chunk_service.insert_chunks = AsyncMock(return_value=False)
+        toc_chunk = {"id": "toc_1"}
+
+        result = await service.insert_toc_chunk(toc_chunk, chunk_service)
+
+        assert result is False

test/unit_test/rag/svr/task_executor_refactor/test_raptor_service.py

@@ -0,0 +1,452 @@
+#
+#  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.
+
+"""
+Tests for RaptorService.
+
+Coverage is driven through the public entry point `run_raptor_for_kb()`.
+
+Design principles:
+- All orchestration behavior is validated through the public API.
+- Only stable pure helpers (`_collect_doc_info`, `_schedule_raptor_cleanup`)
+  are tested directly.
+- Internal methods (`_run_file_level_raptor`, `_run_dataset_level_raptor`,
+  `_should_skip_raptor`, `_load_doc_chunks`, `_load_all_doc_chunks`,
+  `_generate_raptor`, `_get_raptor_chunk_methods`) are NOT tested directly —
+  their behavior is covered by exercising `run_raptor_for_kb()` with
+  appropriate mocks.
+"""
+
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from rag.svr.task_executor_refactor.raptor_service import RaptorService
+
+
+# =============================================================================
+# Stable Pure Helpers (tested directly)
+# =============================================================================
+
+
+class TestRaptorServiceInit:
+    """Tests for RaptorService initialization."""
+
+    def test_init_stores_task_context(self, mock_raptor_context):
+        svc = RaptorService(mock_raptor_context)
+        assert svc._task_context is mock_raptor_context
+
+    def test_init_uses_provided_kb_id(self, mock_raptor_context):
+        mock_raptor_context.kb_id = "custom_kb"
+        svc = RaptorService(mock_raptor_context)
+        assert svc._task_context.kb_id == "custom_kb"
+
+
+class TestRaptorServiceCollectDocInfo:
+    """Tests for _collect_doc_info — stable pure data aggregation (classmethod)."""
+
+    def _make_mock_doc(self, name, type, parser_id, parser_config):
+        """Create a mock document with accessible attributes."""
+        mock_doc = MagicMock()
+        mock_doc.name = name
+        mock_doc.type = type
+        mock_doc.parser_id = parser_id
+        mock_doc.parser_config = parser_config
+        return mock_doc
+
+    def test_collect_doc_info_success(self):
+        doc_ids = ["doc_1", "doc_2"]
+
+        mock_doc_1 = self._make_mock_doc(name="", type="pdf", parser_id="naive", parser_config={})
+        mock_doc_2 = self._make_mock_doc(name="doc2.txt", type="txt", parser_id="manual", parser_config={"chunk_token_num": 512})
+
+        def get_by_id_side_effect(doc_id):
+            if doc_id == "doc_1":
+                return True, mock_doc_1
+            if doc_id == "doc_2":
+                return True, mock_doc_2
+            return False, None
+
+        with patch("rag.svr.task_executor_refactor.raptor_service.DocumentService") as mock_ds:
+            mock_ds.get_by_id = MagicMock(side_effect=get_by_id_side_effect)
+            result = RaptorService._collect_doc_info(doc_ids)
+
+        assert len(result) == 2
+        assert result["doc_1"]["name"] == ""
+        assert result["doc_1"]["type"] == "pdf"
+        assert result["doc_1"]["parser_id"] == "naive"
+        assert result["doc_2"]["name"] == "doc2.txt"
+        assert result["doc_2"]["type"] == "txt"
+        assert result["doc_2"]["parser_id"] == "manual"
+        assert result["doc_2"]["parser_config"] == {"chunk_token_num": 512}
+
+    def test_collect_doc_info_empty_input(self):
+        result = RaptorService._collect_doc_info([])
+        assert result == {}
+
+    def test_collect_doc_info_deduplicates_doc_ids(self):
+        """Duplicate doc_ids should be deduplicated."""
+        doc_ids = ["doc_1", "doc_1", "doc_2"]
+
+        mock_doc = self._make_mock_doc(name="test.pdf", type="pdf", parser_id="naive", parser_config={})
+
+        called_ids = []
+
+        def get_by_id_side_effect(doc_id):
+            called_ids.append(doc_id)
+            return True, mock_doc
+
+        with patch("rag.svr.task_executor_refactor.raptor_service.DocumentService") as mock_ds:
+            mock_ds.get_by_id = MagicMock(side_effect=get_by_id_side_effect)
+            result = RaptorService._collect_doc_info(doc_ids)
+
+        assert sorted(called_ids) == ["doc_1", "doc_2"]
+        assert len(result) == 2
+
+    def test_collect_doc_info_missing_document(self):
+        doc_ids = ["doc_1", "missing_doc"]
+
+        mock_doc = self._make_mock_doc(name="test.pdf", type="pdf", parser_id="naive", parser_config={})
+
+        def get_by_id_side_effect(doc_id):
+            if doc_id == "doc_1":
+                return True, mock_doc
+            return False, None
+
+        with patch("rag.svr.task_executor_refactor.raptor_service.DocumentService") as mock_ds:
+            mock_ds.get_by_id = MagicMock(side_effect=get_by_id_side_effect)
+            result = RaptorService._collect_doc_info(doc_ids)
+
+        assert len(result) == 1
+        assert "doc_1" in result
+        assert "missing_doc" not in result
+
+
+class TestRaptorServiceScheduleRaptorCleanup:
+    """Tests for _schedule_raptor_cleanup — stable pure data operation (classmethod)."""
+
+    def test_schedule_cleanup_adds_entry(self):
+        cleanup_list = []
+        RaptorService._schedule_raptor_cleanup("doc_1", "tree_builder_a", cleanup_list)
+        assert cleanup_list == [("doc_1", "tree_builder_a")]
+
+    def test_schedule_cleanup_deduplicates(self):
+        cleanup_list = [("doc_1", "tree_builder_a")]
+        RaptorService._schedule_raptor_cleanup("doc_1", "tree_builder_a", cleanup_list)
+        assert len(cleanup_list) == 1
+
+    def test_schedule_cleanup_keep_method_none(self):
+        cleanup_list = []
+        RaptorService._schedule_raptor_cleanup("doc_1", None, cleanup_list)
+        assert cleanup_list == [("doc_1", None)]
+
+    def test_schedule_cleanup_multiple_docs(self):
+        cleanup_list = []
+        RaptorService._schedule_raptor_cleanup("doc_1", "t1", cleanup_list)
+        RaptorService._schedule_raptor_cleanup("doc_2", "t2", cleanup_list)
+        RaptorService._schedule_raptor_cleanup("doc_3", None, cleanup_list)
+        assert len(cleanup_list) == 3
+        assert ("doc_1", "t1") in cleanup_list
+        assert ("doc_2", "t2") in cleanup_list
+        assert ("doc_3", None) in cleanup_list
+
+
+# =============================================================================
+# Public Entry Point Tests
+# =============================================================================
+
+
+class TestRaptorServiceRunRaptorForKb:
+    """Tests for run_raptor_for_kb() — the public entry point.
+
+    All orchestration behavior (file-level vs dataset-level dispatch,
+    chunk loading, skip logic, cleanup scheduling) is validated through
+    this method by mocking internal helpers and observing:
+    - Return values (chunks, token_count, cleanup_raptor_chunks)
+    - Mock call patterns (which internal method was invoked, with what args)
+    """
+
+    @pytest.fixture
+    def sample_chunks(self):
+        """Sample RAPTOR summary chunks returned by internal methods."""
+        return [{"id": "chunk_1", "content_with_weight": "Summary 1"}]
+
+    @pytest.fixture
+    def raptor_config_file_scope(self):
+        """RAPTOR config with file-level scope."""
+        return {
+            "raptor": {
+                "tree_builder": "raptor",
+                "clustering_method": "gmm",
+                "scope": "file",
+                "prompt": "summarize",
+                "max_token": 512,
+                "threshold": 0.5,
+                "max_cluster": 64,
+                "random_seed": 42,
+            }
+        }
+
+    @pytest.fixture
+    def raptor_config_dataset_scope(self):
+        """RAPTOR config with dataset-level scope."""
+        return {
+            "raptor": {
+                "tree_builder": "raptor",
+                "clustering_method": "gmm",
+                "scope": "dataset",
+                "prompt": "summarize",
+                "max_token": 512,
+                "threshold": 0.5,
+                "max_cluster": 64,
+                "random_seed": 42,
+            }
+        }
+
+    # ---- Basic dispatch (file-level scope) ----
+
+    def test_run_raptor_for_kb_file_scope_delegates_to_file_level(
+        self, mock_raptor_context, sample_chunks, raptor_config_file_scope
+    ):
+        """When scope='file', _run_file_level_raptor is called."""
+        svc = RaptorService(mock_raptor_context)
+        doc_ids = ["doc_1", "doc_2"]
+        chat_mdl = MagicMock()
+        embd_mdl = MagicMock()
+        vector_size = 128
+
+        with patch.object(svc, "_collect_doc_info", return_value={
+            "doc_1": {"name": "a.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}},
+            "doc_2": {"name": "b.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}},
+        }), patch.object(svc, "_run_file_level_raptor", new_callable=AsyncMock) as mock_file, \
+           patch.object(svc, "_run_dataset_level_raptor", new_callable=AsyncMock) as mock_dataset:
+
+            mock_file.return_value = (sample_chunks, 42)
+
+            AsyncMock(return_value=(sample_chunks, 42, []))
+            with patch.object(RaptorService, "run_raptor_for_kb", new=AsyncMock(wraps=svc.run_raptor_for_kb)):
+                pass  # let's just call directly
+
+            # Direct call since we need to invoke the async method properly
+            import asyncio
+            loop = asyncio.new_event_loop()
+            try:
+                chunks, tk_count, cleanup = loop.run_until_complete(
+                    svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, vector_size, doc_ids)
+                )
+            finally:
+                loop.close()
+
+        mock_file.assert_called_once()
+        mock_dataset.assert_not_called()
+        assert chunks == sample_chunks
+        assert tk_count == 42
+
+    # ---- Basic dispatch (dataset-level scope) ----
+
+    def test_run_raptor_for_kb_dataset_scope_delegates_to_dataset_level(
+        self, mock_raptor_context, sample_chunks, raptor_config_dataset_scope
+    ):
+        """When scope='dataset', _run_dataset_level_raptor is called."""
+        svc = RaptorService(mock_raptor_context)
+        doc_ids = ["doc_1"]
+        chat_mdl = MagicMock()
+        embd_mdl = MagicMock()
+        vector_size = 128
+
+        with patch.object(svc, "_collect_doc_info", return_value={
+            "doc_1": {"name": "a.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}},
+        }), patch.object(svc, "_run_file_level_raptor", new_callable=AsyncMock) as mock_file, \
+           patch.object(svc, "_run_dataset_level_raptor", new_callable=AsyncMock) as mock_dataset:
+
+            mock_dataset.return_value = (sample_chunks, 99)
+
+            import asyncio
+            loop = asyncio.new_event_loop()
+            try:
+                chunks, tk_count, cleanup = loop.run_until_complete(
+                    svc.run_raptor_for_kb(raptor_config_dataset_scope, chat_mdl, embd_mdl, vector_size, doc_ids)
+                )
+            finally:
+                loop.close()
+
+        mock_dataset.assert_called_once()
+        mock_file.assert_not_called()
+        assert chunks == sample_chunks
+        assert tk_count == 99
+
+    # ---- Empty / no documents ----
+
+    def test_run_raptor_for_kb_empty_doc_ids(self, mock_raptor_context, raptor_config_file_scope):
+        """Empty doc_ids returns empty results."""
+        svc = RaptorService(mock_raptor_context)
+        chat_mdl = MagicMock()
+        embd_mdl = MagicMock()
+
+        with patch.object(svc, "_collect_doc_info", return_value={}), \
+           patch.object(svc, "_run_file_level_raptor", new_callable=AsyncMock) as mock_file, \
+           patch.object(svc, "_run_dataset_level_raptor", new_callable=AsyncMock):
+
+            mock_file.return_value = ([], 0)
+
+            import asyncio
+            loop = asyncio.new_event_loop()
+            try:
+                chunks, tk_count, cleanup = loop.run_until_complete(
+                    svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, 128, [])
+                )
+            finally:
+                loop.close()
+
+        assert chunks == []
+        assert tk_count == 0
+        assert cleanup == []
+
+    # ---- Cleanup scheduling through the public API ----
+
+    def test_run_raptor_for_kb_returns_cleanup_list(
+        self, mock_raptor_context, raptor_config_file_scope
+    ):
+        """Cleanup list from internal method is propagated to caller.
+
+        _run_file_level_raptor receives cleanup_raptor_chunks by reference (as
+        a positional arg) and may mutate it. This test verifies the public
+        method propagates whatever ends up in that list.
+        """
+        svc = RaptorService(mock_raptor_context)
+        doc_ids = ["doc_1"]
+        chat_mdl = MagicMock()
+        embd_mdl = MagicMock()
+
+        expected_cleanup = [("doc_1", "tree_builder_a")]
+
+        with patch.object(svc, "_collect_doc_info", return_value={
+            "doc_1": {"name": "a.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}},
+        }), patch.object(svc, "_run_file_level_raptor", new_callable=AsyncMock) as mock_file:
+
+            async def mock_run_file(*args, **kwargs):
+                # _run_file_level_raptor takes 12 positional args;
+                # cleanup_raptor_chunks is args[11] (0-indexed, last positional).
+                cleanup_list = args[11]
+                cleanup_list.append(("doc_1", "tree_builder_a"))
+                return [{"id": "c1"}], 10
+
+            mock_file.side_effect = mock_run_file
+
+            import asyncio
+            loop = asyncio.new_event_loop()
+            try:
+                chunks, tk_count, cleanup = loop.run_until_complete(
+                    svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, 128, doc_ids)
+                )
+            finally:
+                loop.close()
+
+        assert cleanup == expected_cleanup
+
+    # ---- Dispatch with missing raptor config key ----
+
+    def test_run_raptor_for_kb_defaults_to_file_scope_when_no_raptor_key(
+        self, mock_raptor_context
+    ):
+        """When kb_parser_config has no 'raptor' key, defaults to file scope."""
+        svc = RaptorService(mock_raptor_context)
+        doc_ids = ["doc_1"]
+        chat_mdl = MagicMock()
+        embd_mdl = MagicMock()
+        config = {}  # No raptor key at all
+
+        with patch.object(svc, "_collect_doc_info", return_value={
+            "doc_1": {"name": "a.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}},
+        }), patch.object(svc, "_run_file_level_raptor", new_callable=AsyncMock) as mock_file, \
+           patch.object(svc, "_run_dataset_level_raptor", new_callable=AsyncMock) as mock_dataset:
+
+            mock_file.return_value = ([], 0)
+
+            import asyncio
+            loop = asyncio.new_event_loop()
+            try:
+                loop.run_until_complete(
+                    svc.run_raptor_for_kb(config, chat_mdl, embd_mdl, 128, doc_ids)
+                )
+            finally:
+                loop.close()
+
+        mock_file.assert_called_once()
+        mock_dataset.assert_not_called()
+
+    # ---- Vector dimension name construction ----
+
+    def test_run_raptor_for_kb_passes_vector_size_to_file_level(
+        self, mock_raptor_context, sample_chunks, raptor_config_file_scope
+    ):
+        """Vector size is used to construct vctr_nm and passed to internal method."""
+        svc = RaptorService(mock_raptor_context)
+        doc_ids = ["doc_1"]
+        chat_mdl = MagicMock()
+        embd_mdl = MagicMock()
+        vector_size = 256
+
+        with patch.object(svc, "_collect_doc_info", return_value={
+            "doc_1": {"name": "a.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}},
+        }), patch.object(svc, "_run_file_level_raptor", new_callable=AsyncMock) as mock_file:
+
+            mock_file.return_value = (sample_chunks, 10)
+
+            import asyncio
+            loop = asyncio.new_event_loop()
+            try:
+                loop.run_until_complete(
+                    svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, vector_size, doc_ids)
+                )
+            finally:
+                loop.close()
+
+        # Verify _run_file_level_raptor received vctr_nm with the correct vector size
+        # Positional args: 0=raptor_config, 1=tree_builder, 2=clustering_method,
+        #   3=chat_mdl, 4=embd_mdl, 5=vctr_nm
+        positional_args = mock_file.call_args[0]
+        assert positional_args[5] == "q_256_vec"
+
+    # ---- Document info collection through public API ----
+
+    def test_run_raptor_for_kb_collects_doc_info(
+        self, mock_raptor_context, raptor_config_file_scope
+    ):
+        """Document info is collected before dispatching to internal methods."""
+        svc = RaptorService(mock_raptor_context)
+        doc_ids = ["doc_a"]
+        chat_mdl = MagicMock()
+        embd_mdl = MagicMock()
+
+        expected_info = {"doc_a": {"name": "file.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}}}
+
+        with patch.object(svc, "_collect_doc_info", return_value=expected_info) as mock_collect, \
+           patch.object(svc, "_run_file_level_raptor", new_callable=AsyncMock) as mock_file:
+
+            mock_file.return_value = ([], 0)
+
+            import asyncio
+            loop = asyncio.new_event_loop()
+            try:
+                loop.run_until_complete(
+                    svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, 128, doc_ids)
+                )
+            finally:
+                loop.close()
+
+        mock_collect.assert_called_once_with(doc_ids)
+        # Verify doc_info_by_id was passed as positional arg[7] to _run_file_level_raptor
+        positional_args = mock_file.call_args[0]
+        assert positional_args[7] == expected_info

test/unit_test/rag/svr/task_executor_refactor/test_recording_context.py

@@ -0,0 +1,357 @@
+#
+#  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.
+
+"""
+Unit tests for RecordingContext module.
+"""
+
+import time
+import pytest
+from rag.svr.task_executor_refactor.recording_context import (
+    RecordingContext,
+    get_recording_context,
+    set_recording_context,
+    recording_context_manager,
+    timed_with_recording,
+)
+
+
+class TestRecordingContextInit:
+    """Tests for RecordingContext initialization."""
+
+    def test_init_creates_empty_data(self):
+        """Test that __init__ creates empty _data dict."""
+        ctx = RecordingContext()
+        assert ctx._data == {}
+
+    def test_init_creates_empty_records(self):
+        """Test that __init__ creates empty records list."""
+        ctx = RecordingContext()
+        assert ctx.records == []
+
+
+class TestRecordingContextRecord:
+    """Tests for RecordingContext.record method."""
+
+    def test_record_single_value(self):
+        """Test recording a single value."""
+        ctx = RecordingContext()
+        ctx.record("chunk_count", 100)
+        assert ctx.get("chunk_count") == 100
+
+    def test_record_overwrites_existing_value(self):
+        """Test that recording with same key overwrites previous value."""
+        ctx = RecordingContext()
+        ctx.record("key", "value1")
+        ctx.record("key", "value2")
+        assert ctx.get("key") == "value2"
+
+    def test_record_none_value(self):
+        """Test recording None value."""
+        ctx = RecordingContext()
+        ctx.record("key", None)
+        assert ctx.get("key") is None
+
+    def test_record_complex_object(self):
+        """Test recording a complex object like list or dict."""
+        ctx = RecordingContext()
+        ctx.record("chunks", [{"id": 1}, {"id": 2}])
+        assert ctx.get("chunks") == [{"id": 1}, {"id": 2}]
+
+
+class TestRecordingContextFuncReturnValues:
+    """Tests for function return value recording."""
+
+    def test_save_func_return_value_first_call(self):
+        """Test saving first return value for a function."""
+        ctx = RecordingContext()
+        ctx.save_func_return_value("test_func", 42)
+        assert ctx.get_func_return_values("test_func") == [42]
+
+    def test_save_func_return_value_multiple_calls(self):
+        """Test saving multiple return values for same function."""
+        ctx = RecordingContext()
+        ctx.save_func_return_value("test_func", 1)
+        ctx.save_func_return_value("test_func", 2)
+        ctx.save_func_return_value("test_func", 3)
+        assert ctx.get_func_return_values("test_func") == [1, 2, 3]
+
+    def test_get_func_return_values_nonexistent_function(self):
+        """Test getting return values for nonexistent function returns empty list."""
+        ctx = RecordingContext()
+        assert ctx.get_func_return_values("nonexistent") == []
+
+    def test_get_func_return_values_multiple_functions(self):
+        """Test getting return values for different functions."""
+        ctx = RecordingContext()
+        ctx.save_func_return_value("func_a", "a1")
+        ctx.save_func_return_value("func_b", "b1")
+        ctx.save_func_return_value("func_a", "a2")
+        assert ctx.get_func_return_values("func_a") == ["a1", "a2"]
+        assert ctx.get_func_return_values("func_b") == ["b1"]
+
+
+class TestRecordingContextGet:
+    """Tests for RecordingContext.get method."""
+
+    def test_get_existing_key(self):
+        """Test getting an existing key."""
+        ctx = RecordingContext()
+        ctx.record("key", "value")
+        assert ctx.get("key") == "value"
+
+    def test_get_nonexistent_key_returns_none(self):
+        """Test getting nonexistent key returns None."""
+        ctx = RecordingContext()
+        assert ctx.get("missing") is None
+
+    def test_get_nonexistent_key_returns_default(self):
+        """Test getting nonexistent key returns provided default."""
+        ctx = RecordingContext()
+        assert ctx.get("missing", "default") == "default"
+
+    def test_get_with_none_default(self):
+        """Test getting with None as default."""
+        ctx = RecordingContext()
+        assert ctx.get("missing", None) is None
+
+
+class TestRecordingContextGetAllFuncReturnValues:
+    """Tests for get_all_func_return_values method."""
+
+    def test_get_all_func_return_values_empty(self):
+        """Test getting all values when none recorded."""
+        ctx = RecordingContext()
+        assert ctx.get_all_func_return_values() == {}
+
+    def test_get_all_func_return_values_with_data(self):
+        """Test getting all values with some data."""
+        ctx = RecordingContext()
+        ctx.save_func_return_value("func_a", 1)
+        ctx.save_func_return_value("func_b", 2)
+        result = ctx.get_all_func_return_values()
+        assert result == {"func_a": [1], "func_b": [2]}
+
+    def test_get_all_func_return_values_returns_copy(self):
+        """Test that returned dict is a copy, not the original."""
+        ctx = RecordingContext()
+        ctx.save_func_return_value("func", 1)
+        result = ctx.get_all_func_return_values()
+        result["func"] = []
+        # Original should be unchanged
+        assert ctx.get_func_return_values("func") == [1]
+
+
+class TestRecordingContextHas:
+    """Tests for RecordingContext.has method."""
+
+    def test_has_existing_key(self):
+        """Test has returns True for existing key."""
+        ctx = RecordingContext()
+        ctx.record("key", "value")
+        assert ctx.has("key") is True
+
+    def test_has_nonexistent_key(self):
+        """Test has returns False for nonexistent key."""
+        ctx = RecordingContext()
+        assert ctx.has("missing") is False
+
+    def test_has_after_clear(self):
+        """Test has returns False after clear."""
+        ctx = RecordingContext()
+        ctx.record("key", "value")
+        ctx.clear()
+        assert ctx.has("key") is False
+
+
+class TestRecordingContextClear:
+    """Tests for RecordingContext.clear method."""
+
+    def test_clear_removes_all_data(self):
+        """Test that clear removes all recorded data."""
+        ctx = RecordingContext()
+        ctx.record("key1", "value1")
+        ctx.record("key2", "value2")
+        ctx.clear()
+        assert ctx._data == {}
+
+    def test_clear_removes_all_records(self):
+        """Test that clear removes all timing records."""
+        ctx = RecordingContext()
+        with ctx.measure("op1"):
+            pass
+        ctx.clear()
+        assert ctx.records == []
+
+
+class TestRecordingContextMeasure:
+    """Tests for RecordingContext.measure context manager."""
+
+    def test_measure_records_elapsed_time(self):
+        """Test that measure records elapsed time."""
+        ctx = RecordingContext()
+        with ctx.measure("test_op"):
+            time.sleep(0.01)
+        assert len(ctx.records) == 1
+        assert ctx.records[0][0] == "test_op"
+        assert ctx.records[0][1] >= 0.01
+
+    def test_measure_multiple_operations(self):
+        """Test measuring multiple operations."""
+        ctx = RecordingContext()
+        with ctx.measure("op1"):
+            time.sleep(0.01)
+        with ctx.measure("op2"):
+            time.sleep(0.02)
+        assert len(ctx.records) == 2
+        assert ctx.records[0][0] == "op1"
+        assert ctx.records[1][0] == "op2"
+
+    def test_measure_preserves_context_on_exception(self):
+        """Test that measure still records time on exception."""
+        ctx = RecordingContext()
+        with pytest.raises(ValueError):
+            with ctx.measure("failing_op"):
+                raise ValueError("test error")
+        assert len(ctx.records) == 1
+        assert ctx.records[0][0] == "failing_op"
+
+
+class TestRecordingContextReset:
+    """Tests for RecordingContext.reset method."""
+
+    def test_reset_clears_data(self):
+        """Test that reset clears all data."""
+        ctx = RecordingContext()
+        ctx.record("key", "value")
+        ctx.reset()
+        assert ctx._data == {}
+
+    def test_reset_clears_records(self):
+        """Test that reset clears all records."""
+        ctx = RecordingContext()
+        with ctx.measure("op"):
+            pass
+        ctx.reset()
+        assert ctx.records == []
+
+
+class TestRecordingContextRepr:
+    """Tests for RecordingContext.__repr__ method."""
+
+    def test_repr_empty_context(self):
+        """Test repr of empty context."""
+        ctx = RecordingContext()
+        assert "RecordingContext" in repr(ctx)
+
+    def test_repr_with_data(self):
+        """Test repr with some data."""
+        ctx = RecordingContext()
+        ctx.record("key", "value")
+        r = repr(ctx)
+        assert "RecordingContext" in r
+        assert "key" in r
+
+
+class TestContextVariableFunctions:
+    """Tests for context variable functions."""
+
+    def test_set_and_get_recording_context(self):
+        """Test set and get recording context."""
+        ctx = RecordingContext()
+        set_recording_context(ctx)
+        assert get_recording_context() is ctx
+
+    def test_set_recording_context_none_unbinds(self):
+        """Test setting None unbinds the context."""
+        ctx = RecordingContext()
+        set_recording_context(ctx)
+        set_recording_context(None)
+        # After unbinding, get should raise RuntimeError
+        with pytest.raises(RuntimeError, match="no context"):
+            get_recording_context()
+
+
+class TestRecordingContextManager:
+    """Tests for recording_context_manager context manager."""
+
+    def test_context_manager_with_provided_context(self):
+        """Test context manager with provided context."""
+        ctx = RecordingContext()
+        with recording_context_manager(ctx) as mgr:
+            assert mgr is ctx
+            mgr.record("key", "value")
+        assert ctx.get("key") == "value"
+
+    def test_context_manager_creates_new_context(self):
+        """Test context manager creates new context when none provided."""
+        with recording_context_manager() as ctx:
+            assert isinstance(ctx, RecordingContext)
+            ctx.record("key", "value")
+            assert ctx.get("key") == "value"
+
+    def test_context_manager_restores_previous_context(self):
+        """Test context manager restores previous context after exit."""
+        outer_ctx = RecordingContext()
+        set_recording_context(outer_ctx)
+        
+        inner_ctx = RecordingContext()
+        with recording_context_manager(inner_ctx):
+            assert get_recording_context() is inner_ctx
+        
+        # After exiting, should restore outer_ctx
+        assert get_recording_context() is outer_ctx
+
+
+class TestTimedWithRecordingDecorator:
+    """Tests for timed_with_recording decorator."""
+
+    def test_decorator_without_parentheses(self):
+        """Test decorator used without parentheses."""
+        ctx = RecordingContext()
+        set_recording_context(ctx)
+        
+        @timed_with_recording
+        def test_func():
+            time.sleep(0.01)
+            return 42
+        
+        result = test_func()
+        assert result == 42
+
+    def test_decorator_with_parentheses_and_context(self):
+        """Test decorator with explicit context."""
+        ctx = RecordingContext()
+        
+        @timed_with_recording(recording_context=ctx)
+        def test_func():
+            time.sleep(0.01)
+            return "hello"
+        
+        result = test_func()
+        assert result == "hello"
+
+    def test_decorator_without_context_raises_error(self):
+        """Test decorator raises RuntimeError when no context is available."""
+        # Ensure no context is set
+        set_recording_context(None)
+        
+        @timed_with_recording
+        def test_func():
+            return 123
+        
+        # Should raise RuntimeError because no context is available
+        with pytest.raises(RuntimeError, match="no context"):
+            test_func()

test/unit_test/rag/svr/task_executor_refactor/test_task_context.py

@@ -0,0 +1,417 @@
+#
+#  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.
+
+"""
+Unit tests for TaskContext module.
+"""
+
+from unittest.mock import MagicMock
+from rag.svr.task_executor_refactor.task_context import TaskContext, TaskLimiters, TaskCallbacks
+
+
+def _make_ctx(task, **kwargs):
+    """Helper to create TaskContext with default limiters and callbacks."""
+    return TaskContext(
+        task=task,
+        limiters=kwargs.get("limiters", TaskLimiters()),
+        callbacks=kwargs.get("callbacks", TaskCallbacks()),
+        write_interceptor=kwargs.get("write_interceptor", None),
+    )
+
+
+class TestTaskContextInit:
+    """Tests for TaskContext initialization."""
+
+    def test_init_with_minimal_task(self):
+        """Test initialization with minimal task dict."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.id == "task_1"
+
+    def test_init_with_all_parameters(self):
+        """Test initialization with all parameters."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        chat_limiter = MagicMock()
+        minio_limiter = MagicMock()
+        chunk_limiter = MagicMock()
+        embed_limiter = MagicMock()
+        kg_limiter = MagicMock()
+        write_interceptor = MagicMock()
+        progress_callback = MagicMock()
+        has_canceled_func = MagicMock()
+
+        ctx = TaskContext(
+            task=task,
+            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,
+        )
+
+        assert ctx.chat_limiter is chat_limiter
+        assert ctx.minio_limiter is minio_limiter
+        assert ctx.chunk_limiter is chunk_limiter
+        assert ctx.embed_limiter is embed_limiter
+        assert ctx.kg_limiter is kg_limiter
+        assert ctx.write_interceptor is write_interceptor
+        assert ctx.callbacks.progress is progress_callback
+        assert ctx.has_canceled_func is has_canceled_func
+
+    def test_init_defaults_for_callbacks(self):
+        """Test that callbacks default to no-op functions."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        # Should not raise
+        ctx.callbacks.progress()
+        assert ctx.has_canceled_func("task_1") is False
+
+
+class TestTaskContextIdentityProperties:
+    """Tests for task identity properties."""
+
+    def test_id(self):
+        """Test id property."""
+        task = {"id": "task_123", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.id == "task_123"
+
+    def test_tenant_id(self):
+        """Test tenant_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.tenant_id == "tenant_1"
+
+    def test_kb_id_default(self):
+        """Test kb_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.kb_id == ""
+
+    def test_kb_id(self):
+        """Test kb_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "kb_id": "kb_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.kb_id == "kb_1"
+
+    def test_doc_id_default(self):
+        """Test doc_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.doc_id == ""
+
+    def test_doc_id(self):
+        """Test doc_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "doc_id": "doc_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.doc_id == "doc_1"
+
+    def test_doc_ids_default(self):
+        """Test doc_ids property defaults to empty list."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.doc_ids == []
+
+    def test_doc_ids(self):
+        """Test doc_ids property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "doc_ids": ["doc_1", "doc_2"]}
+        ctx = _make_ctx(task=task)
+        assert ctx.doc_ids == ["doc_1", "doc_2"]
+
+
+class TestTaskContextDocumentMetadataProperties:
+    """Tests for document metadata properties."""
+
+    def test_name_default(self):
+        """Test name property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.name == ""
+
+    def test_name(self):
+        """Test name property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "name": "test.pdf"}
+        ctx = _make_ctx(task=task)
+        assert ctx.name == "test.pdf"
+
+    def test_location_default(self):
+        """Test location property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.location == ""
+
+    def test_size_default(self):
+        """Test size property defaults to 0."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.size == 0
+
+    def test_size(self):
+        """Test size property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "size": 1024}
+        ctx = _make_ctx(task=task)
+        assert ctx.size == 1024
+
+
+class TestTaskContextParserProperties:
+    """Tests for parser configuration properties."""
+
+    def test_parser_id_default(self):
+        """Test parser_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.parser_id == ""
+
+    def test_parser_id(self):
+        """Test parser_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "parser_id": "naive"}
+        ctx = _make_ctx(task=task)
+        assert ctx.parser_id == "naive"
+
+    def test_parser_config_default(self):
+        """Test parser_config property defaults to empty dict."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.parser_config == {}
+
+    def test_parser_config(self):
+        """Test parser_config property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "parser_config": {"chunk_size": 512}}
+        ctx = _make_ctx(task=task)
+        assert ctx.parser_config == {"chunk_size": 512}
+
+    def test_kb_parser_config_default(self):
+        """Test kb_parser_config property defaults to empty dict."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.kb_parser_config == {}
+
+    def test_kb_parser_config(self):
+        """Test kb_parser_config property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "kb_parser_config": {"language": "en"}}
+        ctx = _make_ctx(task=task)
+        assert ctx.kb_parser_config == {"language": "en"}
+
+
+class TestTaskContextLanguageAndModelProperties:
+    """Tests for language and model properties."""
+
+    def test_language_default(self):
+        """Test language property defaults to 'en'."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.language == "en"
+
+    def test_language(self):
+        """Test language property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "language": "zh"}
+        ctx = _make_ctx(task=task)
+        assert ctx.language == "zh"
+
+    def test_llm_id_default(self):
+        """Test llm_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.llm_id == ""
+
+    def test_llm_id(self):
+        """Test llm_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "llm_id": "gpt-4"}
+        ctx = _make_ctx(task=task)
+        assert ctx.llm_id == "gpt-4"
+
+    def test_embd_id_default(self):
+        """Test embd_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.embd_id == ""
+
+    def test_embd_id(self):
+        """Test embd_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "embd_id": "text-embedding-ada-002"}
+        ctx = _make_ctx(task=task)
+        assert ctx.embd_id == "text-embedding-ada-002"
+
+
+class TestTaskContextPageRangeProperties:
+    """Tests for page range properties."""
+
+    def test_from_page_default(self):
+        """Test from_page property defaults to 0."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.from_page == 0
+
+    def test_from_page(self):
+        """Test from_page property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "from_page": 10}
+        ctx = _make_ctx(task=task)
+        assert ctx.from_page == 10
+
+    def test_to_page_default(self):
+        """Test to_page property defaults to -1."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.to_page == -1
+
+    def test_to_page(self):
+        """Test to_page property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "to_page": 100}
+        ctx = _make_ctx(task=task)
+        assert ctx.to_page == 100
+
+
+class TestTaskContextTaskTypeAndRoutingProperties:
+    """Tests for task type and routing properties."""
+
+    def test_task_type_default(self):
+        """Test task_type property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.task_type == ""
+
+    def test_task_type(self):
+        """Test task_type property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "task_type": "raptor"}
+        ctx = _make_ctx(task=task)
+        assert ctx.task_type == "raptor"
+
+    def test_dataflow_id_default(self):
+        """Test dataflow_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.dataflow_id == ""
+
+    def test_dataflow_id(self):
+        """Test dataflow_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "dataflow_id": "flow_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.dataflow_id == "flow_1"
+
+
+class TestTaskContextAdditionalProperties:
+    """Tests for additional properties."""
+
+    def test_pagerank_default(self):
+        """Test pagerank property defaults to 0."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.pagerank == 0
+
+    def test_pagerank(self):
+        """Test pagerank property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "pagerank": 10}
+        ctx = _make_ctx(task=task)
+        assert ctx.pagerank == 10
+
+    def test_file_default(self):
+        """Test file property defaults to None."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.file is None
+
+    def test_file(self):
+        """Test file property."""
+        file_obj = MagicMock()
+        task = {"id": "task_1", "tenant_id": "tenant_1", "file": file_obj}
+        ctx = _make_ctx(task=task)
+        assert ctx.file is file_obj
+
+
+class TestTaskContextMemoryProperties:
+    """Tests for memory task properties."""
+
+    def test_memory_id_default(self):
+        """Test memory_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.memory_id == ""
+
+    def test_memory_id(self):
+        """Test memory_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "memory_id": "mem_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.memory_id == "mem_1"
+
+    def test_source_id_default(self):
+        """Test source_id property defaults to empty string."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.source_id == ""
+
+    def test_source_id(self):
+        """Test source_id property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "source_id": "src_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.source_id == "src_1"
+
+    def test_message_dict_default(self):
+        """Test message_dict property defaults to empty dict."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.message_dict == {}
+
+    def test_message_dict(self):
+        """Test message_dict property."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "message_dict": {"key": "value"}}
+        ctx = _make_ctx(task=task)
+        assert ctx.message_dict == {"key": "value"}
+
+
+class TestTaskContextRawTask:
+    """Tests for raw_task property and get method."""
+
+    def test_raw_task_returns_original_dict(self):
+        """Test raw_task returns the original task dict."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "custom_key": "value"}
+        ctx = _make_ctx(task=task)
+        assert ctx.raw_task is task
+
+    def test_get_existing_key(self):
+        """Test get method with existing key."""
+        task = {"id": "task_1", "tenant_id": "tenant_1", "custom_key": "value"}
+        ctx = _make_ctx(task=task)
+        assert ctx.get("custom_key") == "value"
+
+    def test_get_nonexistent_key_returns_none(self):
+        """Test get method with nonexistent key returns None."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.get("missing") is None
+
+    def test_get_with_default(self):
+        """Test get method with default value."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        assert ctx.get("missing", "default") == "default"
+
+
+class TestTaskContextProgressCallback:
+    """Tests for progress callback functionality."""
+
+    def test_progress_cb_is_set_in_init(self):
+        """Test that _progress_cb is set during initialization."""
+        task = {"id": "task_1", "tenant_id": "tenant_1"}
+        ctx = _make_ctx(task=task)
+        # _progress_cb should be set in __init__
+        assert hasattr(ctx, '_progress_cb')
+        assert ctx._progress_cb is not None

test/unit_test/rag/svr/task_executor_refactor/test_task_handler.py

@@ -0,0 +1,300 @@
+#
+#  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.
+
+"""
+Unit tests for TaskHandler module.
+
+All orchestration tests validate behavior through the public handle()/handle_task()
+entry points.  Internal helpers (_run_standard_chunking, _run_dataflow, _run_raptor,
+_run_graphrag, _bind_embedding_model, _get_storage_binary, etc.) are exercised
+implicitly; no test reaches directly into those private orchestration methods.
+
+Stable pure helpers (_build_toc, _get_vector_size) are tested directly since they
+are side-effect-free data transformations.
+"""
+
+import pytest
+import numpy as np
+from unittest.mock import MagicMock, AsyncMock, patch
+
+from rag.svr.task_executor_refactor.task_handler import TaskHandler
+
+
+class TestTaskHandlerHandleTask:
+    """Tests for the public handle_task() entry point."""
+
+    @pytest.mark.asyncio
+    async def test_handle_task_calls_handle(self):
+        """Test handle_task delegates to handle()."""
+        ctx = MagicMock()
+        ctx.id = "task_1"
+        ctx.tenant_id = "tenant_1"
+        ctx.kb_id = "kb_1"
+        ctx.doc_id = "doc_1"
+        ctx.has_canceled_func = MagicMock(return_value=False)
+        handler = TaskHandler(ctx=ctx)
+        handler.handle = AsyncMock()
+        await handler.handle_task()
+        handler.handle.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_handle_task_cleanup_on_cancel(self):
+        """Test handle_task cleans up docStore when canceled."""
+        from common import settings
+        mock_doc_store = MagicMock()
+        mock_doc_store.index_exist = MagicMock(return_value=True)
+        mock_doc_store.delete = MagicMock(return_value=None)
+        orig = settings.docStoreConn
+        settings.docStoreConn = mock_doc_store
+        try:
+            ctx = MagicMock()
+            ctx.id = "task_1"
+            ctx.tenant_id = "tenant_1"
+            ctx.kb_id = "kb_1"
+            ctx.doc_id = "doc_1"
+            ctx.has_canceled_func = MagicMock(return_value=True)
+            ctx.recording_context = MagicMock()
+            handler = TaskHandler(ctx=ctx)
+            handler.handle = AsyncMock(side_effect=Exception("test error"))
+            # Should raise the exception
+            with pytest.raises(Exception, match="test error"):
+                await handler.handle_task()
+            mock_doc_store.delete.assert_called()
+        finally:
+            settings.docStoreConn = orig
+
+
+class TestTaskHandlerHandle:
+    """Tests for the public handle() method.
+
+    Internal orchestration methods (_run_standard_chunking, _run_dataflow,
+    _run_raptor, _run_graphrag, _bind_embedding_model) are exercised through
+    handle() so the suite stays resilient when those private methods change.
+    """
+
+    @pytest.mark.asyncio
+    async def test_handle_memory_task(self):
+        """Test handle dispatches memory tasks correctly."""
+        ctx = MagicMock()
+        ctx.task_type = "memory"
+        ctx.id = "task_1"
+        ctx.raw_task = {"memory_id": "mem_1"}
+        ctx.write_interceptor = None
+        ctx.has_canceled_func = MagicMock(return_value=False)
+
+        with patch("rag.svr.task_executor_refactor.task_handler.handle_save_to_memory_task", new_callable=AsyncMock) as mock_handle:
+            handler = TaskHandler(ctx=ctx)
+            handler._bind_embedding_model = AsyncMock()
+            handler._get_vector_size = MagicMock(return_value=1024)
+            handler._init_kb = MagicMock()
+            handler._run_standard_chunking = AsyncMock()
+            await handler.handle()
+            mock_handle.assert_called_once_with(ctx.raw_task)
+
+    @pytest.mark.asyncio
+    async def test_handle_dataflow_task(self):
+        """Test handle dispatches dataflow tasks."""
+        ctx = MagicMock()
+        ctx.task_type = "dataflow"
+        ctx.id = "task_1"
+        ctx.doc_id = "doc_1"
+        ctx.has_canceled_func = MagicMock(return_value=False)
+
+        handler = TaskHandler(ctx=ctx)
+        handler._run_dataflow = AsyncMock()
+        await handler.handle()
+        handler._run_dataflow.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_handle_canceled_task(self):
+        """Test handle returns early when task is canceled."""
+        ctx = MagicMock()
+        ctx.task_type = "standard"
+        ctx.id = "task_1"
+        ctx.has_canceled_func = MagicMock(return_value=True)
+        ctx.progress_cb = MagicMock()
+
+        handler = TaskHandler(ctx=ctx)
+        await handler.handle()
+        ctx.progress_cb.assert_called_once_with(-1, msg="Task has been canceled.")
+
+    @pytest.mark.asyncio
+    async def test_handle_standard_chunking(self):
+        """Test handle dispatches standard chunking end-to-end."""
+        ctx = MagicMock()
+        ctx.task_type = "standard"
+        ctx.id = "task_1"
+        ctx.tenant_id = "tenant_1"
+        ctx.kb_id = "kb_1"
+        ctx.doc_id = "doc_1"
+        ctx.embd_id = "embd_1"
+        ctx.language = "en"
+        ctx.parser_config = {}
+        ctx.has_canceled_func = MagicMock(return_value=False)
+        ctx.progress_cb = MagicMock()
+        ctx.recording_context = MagicMock()
+        ctx.name = "test.pdf"
+        ctx.from_page = 0
+        ctx.to_page = -1
+
+        handler = TaskHandler(ctx=ctx)
+        handler._bind_embedding_model = AsyncMock(return_value=MagicMock())
+        handler._get_vector_size = MagicMock(return_value=128)
+        handler._init_kb = MagicMock()
+        handler._run_standard_chunking = AsyncMock()
+
+        await handler.handle()
+        handler._run_standard_chunking.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_handle_raptor_task(self):
+        """Test handle dispatches raptor tasks."""
+        ctx = MagicMock()
+        ctx.task_type = "raptor"
+        ctx.id = "task_1"
+        ctx.tenant_id = "tenant_1"
+        ctx.kb_id = "kb_1"
+        ctx.embd_id = "embd_1"
+        ctx.language = "en"
+        ctx.has_canceled_func = MagicMock(return_value=False)
+        ctx.progress_cb = MagicMock()
+        ctx.recording_context = MagicMock()
+
+        handler = TaskHandler(ctx=ctx)
+        handler._bind_embedding_model = AsyncMock(return_value=MagicMock())
+        handler._get_vector_size = MagicMock(return_value=128)
+        handler._init_kb = MagicMock()
+        handler._run_raptor = AsyncMock()
+
+        await handler.handle()
+        handler._run_raptor.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_handle_graphrag_task(self):
+        """Test handle dispatches graphrag tasks."""
+        ctx = MagicMock()
+        ctx.task_type = "graphrag"
+        ctx.id = "task_1"
+        ctx.tenant_id = "tenant_1"
+        ctx.kb_id = "kb_1"
+        ctx.embd_id = "embd_1"
+        ctx.language = "en"
+        ctx.has_canceled_func = MagicMock(return_value=False)
+        ctx.progress_cb = MagicMock()
+        ctx.recording_context = MagicMock()
+
+        handler = TaskHandler(ctx=ctx)
+        handler._bind_embedding_model = AsyncMock(return_value=MagicMock())
+        handler._get_vector_size = MagicMock(return_value=128)
+        handler._init_kb = MagicMock()
+        handler._run_graphrag = AsyncMock()
+
+        await handler.handle()
+        handler._run_graphrag.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_handle_embedding_model_failure(self):
+        """Test handle returns early when embedding model binding fails."""
+        ctx = MagicMock()
+        ctx.task_type = "standard"
+        ctx.id = "task_1"
+        ctx.has_canceled_func = MagicMock(return_value=False)
+
+        handler = TaskHandler(ctx=ctx)
+        handler._bind_embedding_model = AsyncMock(return_value=None)
+
+        await handler.handle()
+        # Should not call _run_standard_chunking when model is None
+        assert not hasattr(handler, '_run_standard_chunking_called')
+
+
+class TestTaskHandlerGetVectorSize:
+    """Tests for _get_vector_size — stable pure helper."""
+
+    def test_get_vector_size(self):
+        mock_model = MagicMock()
+        mock_model.encode.return_value = (np.array([[1.0, 2.0, 3.0]]), 10)
+        result = TaskHandler._get_vector_size(mock_model)
+        assert result == 3
+
+
+class TestTaskHandlerBuildToc:
+    """Tests for _build_toc — stable pure helper (requires LLM mocking)."""
+
+    def test_build_toc_with_empty_docs(self):
+        """Test _build_toc returns None when run_toc_from_text returns empty."""
+        ctx = MagicMock()
+        ctx.tenant_id = "tenant_1"
+        ctx.llm_id = "llm_1"
+        ctx.language = "en"
+
+        docs = [{"id": "chunk_1", "content_with_weight": "text", "page_num_int": [1], "top_int": [0]}]
+
+        def mock_asyncio_run(coro):
+            # Close the coroutine to prevent "never awaited" warnings
+            coro.close()
+            return []
+
+        with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_cfg:
+            mock_cfg.return_value = MagicMock()
+            with patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle:
+                mock_msg = MagicMock()
+                mock_bundle.return_value.__enter__.return_value = mock_msg
+                with patch("rag.svr.task_executor_refactor.task_handler.asyncio.run", side_effect=mock_asyncio_run):
+                    result = TaskHandler._build_toc(ctx, docs, MagicMock())
+                    assert result is None
+
+    def test_build_toc_with_results(self):
+        """Test _build_toc builds TOC chunk when results exist."""
+        ctx = MagicMock()
+        ctx.tenant_id = "tenant_1"
+        ctx.llm_id = "llm_1"
+        ctx.language = "en"
+
+        docs = [{"id": "chunk_0", "content_with_weight": "text", "doc_id": "doc_1", "page_num_int": [1], "top_int": [0]}]
+        toc_result = [{"chunk_id": "0", "title": "Section 1"}]
+
+        def mock_asyncio_run(coro):
+            # Close the coroutine to prevent "never awaited" warnings
+            coro.close()
+            return toc_result
+
+        with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_cfg:
+            mock_cfg.return_value = MagicMock()
+            with patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle:
+                mock_msg = MagicMock()
+                mock_bundle.return_value.__enter__.return_value = mock_msg
+                with patch("rag.svr.task_executor_refactor.task_handler.asyncio.run", side_effect=mock_asyncio_run):
+                    result = TaskHandler._build_toc(ctx, docs, MagicMock())
+                    assert result is not None
+                    assert "toc_kwd" in result
+                    assert result["toc_kwd"] == "toc"
+                    assert result["available_int"] == 0
+
+
+class TestTaskHandlerInit:
+    """Tests for TaskHandler initialization."""
+
+    def test_init_stores_context_and_hook(self):
+        ctx = MagicMock()
+        hook = MagicMock()
+        handler = TaskHandler(ctx=ctx, billing_hook=hook)
+        assert handler._task_context is ctx
+        assert handler._billing_hook is hook
+
+    def test_init_default_hook_none(self):
+        ctx = MagicMock()
+        handler = TaskHandler(ctx=ctx)
+        assert handler._billing_hook is None

test/unit_test/rag/svr/task_executor_refactor/test_task_handler_integration.py

@@ -0,0 +1,993 @@
+#
+#  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.
+
+"""
+Integration tests for TaskHandler orchestration.
+"""
+
+import asyncio
+import gc
+import uuid
+from typing import Any, Dict
+from unittest.mock import MagicMock, AsyncMock, patch
+
+import pytest
+
+from rag.svr.task_executor_refactor.task_handler import TaskHandler
+from rag.svr.task_executor_refactor.task_context import TaskContext, TaskLimiters, TaskCallbacks
+from rag.svr.task_executor_refactor.recording_context import BaseRecordingContext, RecordingContext
+from rag.svr.task_executor_refactor.constants import CANVAS_DEBUG_DOC_ID, GRAPH_RAPTOR_FAKE_DOC_ID
+
+# Import shared helpers from conftest
+from test.unit_test.rag.svr.task_executor_refactor.conftest import (
+    AsyncMockLimiter,
+    create_mock_embedding_model,
+    create_default_chunks,
+    create_mock_settings,
+    create_mock_chunk_service,
+)
+
+
+def create_task_context(
+    task_dict: Dict[str, Any],
+    is_canceled: bool = False,
+    recording_context: BaseRecordingContext | None = None,
+) -> TaskContext:
+    """Create a real TaskContext with mocked limiters and callbacks.
+
+    Args:
+        task_dict: Task dictionary with all task attributes.
+        is_canceled: If True, has_canceled_func returns True.
+        recording_context: RecordingContext to inject. If None, a new one
+            is created automatically so that recording_context access works.
+
+    Returns:
+        TaskContext with all required dependencies injected.
+    """
+    if recording_context is None:
+        recording_context = RecordingContext()
+    limiter = AsyncMockLimiter()
+    progress_callback = MagicMock()
+    ctx = TaskContext(
+        task=task_dict,
+        limiters=TaskLimiters(
+            chat=limiter,
+            minio=limiter,
+            chunk=limiter,
+            embed=limiter,
+            kg=limiter,
+        ),
+        callbacks=TaskCallbacks(
+            progress=progress_callback,
+            has_canceled=MagicMock(return_value=is_canceled),
+        ),
+        recording_context=recording_context,
+    )
+    # Add progress_callback property for task_handler compatibility
+    ctx.progress_callback = progress_callback
+    # Add set_progress_cb method for task_handler compatibility
+    ctx.set_progress_cb = lambda cb: setattr(ctx.callbacks, 'progress_cb', cb)
+    return ctx
+
+
+# Common patcher for _get_storage_binary since it imports settings internally
+def patch_get_storage_binary():
+    return patch.object(TaskHandler, '_get_storage_binary', new_callable=AsyncMock, return_value=b"fake pdf binary")
+
+
+def patch_task_handler_settings(mock_settings):
+    """Patch the settings module-level import in task_handler."""
+    return patch("rag.svr.task_executor_refactor.task_handler.settings", mock_settings)
+
+
+class TestStandardChunkingPipelineIntegration:
+    """P0: Integration tests for the complete standard chunking pipeline."""
+
+    def _create_standard_task_dict(self) -> Dict[str, Any]:
+        return {
+            "id": f"task_{uuid.uuid4().hex[:8]}",
+            "tenant_id": "tenant_test",
+            "kb_id": "kb_test",
+            "doc_id": "doc_test",
+            "name": "test_document.pdf",
+            "location": "/path/to/test_document.pdf",
+            "size": 1024,
+            "parser_id": "naive",
+            "parser_config": {
+                "auto_keywords": 0,
+                "auto_questions": 0,
+                "enable_metadata": False,
+            },
+            "kb_parser_config": {},
+            "language": "en",
+            "llm_id": "llm_test",
+            "embd_id": "embd_test",
+            "from_page": 0,
+            "to_page": -1,
+            "task_type": "standard",
+            "pagerank": 0,
+        }
+
+    @pytest.mark.asyncio
+    async def test_full_chunking_pipeline_records_task_status(self):
+        """Verify that the complete pipeline records task_status as 'completed'."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+
+            async def mock_thread_impl(func, *args, **kwargs):
+                return b"fake pdf binary"
+
+            mock_thread_exec.side_effect = mock_thread_impl
+            mock_chunk_thread_exec.side_effect = mock_thread_impl
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            recording_ctx = ctx.recording_context
+            task_status = recording_ctx.get("task_status")
+            assert task_status == "completed", f"Expected task_status='completed', got {task_status}"
+
+    @pytest.mark.asyncio
+    async def test_full_chunking_pipeline_records_insertion_result(self):
+        """Verify that insertion_result is recorded as 'success'."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+
+            async def mock_thread_impl(func, *args, **kwargs):
+                return b"fake pdf binary"
+
+            mock_thread_exec.side_effect = mock_thread_impl
+            mock_chunk_thread_exec.side_effect = mock_thread_impl
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            recording_ctx = ctx.recording_context
+            insertion_result = recording_ctx.get("insertion_result")
+            assert insertion_result == "success", f"Expected insertion_result='success', got {insertion_result}"
+
+    @pytest.mark.asyncio
+    async def test_full_chunking_pipeline_records_chunk_ids(self):
+        """Verify that chunk_ids_count is recorded after build_chunks."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunks = create_default_chunks(count=3)
+        mock_chunk_service = create_mock_chunk_service(chunks=mock_chunks)
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.run_toc_from_text", new_callable=AsyncMock) as mock_run_toc, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+            mock_run_toc.return_value = []  # TOC returns empty when not enabled
+
+            async def mock_thread_impl(func, *args, **kwargs):
+                return b"fake pdf binary"
+
+            mock_thread_exec.side_effect = mock_thread_impl
+            mock_chunk_thread_exec.side_effect = mock_thread_impl
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            recording_ctx = ctx.recording_context
+            chunk_ids_count = recording_ctx.get("chunk_ids_count")
+            assert chunk_ids_count is not None, "chunk_ids_count should be recorded"
+            assert chunk_ids_count == 3, f"Expected chunk_ids_count=3, got {chunk_ids_count}"
+
+    @pytest.mark.asyncio
+    async def test_full_chunking_pipeline_records_token_count(self):
+        """Verify that token_count and vector_size are recorded after embedding."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+
+            async def mock_thread_impl(func, *args, **kwargs):
+                return b"fake pdf binary"
+
+            mock_thread_exec.side_effect = mock_thread_impl
+            mock_chunk_thread_exec.side_effect = mock_thread_impl
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            recording_ctx = ctx.recording_context
+            token_count = recording_ctx.get("token_count")
+            vector_size = recording_ctx.get("vector_size")
+
+            assert token_count is not None, "token_count should be recorded"
+            assert vector_size is not None, "vector_size should be recorded"
+            assert vector_size == 128, f"Expected vector_size=128, got {vector_size}"
+
+    @pytest.mark.asyncio
+    async def test_full_chunking_pipeline_progress_callback_invoked(self):
+        """Verify that progress_callback is invoked multiple times during pipeline."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+
+            async def mock_thread_impl(func, *args, **kwargs):
+                return b"fake pdf binary"
+
+            mock_thread_exec.side_effect = mock_thread_impl
+            mock_chunk_thread_exec.side_effect = mock_thread_impl
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            ctx.progress_callback.assert_called()
+            call_count = ctx.progress_callback.call_count
+            assert call_count > 0, "progress_callback should have been invoked at least once"
+
+
+class TestTaskCancellationCleanupIntegration:
+    """P0: Integration tests for task cancellation cleanup flow."""
+
+    def _create_standard_task_dict(self) -> Dict[str, Any]:
+        return {
+            "id": f"task_{uuid.uuid4().hex[:8]}",
+            "tenant_id": "tenant_test",
+            "kb_id": "kb_test",
+            "doc_id": "doc_test",
+            "name": "test_document.pdf",
+            "location": "/path/to/test_document.pdf",
+            "size": 1024,
+            "parser_id": "naive",
+            "parser_config": {},
+            "kb_parser_config": {},
+            "language": "en",
+            "llm_id": "llm_test",
+            "embd_id": "embd_test",
+            "from_page": 0,
+            "to_page": -1,
+            "task_type": "standard",
+            "pagerank": 0,
+        }
+
+    @pytest.mark.asyncio
+    async def test_canceled_task_calls_docstore_delete(self):
+        """Verify that docStoreConn.delete is called when task is canceled."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict, is_canceled=True)
+        mock_settings = create_mock_settings()
+
+        call_log = []
+
+        def mock_thread_impl(func, *args, **kwargs):
+            # Get the actual method name from the mock
+            func_repr = repr(func)
+            call_log.append(func_repr)
+            if 'index_exist' in func_repr:
+                return True
+            if 'delete' in func_repr:
+                return {"result": "deleted"}
+            return {"result": "deleted"}
+
+        with patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_index"), \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec", side_effect=mock_thread_impl):
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle_task()
+
+            # Verify delete was called by checking the call log
+            delete_calls = [c for c in call_log if 'delete' in c]
+            assert len(delete_calls) >= 1, f"Expected at least one delete call, got: {call_log}"
+
+    @pytest.mark.asyncio
+    async def test_canceled_task_progress_callback_with_negative_one(self):
+        """Verify that progress_callback is called with -1 when task is canceled."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict, is_canceled=True)
+        mock_settings = create_mock_settings()
+
+        def mock_thread_impl(func, *args, **kwargs):
+            func_repr = repr(func)
+            if 'index_exist' in func_repr:
+                return True
+            if 'delete' in func_repr:
+                return {"result": "deleted"}
+            return {"result": "deleted"}
+
+        with patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_index"), \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec", side_effect=mock_thread_impl):
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle_task()
+
+            ctx.progress_callback.assert_called()
+            call_args_list = ctx.progress_callback.call_args_list
+            # Check for -1 in any position of the call arguments
+            has_negative_progress = False
+            for call in call_args_list:
+                # Check positional args
+                for arg in call[0]:
+                    if arg == -1:
+                        has_negative_progress = True
+                        break
+                # Check keyword args
+                if call[1].get("prog") == -1:
+                    has_negative_progress = True
+                if has_negative_progress:
+                    break
+            assert has_negative_progress, f"progress_callback should have been called with -1 progress. Calls: {call_args_list}"
+
+    @pytest.mark.asyncio
+    async def test_canceled_task_does_not_proceed_to_chunking(self):
+        """Verify that canceled task does not proceed to embedding model binding."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict, is_canceled=True)
+        mock_settings = create_mock_settings()
+
+        with patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default:
+
+            mock_index_name.return_value = "test_index"
+            mock_settings.docStoreConn.index_exist.return_value = True
+            mock_settings.docStoreConn.delete.return_value = {"result": "deleted"}
+
+            async def mock_thread_impl(func, *args, **kwargs):
+                return {"result": "deleted"}
+
+            mock_thread_exec.side_effect = mock_thread_impl
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle_task()
+
+            mock_bundle.assert_not_called()
+
+
+class TestRaptorPipelineIntegration:
+    """P1: Integration tests for the RAPTOR pipeline."""
+
+    def _create_raptor_task_dict(self) -> Dict[str, Any]:
+        return {
+            "id": f"task_{uuid.uuid4().hex[:8]}",
+            "tenant_id": "tenant_test",
+            "kb_id": "kb_test",
+            "doc_id": GRAPH_RAPTOR_FAKE_DOC_ID,
+            "doc_ids": ["doc1", "doc2"],
+            "name": "raptor_task",
+            "parser_id": "naive",
+            "parser_config": {"raptor": {"use_raptor": False}},
+            "kb_parser_config": {"raptor": {"use_raptor": False}},
+            "language": "en",
+            "llm_id": "llm_test",
+            "embd_id": "embd_test",
+            "from_page": 0,
+            "to_page": -1,
+            "task_type": "raptor",
+            "pagerank": 0,
+        }
+
+    @pytest.mark.asyncio
+    async def test_raptor_pipeline_records_task_status(self):
+        """Verify that RAPTOR pipeline records task_status."""
+        task_dict = self._create_raptor_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_kb = MagicMock()
+        mock_kb.id = "kb_test"
+        mock_kb.parser_config = {"raptor": {"use_raptor": False}}
+
+        with patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.KnowledgebaseService") as mock_kb_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.RaptorService") as mock_raptor_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_index_name.return_value = "test_index"
+            mock_kb_service.get_by_id.return_value = (True, mock_kb)
+            mock_kb_service.update_by_id.return_value = True
+            mock_raptor_service.return_value.run_raptor_for_kb = AsyncMock(return_value=([], 0, []))
+            mock_chunk_service.return_value.insert_chunks = AsyncMock(return_value=True)
+            mock_doc_service.increment_chunk_num = MagicMock()
+
+            async def mock_thread_impl(func, *args, **kwargs):
+                return None
+
+            mock_thread_exec.side_effect = mock_thread_impl
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            recording_ctx = ctx.recording_context
+            task_status = recording_ctx.get("task_status")
+            assert task_status == "completed", f"Expected task_status='completed', got {task_status}"
+
+    @pytest.mark.asyncio
+    async def test_raptor_pipeline_enables_raptor_if_not_configured(self):
+        """Verify that RAPTOR is enabled if not already configured."""
+        task_dict = self._create_raptor_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_kb = MagicMock()
+        mock_kb.id = "kb_test"
+        mock_kb.parser_config = {"raptor": {"use_raptor": False}}
+
+        with patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.KnowledgebaseService") as mock_kb_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.RaptorService") as mock_raptor_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_index_name.return_value = "test_index"
+            mock_kb_service.get_by_id.return_value = (True, mock_kb)
+            mock_kb_service.update_by_id.return_value = True
+            mock_raptor_service.return_value.run_raptor_for_kb = AsyncMock(return_value=([], 0, []))
+            mock_chunk_service.return_value.insert_chunks = AsyncMock(return_value=True)
+            mock_doc_service.increment_chunk_num = MagicMock()
+
+            async def mock_thread_impl(func, *args, **kwargs):
+                return None
+
+            mock_thread_exec.side_effect = mock_thread_impl
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            # Check that the kb parser_config was updated
+            mock_kb_service.update_by_id.assert_called_once()
+            call_args = mock_kb_service.update_by_id.call_args
+            update_dict = call_args[0][1]
+            assert update_dict.get("parser_config", {}).get("raptor", {}).get("use_raptor") is True, \
+                "RAPTOR should be enabled in parser_config after running"
+
+
+class TestEmbeddingModelBindingFailureIntegration:
+    """P1: Integration tests for embedding model binding failure."""
+
+    def _create_standard_task_dict(self) -> Dict[str, Any]:
+        return {
+            "id": f"task_{uuid.uuid4().hex[:8]}",
+            "tenant_id": "tenant_test",
+            "kb_id": "kb_test",
+            "doc_id": "doc_test",
+            "name": "test_document.pdf",
+            "location": "/path/to/test_document.pdf",
+            "size": 1024,
+            "parser_id": "naive",
+            "parser_config": {},
+            "kb_parser_config": {},
+            "language": "en",
+            "llm_id": "llm_test",
+            "embd_id": "embd_test",
+            "from_page": 0,
+            "to_page": -1,
+            "task_type": "standard",
+            "pagerank": 0,
+        }
+
+    @pytest.mark.asyncio
+    async def test_embedding_binding_failure_raises_exception(self):
+        """Verify that embedding model binding failure raises an exception."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict)
+
+        with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default:
+
+            mock_get_config.side_effect = Exception("Model not found")
+            mock_get_default.side_effect = Exception("Model not found")
+
+            handler = TaskHandler(ctx=ctx)
+
+            with pytest.raises(Exception, match="Model not found"):
+                await handler.handle()
+
+    @pytest.mark.asyncio
+    async def test_embedding_binding_failure_calls_progress_callback(self):
+        """Verify that embedding model binding failure calls progress_callback."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict)
+
+        with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default:
+
+            mock_get_config.side_effect = Exception("Model not found")
+            mock_get_default.side_effect = Exception("Model not found")
+
+            handler = TaskHandler(ctx=ctx)
+
+            with pytest.raises(Exception):
+                await handler.handle()
+
+            ctx.progress_callback.assert_called()
+
+
+class TestDataflowPipelineIntegration:
+    """P2: Integration tests for the dataflow pipeline."""
+
+    def _create_dataflow_task_dict(self) -> Dict[str, Any]:
+        return {
+            "id": f"task_{uuid.uuid4().hex[:8]}",
+            "tenant_id": "tenant_test",
+            "kb_id": "kb_test",
+            "doc_id": CANVAS_DEBUG_DOC_ID,
+            "name": "dataflow_debug",
+            "parser_id": "naive",
+            "parser_config": {},
+            "kb_parser_config": {},
+            "language": "en",
+            "llm_id": "llm_test",
+            "embd_id": "embd_test",
+            "from_page": 0,
+            "to_page": -1,
+            "task_type": "dataflow",
+            "pagerank": 0,
+        }
+
+    @pytest.mark.asyncio
+    async def test_dataflow_pipeline_calls_dataflow_service(self):
+        """Verify that dataflow pipeline calls DataflowService.run_dataflow()."""
+        task_dict = self._create_dataflow_task_dict()
+        ctx = create_task_context(task_dict)
+
+        with patch("rag.svr.task_executor_refactor.task_handler.DataflowService") as mock_dataflow_service:
+            mock_instance = MagicMock()
+            mock_instance.run_dataflow = AsyncMock(return_value=None)
+            mock_dataflow_service.return_value = mock_instance
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            mock_dataflow_service.assert_called_once()
+            mock_instance.run_dataflow.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_dataflow_debug_mode_calls_dataflow_service(self):
+        """Verify that dataflow debug mode also calls DataflowService."""
+        task_dict = self._create_dataflow_task_dict()
+        ctx = create_task_context(task_dict)
+
+        with patch("rag.svr.task_executor_refactor.task_handler.DataflowService") as mock_dataflow_service:
+            mock_instance = MagicMock()
+            mock_instance.run_dataflow = AsyncMock(return_value=None)
+            mock_dataflow_service.return_value = mock_instance
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            mock_dataflow_service.assert_called_once()
+            mock_instance.run_dataflow.assert_called_once()
+
+
+class TestTocAsyncFlowIntegration:
+    """P2: Integration tests for TOC async flow."""
+
+    def _create_toc_enabled_task_dict(self) -> Dict[str, Any]:
+        return {
+            "id": f"task_{uuid.uuid4().hex[:8]}",
+            "tenant_id": "tenant_test",
+            "kb_id": "kb_test",
+            "doc_id": "doc_test",
+            "name": "test_document.pdf",
+            "location": "/path/to/test_document.pdf",
+            "size": 1024,
+            "parser_id": "naive",
+            "parser_config": {
+                "auto_keywords": 0,
+                "auto_questions": 0,
+                "enable_metadata": False,
+                "toc_extraction": True,
+            },
+            "kb_parser_config": {},
+            "language": "en",
+            "llm_id": "llm_test",
+            "embd_id": "embd_test",
+            "from_page": 0,
+            "to_page": -1,
+            "task_type": "standard",
+            "pagerank": 0,
+        }
+
+    @pytest.mark.asyncio
+    async def test_toc_async_flow_creates_toc_thread(self):
+        """Verify that TOC async flow creates a TOC thread when enabled."""
+
+        task_dict = self._create_toc_enabled_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.run_toc_from_text", new_callable=AsyncMock) as mock_run_toc, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls, \
+             patch("rag.svr.task_executor_refactor.post_processor.DocumentService") as mock_post_doc_service:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+            mock_run_toc.return_value = [{"title": "Test TOC", "level": 1}]
+            mock_post_doc_service.increment_chunk_num = MagicMock()
+
+            async def mock_thread_impl(func, *args, **kwargs):
+                return b"fake pdf binary"
+
+            mock_thread_exec.side_effect = mock_thread_impl
+            mock_chunk_thread_exec.side_effect = mock_thread_impl
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            mock_run_toc.assert_called()
+
+        # Explicit cleanup to prevent resource leaks
+        del mock_embedding, mock_settings, mock_chunk_service
+        del mock_get_config, mock_get_default, mock_bundle, mock_file_service
+        del mock_index_name, mock_doc_service, mock_chunk_service_cls, mock_run_toc, mock_post_doc_service
+        del mock_thread_exec, mock_chunk_thread_exec
+        # Allow pending callbacks to execute
+        await asyncio.sleep(0)
+        gc.collect()
+
+    @pytest.mark.asyncio(loop_scope="function")
+    @pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning")
+    async def test_toc_async_flow_does_not_create_thread_when_disabled(self):
+        """Verify that TOC async flow does not create a thread when disabled.
+        
+        Note: This test has a known issue with resource leaks (unclosed sockets and
+        event loops) when run as part of the full test suite. The warning filter
+        above suppresses these warnings temporarily. The root cause is related to
+        asyncio.to_thread creating new event loops that are not properly cleaned up
+        by pytest-asyncio.
+        """
+
+        task_dict = self._create_toc_enabled_task_dict()
+        task_dict["parser_config"]["toc_extraction"] = False
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.run_toc_from_text", new_callable=AsyncMock) as mock_run_toc, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+
+            async def mock_thread_impl(func, *args, **kwargs):
+                return b"fake pdf binary"
+
+            mock_thread_exec.side_effect = mock_thread_impl
+            mock_chunk_thread_exec.side_effect = mock_thread_impl
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            mock_run_toc.assert_not_called()
+
+        # Explicit cleanup to prevent resource leaks
+        del mock_embedding, mock_settings, mock_chunk_service
+        del mock_get_config, mock_get_default, mock_bundle, mock_file_service
+        del mock_index_name, mock_doc_service, mock_chunk_service_cls, mock_run_toc
+        del mock_thread_exec, mock_chunk_thread_exec
+        # Allow pending callbacks to execute and close event loop
+        await asyncio.sleep(0)
+        # Cancel all pending tasks
+        current_task = asyncio.current_task()
+        pending = [t for t in asyncio.all_tasks() if t is not current_task and not t.done()]
+        for task in pending:
+            task.cancel()
+        if pending:
+            await asyncio.gather(*pending, return_exceptions=True)
+        gc.collect()
+
+
+class TestRecordingContextDataFlowAssertions:
+    """P2: Integration tests for RecordingContext data flow assertions."""
+
+    def _create_standard_task_dict(self) -> Dict[str, Any]:
+        return {
+            "id": f"task_{uuid.uuid4().hex[:8]}",
+            "tenant_id": "tenant_test",
+            "kb_id": "kb_test",
+            "doc_id": "doc_test",
+            "name": "test_document.pdf",
+            "location": "/path/to/test_document.pdf",
+            "size": 1024,
+            "parser_id": "naive",
+            "parser_config": {
+                "auto_keywords": 0,
+                "auto_questions": 0,
+                "enable_metadata": False,
+            },
+            "kb_parser_config": {},
+            "language": "en",
+            "llm_id": "llm_test",
+            "embd_id": "embd_test",
+            "from_page": 0,
+            "to_page": -1,
+            "task_type": "standard",
+            "pagerank": 0,
+        }
+
+    @pytest.mark.asyncio
+    async def test_recording_context_captures_file_size_check(self):
+        """Verify that RecordingContext captures file_size_exceeded result."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+
+            async def mock_thread_impl(func, *args, **kwargs):
+                return b"fake pdf binary"
+
+            mock_thread_exec.side_effect = mock_thread_impl
+            mock_chunk_thread_exec.side_effect = mock_thread_impl
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            recording_ctx = ctx.recording_context
+            file_size_exceeded = recording_ctx.get("file_size_exceeded")
+            assert file_size_exceeded is None or file_size_exceeded is False, \
+                f"Expected file_size_exceeded to be False/None for small file, got {file_size_exceeded}"
+
+    @pytest.mark.asyncio
+    async def test_recording_context_captures_parser_id(self):
+        """Verify that RecordingContext captures parser_id from task context."""
+        task_dict = self._create_standard_task_dict()
+        ctx = create_task_context(task_dict)
+        mock_embedding = create_mock_embedding_model(vector_size=128)
+        mock_settings = create_mock_settings()
+        mock_chunk_service = create_mock_chunk_service()
+
+        with patch_get_storage_binary(), \
+             patch_task_handler_settings(mock_settings), \
+             patch("rag.svr.task_executor_refactor.chunk_service.settings", mock_settings), \
+             patch("rag.svr.task_executor_refactor.task_handler.get_model_config_by_type_and_name") as mock_get_config, \
+             patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \
+             patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_get_default, \
+             patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService") as mock_file_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.thread_pool_exec") as mock_thread_exec, \
+             patch("rag.svr.task_executor_refactor.chunk_service.thread_pool_exec") as mock_chunk_thread_exec, \
+             patch("rag.svr.task_executor_refactor.task_handler.DocumentService") as mock_doc_service, \
+             patch("rag.svr.task_executor_refactor.task_handler.search.index_name") as mock_index_name, \
+             patch("rag.svr.task_executor_refactor.task_handler.ChunkService") as mock_chunk_service_cls:
+
+            mock_get_config.return_value = MagicMock()
+            mock_get_default.return_value = MagicMock()
+            mock_bundle.return_value = mock_embedding
+            mock_file_service.get_storage_address.return_value = ("bucket_test", "name_test")
+            mock_index_name.return_value = "test_index"
+            mock_doc_service.increment_chunk_num = MagicMock()
+            mock_doc_service.get_document_metadata.return_value = {}
+            mock_doc_service.update_document_metadata = MagicMock()
+            mock_chunk_service_cls.return_value = mock_chunk_service
+
+            async def mock_thread_impl(func, *args, **kwargs):
+                return b"fake pdf binary"
+
+            mock_thread_exec.side_effect = mock_thread_impl
+            mock_chunk_thread_exec.side_effect = mock_thread_impl
+
+            handler = TaskHandler(ctx=ctx)
+            await handler.handle()
+
+            recording_ctx = ctx.recording_context
+            # parser_id is available in the task context, verify task completion
+            task_status = recording_ctx.get("task_status")
+            assert task_status == "completed", f"Expected task_status='completed', got {task_status}"
+            # Verify the parser_id is accessible from the task context
+            assert ctx.parser_id == "naive", f"Expected parser_id='naive', got {ctx.parser_id}"

test/unit_test/rag/svr/task_executor_refactor/test_task_raptor_utils.py

@@ -0,0 +1,219 @@
+#
+#  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.
+
+"""
+Unit tests for rag/svr/task_executor_refactor/raptor_utils.py module.
+"""
+
+import pytest
+from unittest.mock import MagicMock, patch
+from rag.svr.task_executor_refactor.raptor_utils import (
+    get_raptor_chunk_field_map,
+    delete_raptor_chunks,
+)
+
+
+class TestGetRaptorChunkFieldMap:
+    """Tests for get_raptor_chunk_field_map function."""
+
+    @pytest.mark.asyncio
+    async def test_returns_primary_result_when_raptor_chunks_exist(self):
+        """Test that primary result is returned when RAPTOR chunks exist."""
+        from common import settings
+        original_retriever = settings.docStoreConn
+        
+        mock_doc_store = MagicMock()
+        mock_doc_store.search.return_value = {"chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}}}
+        mock_doc_store.get_fields.return_value = {"chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}}}
+        settings.docStoreConn = mock_doc_store
+
+        try:
+            with patch("rag.svr.task_executor_refactor.raptor_utils.thread_pool_exec") as mock_thread:
+                async def mock_exec(*args, **kwargs):
+                    return {"chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}}}
+                mock_thread.side_effect = mock_exec
+
+                with patch("rag.svr.task_executor_refactor.raptor_utils.collect_raptor_chunk_ids") as mock_collect:
+                    mock_collect.return_value = {"chunk_1"}
+
+                    result = await get_raptor_chunk_field_map("doc_1", "tenant_1", "kb_1")
+
+                    assert "chunk_1" in result
+        finally:
+            settings.docStoreConn = original_retriever
+
+    @pytest.mark.asyncio
+    async def test_falls_back_to_secondary_search_when_no_raptor_chunks(self):
+        """Test that fallback search is used when no RAPTOR chunks found."""
+        from common import settings
+        original_retriever = settings.docStoreConn
+        
+        mock_doc_store = MagicMock()
+        settings.docStoreConn = mock_doc_store
+
+        try:
+            call_count = 0
+            async def mock_exec(*args, **kwargs):
+                nonlocal call_count
+                call_count += 1
+                if call_count == 1:
+                    return {}  # Primary returns empty
+                else:
+                    return {"chunk_1": {"raptor_kwd": "raptor"}}  # Fallback
+
+            with patch("rag.svr.task_executor_refactor.raptor_utils.thread_pool_exec") as mock_thread:
+                mock_thread.side_effect = mock_exec
+
+                with patch("rag.svr.task_executor_refactor.raptor_utils.collect_raptor_chunk_ids") as mock_collect:
+                    mock_collect.return_value = set()  # Primary has no RAPTOR chunks
+
+                    _ = await get_raptor_chunk_field_map("doc_1", "tenant_1", "kb_1")
+
+                    # Should have called thread_pool_exec twice (primary + fallback)
+                    assert mock_thread.call_count == 2
+        finally:
+            settings.docStoreConn = original_retriever
+
+    @pytest.mark.asyncio
+    async def test_handles_fallback_search_exception(self):
+        """Test that exception in fallback search is handled gracefully."""
+        from common import settings
+        original_retriever = settings.docStoreConn
+        
+        mock_doc_store = MagicMock()
+        mock_doc_store.get_fields.return_value = {}
+        settings.docStoreConn = mock_doc_store
+
+        try:
+            call_count = 0
+            async def mock_exec(*args, **kwargs):
+                nonlocal call_count
+                call_count += 1
+                if call_count == 1:
+                    return {}  # Primary returns empty
+                else:
+                    raise Exception("Fallback search failed")  # Fallback will raise exception
+
+            with patch("rag.svr.task_executor_refactor.raptor_utils.thread_pool_exec") as mock_thread:
+                mock_thread.side_effect = mock_exec
+
+                with patch("rag.svr.task_executor_refactor.raptor_utils.collect_raptor_chunk_ids") as mock_collect:
+                    mock_collect.return_value = set()  # Primary has no RAPTOR chunks
+
+                    # Fallback will raise exception, but it should be caught
+                    result = await get_raptor_chunk_field_map("doc_1", "tenant_1", "kb_1")
+
+                    # Should return primary result (empty)
+                    assert result == {}
+        finally:
+            settings.docStoreConn = original_retriever
+
+
+class TestDeleteRaptorChunks:
+    """Tests for delete_raptor_chunks function."""
+
+    @pytest.mark.asyncio
+    async def test_deletes_all_chunks_when_keep_method_is_none(self):
+        """Test that all RAPTOR chunks are deleted when keep_method is None."""
+        from common import settings
+        original_retriever = settings.docStoreConn
+        
+        mock_doc_store = MagicMock()
+        settings.docStoreConn = mock_doc_store
+
+        try:
+            with patch("rag.svr.task_executor_refactor.raptor_utils.thread_pool_exec") as mock_thread:
+                mock_thread.return_value = 0
+
+                _ = await delete_raptor_chunks("doc_1", "tenant_1", "kb_1", keep_method=None)
+
+                mock_thread.assert_called_once()
+                # Verify delete was called with correct condition
+                call_args = mock_thread.call_args
+                assert call_args[0][0] == settings.docStoreConn.delete
+        finally:
+            settings.docStoreConn = original_retriever
+
+    @pytest.mark.asyncio
+    async def test_returns_0_when_no_stale_chunks(self):
+        """Test that 0 is returned when no stale chunks to delete."""
+        with patch("rag.svr.task_executor_refactor.raptor_utils.get_raptor_chunk_field_map") as mock_get_map:
+            mock_get_map.return_value = {}
+
+            with patch("rag.svr.task_executor_refactor.raptor_utils.collect_raptor_chunk_ids") as mock_collect:
+                mock_collect.return_value = set()  # No stale chunks
+
+                result = await delete_raptor_chunks("doc_1", "tenant_1", "kb_1", keep_method="raptor")
+
+                assert result == 0
+                mock_collect.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_deletes_stale_chunks_when_keep_method_specified(self):
+        """Test that stale chunks are deleted when keep_method is specified."""
+        from common import settings
+        original_retriever = settings.docStoreConn
+        
+        mock_doc_store = MagicMock()
+        settings.docStoreConn = mock_doc_store
+
+        try:
+            with patch("rag.svr.task_executor_refactor.raptor_utils.get_raptor_chunk_field_map") as mock_get_map:
+                mock_get_map.return_value = {
+                    "chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}},
+                    "chunk_2": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}}
+                }
+
+                with patch("rag.svr.task_executor_refactor.raptor_utils.collect_raptor_chunk_ids") as mock_collect:
+                    mock_collect.return_value = {"chunk_1"}  # Only chunk_1 is stale (psi, not raptor)
+
+                    with patch("rag.svr.task_executor_refactor.raptor_utils.thread_pool_exec") as mock_thread:
+                        mock_thread.return_value = 0
+
+                        _ = await delete_raptor_chunks("doc_1", "tenant_1", "kb_1", keep_method="raptor")
+
+                        # Should have called delete for stale chunks
+                        mock_thread.assert_called_once()
+        finally:
+            settings.docStoreConn = original_retriever
+
+    @pytest.mark.asyncio
+    async def test_logs_info_when_removing_stale_chunks(self):
+        """Test that info is logged when removing stale chunks."""
+        from common import settings
+        original_retriever = settings.docStoreConn
+        
+        mock_doc_store = MagicMock()
+        settings.docStoreConn = mock_doc_store
+
+        try:
+            with patch("rag.svr.task_executor_refactor.raptor_utils.get_raptor_chunk_field_map") as mock_get_map:
+                mock_get_map.return_value = {
+                    "chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}}
+                }
+
+                with patch("rag.svr.task_executor_refactor.raptor_utils.collect_raptor_chunk_ids") as mock_collect:
+                    mock_collect.return_value = {"chunk_1"}
+
+                    with patch("rag.svr.task_executor_refactor.raptor_utils.thread_pool_exec") as mock_thread:
+                        mock_thread.return_value = 0
+
+                        with patch("rag.svr.task_executor_refactor.raptor_utils.logging.info") as mock_log:
+                            await delete_raptor_chunks("doc_1", "tenant_1", "kb_1", keep_method="raptor")
+
+                            # Should have logged the removal
+                            mock_log.assert_called()
+        finally:
+            settings.docStoreConn = original_retriever

test/unit_test/rag/svr/task_executor_refactor/test_write_operation_interceptor.py

@@ -0,0 +1,228 @@
+#
+#  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.
+
+"""
+Unit tests for WriteOperationInterceptor module.
+"""
+
+import pytest
+from rag.svr.task_executor_refactor.write_operation_interceptor import (
+    WriteOperationInterceptor,
+    ALLOWED_METHOD_NAMES,
+)
+
+
+def _create_valid_recorded_values():
+    """Helper to create valid recorded_values dict."""
+    return {method: [] for method in ALLOWED_METHOD_NAMES}
+
+
+@pytest.fixture
+def valid_recorded_values():
+    """Provide a valid recorded_values dict for testing."""
+    return _create_valid_recorded_values()
+
+
+class TestAllowedMethodNames:
+    """Tests for ALLOWED_METHOD_NAMES constant."""
+
+    def test_allowed_method_names_count(self):
+        """Test that ALLOWED_METHOD_NAMES contains exactly 8 methods."""
+        assert len(ALLOWED_METHOD_NAMES) == 10
+
+    def test_allowed_method_names_contains_expected_methods(self):
+        """Test that ALLOWED_METHOD_NAMES contains all expected methods."""
+        expected_methods = {
+            "KnowledgebaseService.update_by_id",
+            "TaskService.update_chunk_ids",
+            "DocumentService.increment_chunk_num",
+            "DocMetadataService.update_document_metadata",
+            "PipelineOperationLogService.record_pipeline_operation",
+            "handle_save_to_memory_task",
+            "PipelineOperationLogService.create",
+            "delete_raptor_chunks",
+            "docStoreConn.insert",
+            "docStoreConn.delete"
+        }
+        assert ALLOWED_METHOD_NAMES == expected_methods
+
+
+class TestWriteOperationInterceptorInit:
+    """Tests for WriteOperationInterceptor.__init__."""
+
+    def test_init_with_valid_empty_values(self, valid_recorded_values):
+        """Test initialization with valid but empty values for all methods."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor is not None
+
+    def test_init_with_valid_values(self, valid_recorded_values):
+        """Test initialization with valid recorded values."""
+        valid_recorded_values["KnowledgebaseService.update_by_id"] = [1, 0]
+        valid_recorded_values["handle_save_to_memory_task"] = [None]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor is not None
+
+    def test_init_with_extra_keys_ignored(self, valid_recorded_values):
+        """Test that extra keys in recorded_values are ignored."""
+        valid_recorded_values["invalid_method_name"] = [1, 2, 3]
+        # Should not raise an error, extra keys are simply ignored
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor is not None
+        # The extra key should not be accessible
+        assert "invalid_method_name" not in interceptor._recorded_values
+
+
+class TestWriteOperationInterceptorIntercept:
+    """Tests for WriteOperationInterceptor.intercept."""
+
+    def test_intercept_returns_first_value(self, valid_recorded_values):
+        """Test that intercept returns the first value in the list."""
+        valid_recorded_values["KnowledgebaseService.update_by_id"] = [1, 0, 2]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        result = interceptor.intercept("KnowledgebaseService.update_by_id")
+        assert result == 1
+
+    def test_intercept_returns_subsequent_values(self, valid_recorded_values):
+        """Test that intercept returns subsequent values on each call."""
+        valid_recorded_values["KnowledgebaseService.update_by_id"] = [1, 0, 2]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor.intercept("KnowledgebaseService.update_by_id") == 1
+        assert interceptor.intercept("KnowledgebaseService.update_by_id") == 0
+        assert interceptor.intercept("KnowledgebaseService.update_by_id") == 2
+
+    def test_intercept_invalid_method_raises_value_error(self, valid_recorded_values):
+        """Test that intercepting an invalid method raises ValueError."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        with pytest.raises(ValueError, match="Cannot intercept method"):
+            interceptor.intercept("invalid_method_name")
+
+    def test_intercept_empty_list_raises_index_error(self, valid_recorded_values):
+        """Test that intercepting when list is empty raises IndexError."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        with pytest.raises(IndexError, match="No more recorded values"):
+            interceptor.intercept("KnowledgebaseService.update_by_id")
+
+    def test_intercept_pops_value(self, valid_recorded_values):
+        """Test that intercept pops the value from the internal list."""
+        valid_recorded_values["KnowledgebaseService.update_by_id"] = [42]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        interceptor.intercept("KnowledgebaseService.update_by_id")
+        # Check internal state, not the original input list (which is copied)
+        assert len(interceptor._recorded_values["KnowledgebaseService.update_by_id"]) == 0
+
+    def test_intercept_with_none_value(self, valid_recorded_values):
+        """Test that intercept can return None values."""
+        valid_recorded_values["handle_save_to_memory_task"] = [None]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        result = interceptor.intercept("handle_save_to_memory_task")
+        assert result is None
+
+    def test_intercept_with_default_value_when_empty(self, valid_recorded_values):
+        """Test that intercept returns default_value when list is empty."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        result = interceptor.intercept("KnowledgebaseService.update_by_id", default_value=42)
+        assert result == 42
+
+    def test_intercept_with_default_value_none_when_empty(self, valid_recorded_values):
+        """Test that intercept returns None when default_value is None and list is empty."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        # When default_value is None, it should return None (not raise IndexError)
+        # because None is a valid default value (different from _NO_DEFAULT sentinel)
+        result = interceptor.intercept("KnowledgebaseService.update_by_id", default_value=None)
+        assert result is None
+
+    def test_intercept_default_value_does_not_affect_existing_values(self, valid_recorded_values):
+        """Test that default_value is only used when list is empty."""
+        valid_recorded_values["KnowledgebaseService.update_by_id"] = [100]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        # Should return the recorded value, not the default_value
+        result = interceptor.intercept("KnowledgebaseService.update_by_id", default_value=999)
+        assert result == 100
+
+    @pytest.mark.parametrize("default_value", [
+        "default_string",
+        {"status": "success", "data": [1, 2, 3]},
+        [1, 2, 3, 4, 5],
+        (1, "two", 3.0),
+        True,
+        False,
+        0,
+        "",
+        [],
+        {},
+    ])
+    def test_intercept_with_various_default_values(self, valid_recorded_values, default_value):
+        """Test that intercept returns various default_value types when list is empty."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        result = interceptor.intercept("KnowledgebaseService.update_by_id", default_value=default_value)
+        assert result == default_value
+
+    def test_intercept_with_complex_values(self, valid_recorded_values):
+        """Test that intercept can return complex values like dicts and tuples."""
+        complex_value = {"key": "value", "nested": [1, 2, 3]}
+        valid_recorded_values["DocMetadataService.update_document_metadata"] = [complex_value]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        result = interceptor.intercept("DocMetadataService.update_document_metadata")
+        assert result == complex_value
+
+class TestWriteOperationInterceptorRemainingCount:
+    """Tests for WriteOperationInterceptor.remaining_count."""
+
+    def test_remaining_count_with_values(self, valid_recorded_values):
+        """Test remaining_count returns correct count."""
+        valid_recorded_values["KnowledgebaseService.update_by_id"] = [1, 2, 3]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor.remaining_count("KnowledgebaseService.update_by_id") == 3
+
+    def test_remaining_count_empty_list(self, valid_recorded_values):
+        """Test remaining_count returns 0 for empty list."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor.remaining_count("KnowledgebaseService.update_by_id") == 0
+
+        with pytest.raises(IndexError):
+            interceptor.intercept("KnowledgebaseService.update_by_id")
+
+    def test_remaining_count_after_intercept(self, valid_recorded_values):
+        """Test remaining_count decreases after intercept calls."""
+        valid_recorded_values["KnowledgebaseService.update_by_id"] = [1, 2, 3]
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor.remaining_count("KnowledgebaseService.update_by_id") == 3
+        interceptor.intercept("KnowledgebaseService.update_by_id")
+        assert interceptor.remaining_count("KnowledgebaseService.update_by_id") == 2
+        interceptor.intercept("KnowledgebaseService.update_by_id")
+        assert interceptor.remaining_count("KnowledgebaseService.update_by_id") == 1
+        interceptor.intercept("KnowledgebaseService.update_by_id")
+        assert interceptor.remaining_count("KnowledgebaseService.update_by_id") == 0
+
+    def test_remaining_count_invalid_method(self, valid_recorded_values):
+        """Test remaining_count returns 0 for invalid method names."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        assert interceptor.remaining_count("invalid_method") == 0
+
+
+class TestWriteOperationInterceptorRepr:
+    """Tests for WriteOperationInterceptor.__repr__."""
+
+    def test_repr_contains_class_name(self, valid_recorded_values):
+        """Test that repr contains the class name."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        repr_str = repr(interceptor)
+        assert "WriteOperationInterceptor" in repr_str
+
+    def test_repr_contains_total_recorded(self, valid_recorded_values):
+        """Test that repr contains total_recorded."""
+        interceptor = WriteOperationInterceptor(valid_recorded_values)
+        repr_str = repr(interceptor)
+        assert "total_recorded=" in repr_str

test/unit_test/rag/utils/test_raptor_utils.py

@@ -1,5 +1,5 @@
 #
-#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#  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.
@@ -12,395 +12,441 @@
 #  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.
-#
 
 """
-Unit tests for Raptor utility functions.
+Unit tests for rag/utils/raptor_utils.py module.
 """
 
-import logging
-
 import pytest
 from rag.utils.raptor_utils import (
-    CSV_EXTENSIONS,
-    EXCEL_EXTENSIONS,
-    STRUCTURED_EXTENSIONS,
-    collect_raptor_chunk_ids,
-    collect_raptor_methods,
-    get_raptor_clustering_method,
+    RAPTOR_TREE_BUILDER,
+    PSI_TREE_BUILDER,
+    GMM_CLUSTERING_METHOD,
+    AHC_CLUSTERING_METHOD,
     get_raptor_tree_builder,
-    get_skip_reason,
+    get_raptor_clustering_method,
+    _as_extra_dict,
+    _has_raptor_marker,
+    _raptor_methods_from_fields,
+    collect_raptor_methods,
+    collect_raptor_chunk_ids,
+    make_raptor_summary_chunk_id,
     is_structured_file_type,
     is_tabular_pdf,
-    make_raptor_summary_chunk_id,
     should_skip_raptor,
+    get_skip_reason,
 )
 
 
+class TestGetRaptorTreeBuilder:
+    """Tests for get_raptor_tree_builder function."""
+
+    def test_returns_default_raptor_tree_builder(self):
+        """Test that default tree builder is 'raptor'."""
+        result = get_raptor_tree_builder(None)
+        assert result == RAPTOR_TREE_BUILDER
+
+    def test_returns_default_with_empty_config(self):
+        """Test that empty config returns default."""
+        result = get_raptor_tree_builder({})
+        assert result == RAPTOR_TREE_BUILDER
+
+    def test_returns_configured_tree_builder(self):
+        """Test that configured tree builder is returned."""
+        config = {"tree_builder": PSI_TREE_BUILDER}
+        result = get_raptor_tree_builder(config)
+        assert result == PSI_TREE_BUILDER
+
+    def test_returns_ext_tree_builder(self):
+        """Test that ext.tree_builder takes precedence."""
+        config = {"tree_builder": "old", "ext": {"tree_builder": PSI_TREE_BUILDER}}
+        result = get_raptor_tree_builder(config)
+        assert result == PSI_TREE_BUILDER
+
+    def test_raises_error_for_unsupported_tree_builder(self):
+        """Test that unsupported tree builder raises ValueError."""
+        config = {"tree_builder": "unknown"}
+        with pytest.raises(ValueError, match="Unsupported RAPTOR tree builder"):
+            get_raptor_tree_builder(config)
+
+
+class TestGetRaptorClusteringMethod:
+    """Tests for get_raptor_clustering_method function."""
+
+    def test_returns_default_gmm(self):
+        """Test that default clustering method is 'gmm'."""
+        result = get_raptor_clustering_method(None)
+        assert result == GMM_CLUSTERING_METHOD
+
+    def test_returns_configured_clustering_method(self):
+        """Test that configured clustering method is returned."""
+        config = {"clustering_method": AHC_CLUSTERING_METHOD}
+        result = get_raptor_clustering_method(config)
+        assert result == AHC_CLUSTERING_METHOD
+
+    def test_returns_ext_clustering_method(self):
+        """Test that ext.clustering_method takes precedence."""
+        config = {"clustering_method": "old", "ext": {"clustering_method": AHC_CLUSTERING_METHOD}}
+        result = get_raptor_clustering_method(config)
+        assert result == AHC_CLUSTERING_METHOD
+
+    def test_raises_error_for_unsupported_clustering_method(self):
+        """Test that unsupported clustering method raises ValueError."""
+        config = {"clustering_method": "unknown"}
+        with pytest.raises(ValueError, match="Unsupported RAPTOR clustering method"):
+            get_raptor_clustering_method(config)
+
+
+class TestAsExtraDict:
+    """Tests for _as_extra_dict function."""
+
+    def test_returns_dict_as_is(self):
+        """Test that dict input is returned as-is."""
+        input_dict = {"key": "value"}
+        result = _as_extra_dict(input_dict)
+        assert result == input_dict
+
+    def test_returns_empty_dict_for_none(self):
+        """Test that None input returns empty dict."""
+        result = _as_extra_dict(None)
+        assert result == {}
+
+    def test_returns_empty_dict_for_empty_string(self):
+        """Test that empty string input returns empty dict."""
+        result = _as_extra_dict("")
+        assert result == {}
+
+    def test_parses_valid_json_string(self):
+        """Test that valid JSON string is parsed correctly."""
+        input_str = '{"key": "value"}'
+        result = _as_extra_dict(input_str)
+        assert result == {"key": "value"}
+
+    def test_returns_empty_dict_for_non_dict_json(self):
+        """Test that non-dict JSON returns empty dict."""
+        input_str = '[1, 2, 3]'
+        result = _as_extra_dict(input_str)
+        assert result == {}
+
+    def test_parses_python_dict_literal(self):
+        """Test that Python dict literal is parsed."""
+        input_str = "{'key': 'value'}"
+        result = _as_extra_dict(input_str)
+        assert result == {"key": "value"}
+
+    def test_returns_empty_dict_for_malformed_string(self):
+        """Test that malformed string returns empty dict."""
+        input_str = "{invalid json}"
+        result = _as_extra_dict(input_str)
+        assert result == {}
+
+
+class TestHasRaptorMarker:
+    """Tests for _has_raptor_marker function."""
+
+    def test_returns_true_for_raptor_string(self):
+        """Test that 'raptor' string returns True."""
+        assert _has_raptor_marker("raptor") is True
+
+    def test_returns_true_for_raptor_in_list(self):
+        """Test that 'raptor' in list returns True."""
+        assert _has_raptor_marker(["raptor", "other"]) is True
+
+    def test_returns_false_for_other_string(self):
+        """Test that other string returns False."""
+        assert _has_raptor_marker("other") is False
+
+    def test_returns_false_for_empty_list(self):
+        """Test that empty list returns False."""
+        assert _has_raptor_marker([]) is False
+
+    def test_returns_false_for_list_without_raptor(self):
+        """Test that list without 'raptor' returns False."""
+        assert _has_raptor_marker(["psi", "other"]) is False
+
+
+class TestRaptorMethodsFromFields:
+    """Tests for _raptor_methods_from_fields function."""
+
+    def test_returns_default_raptor_method(self):
+        """Test that default method is 'raptor'."""
+        result = _raptor_methods_from_fields({})
+        assert result == {RAPTOR_TREE_BUILDER}
+
+    def test_returns_method_from_extra_dict(self):
+        """Test that method is extracted from extra dict."""
+        fields = {"extra": {"raptor_method": PSI_TREE_BUILDER}}
+        result = _raptor_methods_from_fields(fields)
+        assert result == {PSI_TREE_BUILDER}
+
+    def test_returns_method_from_extra_field(self):
+        """Test that method is extracted from extra field directly."""
+        fields = {"extra": "{'raptor_method': 'psi'}"}
+        result = _raptor_methods_from_fields(fields)
+        assert result == {PSI_TREE_BUILDER}
+
+    def test_handles_list_method(self):
+        """Test that list method is converted to set."""
+        fields = {"extra": {"raptor_method": ["raptor", "psi"]}}
+        result = _raptor_methods_from_fields(fields)
+        assert result == {RAPTOR_TREE_BUILDER, PSI_TREE_BUILDER}
+
+    def test_handles_empty_method(self):
+        """Test that empty method returns default."""
+        fields = {"extra": {"raptor_method": ""}}
+        result = _raptor_methods_from_fields(fields)
+        assert result == {RAPTOR_TREE_BUILDER}
+
+
+class TestCollectRaptorMethods:
+    """Tests for collect_raptor_methods function."""
+
+    def test_returns_empty_set_for_empty_map(self):
+        """Test that empty field map returns empty set."""
+        result = collect_raptor_methods({})
+        assert result == set()
+
+    def test_collects_methods_from_raptor_chunks(self):
+        """Test that methods are collected from RAPTOR chunks."""
+        field_map = {
+            "chunk_1": {
+                "raptor_kwd": "raptor",
+                "extra": {"raptor_method": PSI_TREE_BUILDER}
+            }
+        }
+        result = collect_raptor_methods(field_map)
+        assert result == {PSI_TREE_BUILDER}
+
+    def test_skips_non_raptor_chunks(self):
+        """Test that non-RAPTOR chunks are skipped."""
+        field_map = {
+            "chunk_1": {
+                "raptor_kwd": "other",
+                "extra": {"raptor_method": PSI_TREE_BUILDER}
+            }
+        }
+        result = collect_raptor_methods(field_map)
+        assert result == set()
+
+    def test_collects_multiple_methods(self):
+        """Test that multiple methods are collected."""
+        field_map = {
+            "chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}},
+            "chunk_2": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}}
+        }
+        result = collect_raptor_methods(field_map)
+        assert result == {RAPTOR_TREE_BUILDER, PSI_TREE_BUILDER}
+
+
+class TestCollectRaptorChunkIds:
+    """Tests for collect_raptor_chunk_ids function."""
+
+    def test_returns_empty_set_for_empty_map(self):
+        """Test that empty field map returns empty set."""
+        result = collect_raptor_chunk_ids({})
+        assert result == set()
+
+    def test_collects_ids_of_raptor_chunks(self):
+        """Test that IDs of RAPTOR chunks are collected."""
+        field_map = {
+            "chunk_1": {"raptor_kwd": "raptor"},
+            "chunk_2": {"raptor_kwd": "raptor"}
+        }
+        result = collect_raptor_chunk_ids(field_map)
+        assert result == {"chunk_1", "chunk_2"}
+
+    def test_excludes_specified_methods(self):
+        """Test that specified methods are excluded."""
+        field_map = {
+            "chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}},
+            "chunk_2": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}}
+        }
+        result = collect_raptor_chunk_ids(field_map, exclude_methods={"raptor"})
+        assert result == {"chunk_2"}
+
+    def test_skips_non_raptor_chunks(self):
+        """Test that non-RAPTOR chunks are skipped."""
+        field_map = {
+            "chunk_1": {"raptor_kwd": "raptor"},
+            "chunk_2": {"raptor_kwd": "other"}
+        }
+        result = collect_raptor_chunk_ids(field_map)
+        assert result == {"chunk_1"}
+
+
+class TestMakeRaptorSummaryChunkId:
+    """Tests for make_raptor_summary_chunk_id function."""
+
+    def test_generates_consistent_id(self):
+        """Test that same input generates same ID."""
+        id1 = make_raptor_summary_chunk_id("content", "doc_1")
+        id2 = make_raptor_summary_chunk_id("content", "doc_1")
+        assert id1 == id2
+
+    def test_generates_different_ids_for_different_content(self):
+        """Test that different content generates different ID."""
+        id1 = make_raptor_summary_chunk_id("content1", "doc_1")
+        id2 = make_raptor_summary_chunk_id("content2", "doc_1")
+        assert id1 != id2
+
+    def test_generates_different_ids_for_different_doc(self):
+        """Test that different doc_id generates different ID."""
+        id1 = make_raptor_summary_chunk_id("content", "doc_1")
+        id2 = make_raptor_summary_chunk_id("content", "doc_2")
+        assert id1 != id2
+
+    def test_returns_string(self):
+        """Test that result is a string."""
+        result = make_raptor_summary_chunk_id("content", "doc_1")
+        assert isinstance(result, str)
+
+
 class TestIsStructuredFileType:
-    """Test file type detection for structured data"""
+    """Tests for is_structured_file_type function."""
 
-    @pytest.mark.parametrize("file_type,expected", [
-        (".xlsx", True),
-        (".xls", True),
-        (".xlsm", True),
-        (".xlsb", True),
-        (".csv", True),
-        (".tsv", True),
-        ("xlsx", True),  # Without leading dot
-        ("XLSX", True),  # Uppercase
-        (".pdf", False),
-        (".docx", False),
-        (".txt", False),
-        ("", False),
-        (None, False),
-    ])
-    def test_file_type_detection(self, file_type, expected):
-        """Test detection of various file types"""
-        assert is_structured_file_type(file_type) == expected
+    def test_returns_true_for_xlsx(self):
+        """Test that .xlsx is recognized as structured."""
+        assert is_structured_file_type(".xlsx") is True
 
-    def test_excel_extensions_defined(self):
-        """Test that Excel extensions are properly defined"""
-        assert ".xlsx" in EXCEL_EXTENSIONS
-        assert ".xls" in EXCEL_EXTENSIONS
-        assert len(EXCEL_EXTENSIONS) >= 4
+    def test_returns_true_for_xls(self):
+        """Test that .xls is recognized as structured."""
+        assert is_structured_file_type(".xls") is True
 
-    def test_csv_extensions_defined(self):
-        """Test that CSV extensions are properly defined"""
-        assert ".csv" in CSV_EXTENSIONS
-        assert ".tsv" in CSV_EXTENSIONS
+    def test_returns_true_for_csv(self):
+        """Test that .csv is recognized as structured."""
+        assert is_structured_file_type(".csv") is True
 
-    def test_structured_extensions_combined(self):
-        """Test that structured extensions include both Excel and CSV"""
-        assert EXCEL_EXTENSIONS.issubset(STRUCTURED_EXTENSIONS)
-        assert CSV_EXTENSIONS.issubset(STRUCTURED_EXTENSIONS)
+    def test_returns_true_for_tsv(self):
+        """Test that .tsv is recognized as structured."""
+        assert is_structured_file_type(".tsv") is True
+
+    def test_returns_false_for_pdf(self):
+        """Test that .pdf is not structured."""
+        assert is_structured_file_type(".pdf") is False
+
+    def test_returns_false_for_txt(self):
+        """Test that .txt is not structured."""
+        assert is_structured_file_type(".txt") is False
+
+    def test_returns_false_for_none(self):
+        """Test that None is not structured."""
+        assert is_structured_file_type(None) is False
+
+    def test_returns_false_for_empty_string(self):
+        """Test that empty string is not structured."""
+        assert is_structured_file_type("") is False
+
+    def test_handles_case_insensitive(self):
+        """Test that case is handled insensitively."""
+        assert is_structured_file_type(".XLSX") is True
+        assert is_structured_file_type("xlsx") is True
+
+    def test_handles_missing_dot(self):
+        """Test that missing dot is handled."""
+        assert is_structured_file_type("xlsx") is True
 
 
-class TestIsTabularPDF:
-    """Test tabular PDF detection"""
+class TestIsTabularPdf:
+    """Tests for is_tabular_pdf function."""
 
-    def test_table_parser_detected(self):
-        """Test that table parser is detected as tabular"""
+    def test_returns_true_for_table_parser(self):
+        """Test that table parser returns True."""
         assert is_tabular_pdf("table", {}) is True
-        assert is_tabular_pdf("TABLE", {}) is True
 
-    def test_html4excel_detected(self):
-        """Test that html4excel config is detected as tabular"""
+    def test_returns_true_for_html4excel(self):
+        """Test that html4excel enabled returns True."""
         assert is_tabular_pdf("naive", {"html4excel": True}) is True
-        assert is_tabular_pdf("", {"html4excel": True}) is True
 
-    def test_non_tabular_pdf(self):
-        """Test that non-tabular PDFs are not detected"""
+    def test_returns_false_for_naive_parser(self):
+        """Test that naive parser returns False."""
         assert is_tabular_pdf("naive", {}) is False
-        assert is_tabular_pdf("naive", {"html4excel": False}) is False
+
+    def test_returns_false_for_empty_parser_id(self):
+        """Test that empty parser_id returns False."""
         assert is_tabular_pdf("", {}) is False
 
-    def test_combined_conditions(self):
-        """Test combined table parser and html4excel"""
-        assert is_tabular_pdf("table", {"html4excel": True}) is True
-        assert is_tabular_pdf("table", {"html4excel": False}) is True
+    def test_returns_false_for_html4excel_false(self):
+        """Test that html4excel=False returns False."""
+        assert is_tabular_pdf("naive", {"html4excel": False}) is False
+
+    def test_handles_case_insensitive_parser_id(self):
+        """Test that parser_id case is handled."""
+        assert is_tabular_pdf("TABLE", {}) is True
+        assert is_tabular_pdf("Table", {}) is True
 
 
 class TestShouldSkipRaptor:
-    """Test Raptor skip logic"""
+    """Tests for should_skip_raptor function."""
 
-    def test_skip_excel_files(self):
-        """Test that Excel files skip Raptor"""
-        assert should_skip_raptor(".xlsx") is True
-        assert should_skip_raptor(".xls") is True
-        assert should_skip_raptor(".xlsm") is True
+    def test_skips_for_xlsx_file(self):
+        """Test that .xlsx file skips Raptor."""
+        assert should_skip_raptor(file_type=".xlsx") is True
 
-    def test_skip_csv_files(self):
-        """Test that CSV files skip Raptor"""
-        assert should_skip_raptor(".csv") is True
-        assert should_skip_raptor(".tsv") is True
+    def test_skips_for_csv_file(self):
+        """Test that .csv file skips Raptor."""
+        assert should_skip_raptor(file_type=".csv") is True
 
-    def test_skip_tabular_pdf_with_table_parser(self):
-        """Test that tabular PDFs skip Raptor"""
-        assert should_skip_raptor(".pdf", parser_id="table") is True
-        assert should_skip_raptor("pdf", parser_id="TABLE") is True
+    def test_skips_for_tabular_pdf(self):
+        """Test that tabular PDF skips Raptor."""
+        assert should_skip_raptor(file_type=".pdf", parser_id="table") is True
 
-    def test_skip_tabular_pdf_with_html4excel(self):
-        """Test that PDFs with html4excel skip Raptor"""
-        assert should_skip_raptor(".pdf", parser_config={"html4excel": True}) is True
+    def test_does_not_skip_for_normal_pdf(self):
+        """Test that normal PDF does not skip Raptor."""
+        assert should_skip_raptor(file_type=".pdf", parser_id="naive") is False
 
-    def test_dont_skip_regular_pdf(self):
-        """Test that regular PDFs don't skip Raptor"""
-        assert should_skip_raptor(".pdf", parser_id="naive") is False
-        assert should_skip_raptor(".pdf", parser_config={}) is False
+    def test_does_not_skip_for_txt_file(self):
+        """Test that .txt file does not skip Raptor."""
+        assert should_skip_raptor(file_type=".txt") is False
 
-    def test_dont_skip_text_files(self):
-        """Test that text files don't skip Raptor"""
-        assert should_skip_raptor(".txt") is False
-        assert should_skip_raptor(".docx") is False
-        assert should_skip_raptor(".md") is False
+    def test_respects_auto_disable_config_false(self):
+        """Test that auto_disable_for_structured_data=False disables skipping."""
+        assert should_skip_raptor(
+            file_type=".xlsx",
+            raptor_config={"auto_disable_for_structured_data": False}
+        ) is False
 
-    def test_override_with_config(self):
-        """Test that auto-disable can be overridden"""
-        raptor_config = {"auto_disable_for_structured_data": False}
-        
-        # Should not skip even for Excel files
-        assert should_skip_raptor(".xlsx", raptor_config=raptor_config) is False
-        assert should_skip_raptor(".csv", raptor_config=raptor_config) is False
-        assert should_skip_raptor(".pdf", parser_id="table", raptor_config=raptor_config) is False
+    def test_respects_auto_disable_config_true(self):
+        """Test that auto_disable_for_structured_data=True enables skipping."""
+        assert should_skip_raptor(
+            file_type=".xlsx",
+            raptor_config={"auto_disable_for_structured_data": True}
+        ) is True
 
-    def test_default_auto_disable_enabled(self):
-        """Test that auto-disable is enabled by default"""
-        # Empty raptor_config should default to auto_disable=True
-        assert should_skip_raptor(".xlsx", raptor_config={}) is True
-        assert should_skip_raptor(".xlsx", raptor_config=None) is True
+    def test_default_auto_disable_is_true(self):
+        """Test that default auto_disable is True."""
+        assert should_skip_raptor(file_type=".xlsx") is True
 
-    def test_explicit_auto_disable_enabled(self):
-        """Test explicit auto-disable enabled"""
-        raptor_config = {"auto_disable_for_structured_data": True}
-        assert should_skip_raptor(".xlsx", raptor_config=raptor_config) is True
+    def test_returns_false_for_none_file_type(self):
+        """Test that None file_type does not skip."""
+        assert should_skip_raptor(file_type=None) is False
 
 
 class TestGetSkipReason:
-    """Test skip reason generation"""
+    """Tests for get_skip_reason function."""
 
-    def test_excel_skip_reason(self):
-        """Test skip reason for Excel files"""
-        reason = get_skip_reason(".xlsx")
+    def test_returns_reason_for_structured_file(self):
+        """Test that reason is returned for structured file."""
+        reason = get_skip_reason(file_type=".xlsx")
         assert "Structured data file" in reason
         assert ".xlsx" in reason
-        assert "auto-disabled" in reason.lower()
 
-    def test_csv_skip_reason(self):
-        """Test skip reason for CSV files"""
-        reason = get_skip_reason(".csv")
-        assert "Structured data file" in reason
-        assert ".csv" in reason
-
-    def test_tabular_pdf_skip_reason(self):
-        """Test skip reason for tabular PDFs"""
-        reason = get_skip_reason(".pdf", parser_id="table")
+    def test_returns_reason_for_tabular_pdf(self):
+        """Test that reason is returned for tabular PDF."""
+        reason = get_skip_reason(file_type=".pdf", parser_id="table")
         assert "Tabular PDF" in reason
-        assert "table" in reason.lower()
-        assert "auto-disabled" in reason.lower()
+        assert "table" in reason
 
-    def test_html4excel_skip_reason(self):
-        """Test skip reason for html4excel PDFs"""
-        reason = get_skip_reason(".pdf", parser_config={"html4excel": True})
-        assert "Tabular PDF" in reason
-
-    def test_no_skip_reason_for_regular_files(self):
-        """Test that regular files have no skip reason"""
-        assert get_skip_reason(".txt") == ""
-        assert get_skip_reason(".docx") == ""
-        assert get_skip_reason(".pdf", parser_id="naive") == ""
-
-
-class TestEdgeCases:
-    """Test edge cases and error handling"""
-
-    def test_none_values(self):
-        """Test handling of None values"""
-        assert should_skip_raptor(None) is False
-        assert should_skip_raptor("") is False
-        assert get_skip_reason(None) == ""
-
-    def test_empty_strings(self):
-        """Test handling of empty strings"""
-        assert should_skip_raptor("") is False
-        assert get_skip_reason("") == ""
-
-    def test_case_insensitivity(self):
-        """Test case insensitive handling"""
-        assert is_structured_file_type("XLSX") is True
-        assert is_structured_file_type("XlSx") is True
-        assert is_tabular_pdf("TABLE", {}) is True
-        assert is_tabular_pdf("TaBlE", {}) is True
-
-    def test_with_and_without_dot(self):
-        """Test file extensions with and without leading dot"""
-        assert should_skip_raptor(".xlsx") is True
-        assert should_skip_raptor("xlsx") is True
-        assert should_skip_raptor(".CSV") is True
-        assert should_skip_raptor("csv") is True
-
-
-class TestIntegrationScenarios:
-    """Test real-world integration scenarios"""
-
-    def test_financial_excel_report(self):
-        """Test scenario: Financial quarterly Excel report"""
-        file_type = ".xlsx"
-        parser_id = "naive"
-        parser_config = {}
-        raptor_config = {"use_raptor": True}
-        
-        # Should skip Raptor
-        assert should_skip_raptor(file_type, parser_id, parser_config, raptor_config) is True
-        reason = get_skip_reason(file_type, parser_id, parser_config)
-        assert "Structured data file" in reason
-
-    def test_scientific_csv_data(self):
-        """Test scenario: Scientific experimental CSV results"""
-        file_type = ".csv"
-        
-        # Should skip Raptor
-        assert should_skip_raptor(file_type) is True
-        reason = get_skip_reason(file_type)
-        assert ".csv" in reason
-
-    def test_legal_contract_with_tables(self):
-        """Test scenario: Legal contract PDF with tables"""
-        file_type = ".pdf"
-        parser_id = "table"
-        parser_config = {}
-        
-        # Should skip Raptor
-        assert should_skip_raptor(file_type, parser_id, parser_config) is True
-        reason = get_skip_reason(file_type, parser_id, parser_config)
-        assert "Tabular PDF" in reason
-
-    def test_text_heavy_pdf_document(self):
-        """Test scenario: Text-heavy PDF document"""
-        file_type = ".pdf"
-        parser_id = "naive"
-        parser_config = {}
-        
-        # Should NOT skip Raptor
-        assert should_skip_raptor(file_type, parser_id, parser_config) is False
-        reason = get_skip_reason(file_type, parser_id, parser_config)
+    def test_returns_empty_for_normal_pdf(self):
+        """Test that empty reason is returned for normal PDF."""
+        reason = get_skip_reason(file_type=".pdf", parser_id="naive")
         assert reason == ""
 
-    def test_mixed_dataset_processing(self):
-        """Test scenario: Mixed dataset with various file types"""
-        files = [
-            (".xlsx", "naive", {}, True),  # Excel - skip
-            (".csv", "naive", {}, True),   # CSV - skip
-            (".pdf", "table", {}, True),   # Tabular PDF - skip
-            (".pdf", "naive", {}, False),  # Regular PDF - don't skip
-            (".docx", "naive", {}, False), # Word doc - don't skip
-            (".txt", "naive", {}, False),  # Text file - don't skip
-        ]
-        
-        for file_type, parser_id, parser_config, expected_skip in files:
-            result = should_skip_raptor(file_type, parser_id, parser_config)
-            assert result == expected_skip, f"Failed for {file_type}"
+    def test_returns_empty_for_txt_file(self):
+        """Test that empty reason is returned for .txt file."""
+        reason = get_skip_reason(file_type=".txt")
+        assert reason == ""
 
-    def test_override_for_special_excel(self):
-        """Test scenario: Override auto-disable for special Excel processing"""
-        file_type = ".xlsx"
-        raptor_config = {"auto_disable_for_structured_data": False}
-        
-        # Should NOT skip when explicitly disabled
-        assert should_skip_raptor(file_type, raptor_config=raptor_config) is False
-
-
-class TestRaptorTreeBuilderConfig:
-    """Test RAPTOR tree builder config resolution"""
-
-    def test_defaults_to_original_raptor_builder(self):
-        assert get_raptor_tree_builder({}) == "raptor"
-        assert get_raptor_tree_builder(None) == "raptor"
-
-    def test_reads_top_level_tree_builder(self):
-        assert get_raptor_tree_builder({"tree_builder": "psi"}) == "psi"
-
-    def test_reads_legacy_ext_tree_builder(self):
-        assert get_raptor_tree_builder({"ext": {"tree_builder": "psi"}}) == "psi"
-
-    def test_ext_tree_builder_overrides_stale_top_level_value(self):
-        assert get_raptor_tree_builder({"tree_builder": "psi", "ext": {"tree_builder": "raptor"}}) == "raptor"
-
-    def test_rejects_unknown_tree_builder(self):
-        with pytest.raises(ValueError, match="Unsupported RAPTOR tree builder"):
-            get_raptor_tree_builder({"tree_builder": "ahc"})
-
-
-class TestRaptorClusteringMethodConfig:
-    """Test RAPTOR clustering method config resolution"""
-
-    def test_defaults_to_gmm(self):
-        assert get_raptor_clustering_method({}) == "gmm"
-        assert get_raptor_clustering_method(None) == "gmm"
-
-    def test_reads_top_level_clustering_method(self):
-        assert get_raptor_clustering_method({"clustering_method": "gmm"}) == "gmm"
-        assert get_raptor_clustering_method({"clustering_method": "ahc"}) == "ahc"
-
-    def test_reads_legacy_ext_clustering_method(self):
-        assert get_raptor_clustering_method({"ext": {"clustering_method": "ahc"}}) == "ahc"
-
-    def test_ext_clustering_method_overrides_stale_top_level_value(self):
-        assert get_raptor_clustering_method({"clustering_method": "gmm", "ext": {"clustering_method": "ahc"}}) == "ahc"
-
-    def test_rejects_unknown_clustering_method(self):
-        with pytest.raises(ValueError, match="Unsupported RAPTOR clustering method"):
-            get_raptor_clustering_method({"clustering_method": "unknown"})
-
-
-class TestRaptorMethodCollection:
-    """Test RAPTOR summary method extraction from doc-store fields"""
-
-    def test_legacy_summary_without_method_is_original_raptor(self):
-        field_map = {"chunk_1": {"raptor_kwd": "raptor"}}
-
-        assert collect_raptor_methods(field_map) == {"raptor"}
-        assert collect_raptor_chunk_ids(field_map) == {"chunk_1"}
-
-    def test_extra_method_is_preserved(self):
-        field_map = {"chunk_1": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}}}
-
-        assert collect_raptor_methods(field_map) == {"psi"}
-        assert collect_raptor_chunk_ids(field_map) == {"chunk_1"}
-
-    def test_extra_field_supports_oceanbase_legacy_rows(self):
-        field_map = {
-            "chunk_1": {
-                "extra": {
-                    "raptor_kwd": "raptor",
-                    "raptor_method": "psi",
-                }
-            },
-            "chunk_2": {
-                "extra": "{\"raptor_kwd\": \"raptor\"}",
-            },
-            "chunk_3": {
-                "extra": {"raptor_kwd": ""},
-            },
-        }
-
-        assert collect_raptor_methods(field_map) == {"psi", "raptor"}
-        assert collect_raptor_chunk_ids(field_map) == {"chunk_1", "chunk_2"}
-
-    def test_non_raptor_rows_are_ignored(self):
-        field_map = {
-            "chunk_1": {"raptor_kwd": ""},
-            "chunk_2": {"extra": {"raptor_kwd": "graph"}},
-            "chunk_3": {},
-        }
-
-        assert collect_raptor_methods(field_map) == set()
-        assert collect_raptor_chunk_ids(field_map) == set()
-
-    def test_malformed_extra_payload_is_logged_and_ignored(self, caplog):
-        field_map = {"chunk_1": {"extra": "{bad json"}}
-
-        with caplog.at_level(logging.WARNING):
-            assert collect_raptor_methods(field_map) == set()
-            assert collect_raptor_chunk_ids(field_map) == set()
-
-        assert "Ignoring malformed RAPTOR extra payload" in caplog.text
-
-    def test_chunk_id_collection_can_preserve_current_method(self):
-        field_map = {
-            "legacy": {"raptor_kwd": "raptor"},
-            "old": {"raptor_kwd": "raptor", "extra": {"raptor_method": "raptor"}},
-            "current": {"raptor_kwd": "raptor", "extra": {"raptor_method": "psi"}},
-        }
-
-        assert collect_raptor_chunk_ids(field_map, exclude_methods={"psi"}) == {"legacy", "old"}
-        assert collect_raptor_chunk_ids(field_map, exclude_methods={"raptor"}) == {"current"}
-
-    def test_summary_chunk_ids_include_real_document_id(self):
-        content = "same generated summary"
-
-        assert make_raptor_summary_chunk_id(content, "doc-a") != make_raptor_summary_chunk_id(content, "doc-b")
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-v"])
+    def test_returns_empty_for_none_file_type(self):
+        """Test that empty reason is returned for None file_type."""
+        reason = get_skip_reason(file_type=None)
+        assert reason == ""