diff --git a/api/db/joint_services/memory_message_service.py b/api/db/joint_services/memory_message_service.py index 1a6da3a8d6..fd86342181 100644 --- a/api/db/joint_services/memory_message_service.py +++ b/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): diff --git a/api/db/services/document_service.py b/api/db/services/document_service.py index 348d8a3a60..12fdc19fef 100644 --- a/api/db/services/document_service.py +++ b/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) diff --git a/api/db/services/llm_service.py b/api/db/services/llm_service.py index 9b6b5bd4f1..acd35bfe6e 100644 --- a/api/db/services/llm_service.py +++ b/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 /EMBEDDING used_tokens: {}".format(used_tokens)) diff --git a/api/db/services/task_service.py b/api/db/services/task_service.py index 640c8fbd25..1ce0cd4dd6 100644 --- a/api/db/services/task_service.py +++ b/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." diff --git a/api/db/services/tenant_llm_service.py b/api/db/services/tenant_llm_service.py index f14f97fcef..7fae4571cc 100644 --- a/api/db/services/tenant_llm_service.py +++ b/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 diff --git a/common/constants.py b/common/constants.py index c76dcdbb09..b222c4caf7 100644 --- a/common/constants.py +++ b/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" diff --git a/common/decorator.py b/common/decorator.py index f45a41a9d8..7dd0319f43 100644 --- a/common/decorator.py +++ b/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 \ No newline at end of file + 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 \ No newline at end of file diff --git a/common/settings.py b/common/settings.py index 49693b9370..1c313b3494 100644 --- a/common/settings.py +++ b/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") diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh index 79f77fe43a..99ae05fb6b 100755 --- a/docker/entrypoint.sh +++ b/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 diff --git a/docker/launch_backend_service.sh b/docker/launch_backend_service.sh index c76381fa85..2f5ddb14c8 100755 --- a/docker/launch_backend_service.sh +++ b/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." diff --git a/rag/flow/pipeline.py b/rag/flow/pipeline.py index cc4bed0fab..76e19084e0 100644 --- a/rag/flow/pipeline.py +++ b/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( diff --git a/rag/flow/tokenizer/tokenizer.py b/rag/flow/tokenizer/tokenizer.py index d38554aec4..e4be09f005 100644 --- a/rag/flow/tokenizer/tokenizer.py +++ b/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 diff --git a/rag/prompts/assign_toc_levels.md b/rag/prompts/assign_toc_levels.md index ce80c22622..5e6ab22911 100644 --- a/rag/prompts/assign_toc_levels.md +++ b/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": }. -- 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 diff --git a/rag/prompts/generator.py b/rag/prompts/generator.py index fc4999dbe4..83e277d4b8 100644 --- a/rag/prompts/generator.py +++ b/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" diff --git a/rag/svr/task_executor.py b/rag/svr/task_executor.py index ce41c2b28b..ded3a0141a 100644 --- a/rag/svr/task_executor.py +++ b/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: diff --git a/rag/svr/task_executor_limiter.py b/rag/svr/task_executor_limiter.py new file mode 100644 index 0000000000..61b50849b3 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/rag/svr/task_executor_refactor/chunk_builder.py b/rag/svr/task_executor_refactor/chunk_builder.py new file mode 100644 index 0000000000..b9dc353b4e --- /dev/null +++ b/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) diff --git a/rag/svr/task_executor_refactor/chunk_post_processor.py b/rag/svr/task_executor_refactor/chunk_post_processor.py new file mode 100644 index 0000000000..fc12453357 --- /dev/null +++ b/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)) diff --git a/rag/svr/task_executor_refactor/chunk_service.py b/rag/svr/task_executor_refactor/chunk_service.py new file mode 100644 index 0000000000..060a99ed5f --- /dev/null +++ b/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 diff --git a/rag/svr/task_executor_refactor/comparator.py b/rag/svr/task_executor_refactor/comparator.py new file mode 100644 index 0000000000..0ace8d9aa0 --- /dev/null +++ b/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}", + ) diff --git a/rag/svr/task_executor_refactor/constants.py b/rag/svr/task_executor_refactor/constants.py new file mode 100644 index 0000000000..1c90bcaef7 --- /dev/null +++ b/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" diff --git a/rag/svr/task_executor_refactor/dataflow_service.py b/rag/svr/task_executor_refactor/dataflow_service.py new file mode 100644 index 0000000000..006c6edff3 --- /dev/null +++ b/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 diff --git a/rag/svr/task_executor_refactor/embedding_service.py b/rag/svr/task_executor_refactor/embedding_service.py new file mode 100644 index 0000000000..35ddde7002 --- /dev/null +++ b/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()) diff --git a/rag/svr/task_executor_refactor/embedding_utils.py b/rag/svr/task_executor_refactor/embedding_utils.py new file mode 100644 index 0000000000..44eec7e125 --- /dev/null +++ b/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"]{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 diff --git a/rag/svr/task_executor_refactor/post_processor.py b/rag/svr/task_executor_refactor/post_processor.py new file mode 100644 index 0000000000..31a73a5bed --- /dev/null +++ b/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) diff --git a/rag/svr/task_executor_refactor/raptor_service.py b/rag/svr/task_executor_refactor/raptor_service.py new file mode 100644 index 0000000000..ba80bac3a5 --- /dev/null +++ b/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 diff --git a/rag/svr/task_executor_refactor/raptor_utils.py b/rag/svr/task_executor_refactor/raptor_utils.py new file mode 100644 index 0000000000..f98975bfb2 --- /dev/null +++ b/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) diff --git a/rag/svr/task_executor_refactor/recording_context.py b/rag/svr/task_executor_refactor/recording_context.py new file mode 100644 index 0000000000..bf64a68644 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/rag/svr/task_executor_refactor/report_generator.py b/rag/svr/task_executor_refactor/report_generator.py new file mode 100644 index 0000000000..725fb6ff4b --- /dev/null +++ b/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) diff --git a/rag/svr/task_executor_refactor/task_context.py b/rag/svr/task_executor_refactor/task_context.py new file mode 100644 index 0000000000..2a15d6b50c --- /dev/null +++ b/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 diff --git a/rag/svr/task_executor_refactor/task_executor_refactoring_plan.md b/rag/svr/task_executor_refactor/task_executor_refactoring_plan.md new file mode 100644 index 0000000000..9bf0d35fb3 --- /dev/null +++ b/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. diff --git a/rag/svr/task_executor_refactor/task_handler.py b/rag/svr/task_executor_refactor/task_handler.py new file mode 100644 index 0000000000..deee1b4b36 --- /dev/null +++ b/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) diff --git a/rag/svr/task_executor_refactor/task_manager.py b/rag/svr/task_executor_refactor/task_manager.py new file mode 100644 index 0000000000..041a1b8924 --- /dev/null +++ b/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 ") \ No newline at end of file diff --git a/rag/svr/task_executor_refactor/write_operation_interceptor.py b/rag/svr/task_executor_refactor/write_operation_interceptor.py new file mode 100644 index 0000000000..fe57f8d9c8 --- /dev/null +++ b/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})" diff --git a/rag/utils/base64_image.py b/rag/utils/base64_image.py index bc73c0433f..4aa1e623c4 100644 --- a/rag/utils/base64_image.py +++ b/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 diff --git a/rag/utils/raptor_utils.py b/rag/utils/raptor_utils.py index 91d43cd937..e461841abf 100644 --- a/rag/utils/raptor_utils.py +++ b/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 {} diff --git a/test/unit_test/common/test_settings_queue.py b/test/unit_test/common/test_settings_queue.py new file mode 100644 index 0000000000..6ec582ab9f --- /dev/null +++ b/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 diff --git a/test/unit_test/rag/svr/task_executor_refactor/conftest.py b/test/unit_test/rag/svr/task_executor_refactor/conftest.py new file mode 100644 index 0000000000..cb20ba9998 --- /dev/null +++ b/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() diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_chunk_builder.py b/test/unit_test/rag/svr/task_executor_refactor/test_chunk_builder.py new file mode 100644 index 0000000000..1c4f844e8f --- /dev/null +++ b/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) diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_chunk_post_processor.py b/test/unit_test/rag/svr/task_executor_refactor/test_chunk_post_processor.py new file mode 100644 index 0000000000..1d684fc0ef --- /dev/null +++ b/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 == [] diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_chunk_service.py b/test/unit_test/rag/svr/task_executor_refactor/test_chunk_service.py new file mode 100644 index 0000000000..60d937ab29 --- /dev/null +++ b/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] diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_comparator.py b/test/unit_test/rag/svr/task_executor_refactor/test_comparator.py new file mode 100644 index 0000000000..207795d798 --- /dev/null +++ b/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 diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_constants.py b/test/unit_test/rag/svr/task_executor_refactor/test_constants.py new file mode 100644 index 0000000000..28b9dc7071 --- /dev/null +++ b/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 diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_dataflow_service.py b/test/unit_test/rag/svr/task_executor_refactor/test_dataflow_service.py new file mode 100644 index 0000000000..1ec8a75063 --- /dev/null +++ b/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": "

content

"}) + assert result == [{"text": ["

content

"]}] + + 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 \ No newline at end of file diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_embedding_service.py b/test/unit_test/rag/svr/task_executor_refactor/test_embedding_service.py new file mode 100644 index 0000000000..b6bfccb11c --- /dev/null +++ b/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] diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_embedding_utils.py b/test/unit_test/rag/svr/task_executor_refactor/test_embedding_utils.py new file mode 100644 index 0000000000..2c7601e852 --- /dev/null +++ b/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": "
Cell
"}, + ] + titles, contents = EmbeddingUtils.prepare_texts_for_embedding(docs) + # Table tags should be replaced with spaces + assert "" 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 = "
Cell
" + result = EmbeddingUtils._normalize_table_html(html) + assert "" not in result + assert "" not in result + assert "
" 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" diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_post_processor.py b/test/unit_test/rag/svr/task_executor_refactor/test_post_processor.py new file mode 100644 index 0000000000..ce808b2e52 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_raptor_service.py b/test/unit_test/rag/svr/task_executor_refactor/test_raptor_service.py new file mode 100644 index 0000000000..2399d13ed6 --- /dev/null +++ b/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 diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_recording_context.py b/test/unit_test/rag/svr/task_executor_refactor/test_recording_context.py new file mode 100644 index 0000000000..d5336393b1 --- /dev/null +++ b/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() diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_task_context.py b/test/unit_test/rag/svr/task_executor_refactor/test_task_context.py new file mode 100644 index 0000000000..b7af25499a --- /dev/null +++ b/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 diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_task_handler.py b/test/unit_test/rag/svr/task_executor_refactor/test_task_handler.py new file mode 100644 index 0000000000..d8149867a9 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_task_handler_integration.py b/test/unit_test/rag/svr/task_executor_refactor/test_task_handler_integration.py new file mode 100644 index 0000000000..dbacf251ae --- /dev/null +++ b/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}" diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_task_raptor_utils.py b/test/unit_test/rag/svr/task_executor_refactor/test_task_raptor_utils.py new file mode 100644 index 0000000000..01bc3ecfa8 --- /dev/null +++ b/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 diff --git a/test/unit_test/rag/svr/task_executor_refactor/test_write_operation_interceptor.py b/test/unit_test/rag/svr/task_executor_refactor/test_write_operation_interceptor.py new file mode 100644 index 0000000000..59df761051 --- /dev/null +++ b/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 diff --git a/test/unit_test/rag/utils/test_raptor_utils.py b/test/unit_test/rag/utils/test_raptor_utils.py index 95abe21097..b0b8581e31 100644 --- a/test/unit_test/rag/utils/test_raptor_utils.py +++ b/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 == ""