From b36314699774b5439bf06c078060dbccde38ef4e Mon Sep 17 00:00:00 2001 From: Jack Date: Wed, 3 Jun 2026 17:18:31 +0800 Subject: [PATCH] refactor: overhaul task executor with layered architecture and comprehensive test suite (#15471) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Decomposes the monolithic `task_executor.py` (1945 lines) into a 6-layer architecture with clear separation of concerns. The refactored code is functionally equivalent to the original, verified through 400 passing tests and a production-vs-dry-run comparison framework. ## Architecture ``` entry (task_manager) └─ orchestration (task_handler) ├─ services (chunk_service, embedding_service, dataflow_service, raptor_service, post_processor) │ └─ utilities (chunk_builder, chunk_post_processor, embedding_utils) └─ infrastructure (task_context, recording_context, interceptor) ``` Key design decisions: - **TaskContext** — typed facade over raw task dict, injects rate limiters + callbacks via composition - **RecordingContext + Comparator** — enables side-by-side production vs dry-run execution for safe migration - **NullRecordingContext** — zero-allocation no-op for production, uses `__slots__` - **WriteOperationInterceptor** — FIFO replay of previous runs function returns for comparison mode ## Migration Strategy The original `handle_task()` in `task_executor.py` uses a 3-way switch via `TE_RUN_MODE`: - `TE_RUN_MODE=0` (default) → runs refactored code - `TE_RUN_MODE=1` → runs both original + refactored, compares all intermediate results - `TE_RUN_MODE=2` → runs original code (fallback) The comparison mode (`TE_RUN_MODE=1`) records ~40 intermediate values (chunks, vectors, token counts, func return values) from the production run and replays them during dry-run, then uses `ContextComparator` to report mismatches. ## Functional Equivalence Fixes All divergences between original and refactored code were identified and fixed: - Timeout decorators (handle/build_chunks/raptor/embedding) - NullRecordingContext leak in finally block causing RuntimeError - MinIO None-binary check with proper FileNotFoundError - Dataflow dispatch after embedding binding + init_kb - Memory task missing return after processing - RAPTOR checkpoint progress reporting - Tag cache (get_tags_from_cache/set_tags_to_cache) restoration - dataflow_id correction in _load_dsl - Language default Chinese, dead code guard removal - embed_chunks made async with proper thread_pool_exec - Full GraphRAG default configuration (10 parameters) - Hardcoded q_768_vec fallback removal in RAPTOR ## Test Changes - 20 new tests covering table parser manual mode, tag cache, embedding edge cases, RAPTOR checkpoint, dataflow_id correction, storage binary None, cancel cleanup, metadata=None boundary - Unified `make_task_context`/`make_task_dict` factories eliminated 10+ duplicated helpers - DataflowService tests migrated from internal method mocks to IO boundary mocks (real orchestration code executes) - Parametrized duplicate build_chunks post-processor tests - 7 raptor tests modernized to @pytest.mark.asyncio - Mock count per test reduced through boundary-level mocking strategy **Test count: 400 passing, 0 warnings, 0 skips** ## Files Changed | File | Change | |------|--------| | `rag/svr/task_executor.py` | +1 line (NullRecordingContext fix) | | `rag/svr/task_executor_refactor/task_handler.py` | Orchestration layer, 8 logic fixes | | `rag/svr/task_executor_refactor/chunk_service.py` | +timeout + None-check | | `rag/svr/task_executor_refactor/embedding_service.py` | sync→async rewrite | | `rag/svr/task_executor_refactor/dataflow_service.py` | dataflow_id fix + timeout | | `rag/svr/task_executor_refactor/raptor_service.py` | checkpoint fix + assert | | `rag/svr/task_executor_refactor/chunk_post_processor.py` | tag cache restore | | `rag/svr/task_executor_refactor/task_context.py` | language default fix | | `test/.../conftest.py` | +294 lines shared helpers | | `test/.../*.py` | 15 test files refactored, 20 new tests | --------- Co-authored-by: Claude Opus 4.8 --- CLAUDE.md | 64 ++- rag/svr/task_executor.py | 1 + .../chunk_post_processor.py | 10 +- .../task_executor_refactor/chunk_service.py | 3 + .../dataflow_service.py | 19 +- .../embedding_service.py | 35 +- .../task_executor_refactor/raptor_service.py | 9 +- .../task_executor_refactor/task_context.py | 2 +- .../task_executor_refactor/task_handler.py | 68 ++- .../svr/task_executor_refactor/conftest.py | 296 +++++++++- .../test_chunk_builder.py | 103 +--- .../test_chunk_post_processor.py | 524 ++++++++---------- .../test_chunk_service.py | 199 ++----- .../task_executor_refactor/test_comparator.py | 15 +- .../test_dataflow_service.py | 140 ++--- .../test_embedding_service.py | 179 +++++- .../test_post_processor.py | 52 ++ .../test_raptor_service.py | 163 +++--- .../test_task_context.py | 4 +- .../test_task_handler.py | 282 +++++----- .../test_task_handler_integration.py | 371 ++++--------- 21 files changed, 1317 insertions(+), 1222 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 7cb61ad126..2302d23de0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -6,44 +6,62 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co RAGFlow is an open-source RAG (Retrieval-Augmented Generation) engine based on deep document understanding. It's a full-stack application with: -- Python backend (Flask-based API server) +- Python backend (Quart-based async API server — Quart is the async reimplementation of Flask) - React/TypeScript frontend (built with vitejs) -- Microservices architecture with Docker deployment -- Multiple data stores (MySQL, Elasticsearch/Infinity, Redis, MinIO) +- Background task executor workers (separate Python processes, Redis-queue-driven) +- Peewee ORM for database models (not SQLAlchemy) +- Multiple data stores (MySQL/PostgreSQL, Elasticsearch/Infinity/OpenSearch/OceanBase, Redis, MinIO) ## Architecture -### Backend (`/api/`) +### Runtime Architecture -- **Main Server**: `api/ragflow_server.py` - Flask application entry point -- **Apps**: Modular Flask blueprints in `api/apps/` for different functionalities: - - `kb_app.py` - Knowledge base management - - `dialog_app.py` - Chat/conversation handling - - `document_app.py` - Document processing - - `canvas_app.py` - Agent workflow canvas - - `file_app.py` - File upload/management -- **Services**: Business logic in `api/db/services/` -- **Models**: Database models in `api/db/db_models.py` +RAGFlow runs as **two separate Python process types**, orchestrated by `docker/launch_backend_service.sh`: + +- **API Server** (`api/ragflow_server.py`): Quart-based async HTTP server +- **Task Executors** (`rag/svr/task_executor.py`): Background workers processing documents from Redis streams. Multiple instances run in parallel (controlled by `WS` env var). Each consumes from priority-ordered Redis streams (`te.1.common`, `te.0.common`), using consumer groups for load distribution. + +Key consequence: task executors import a different code surface than the API server, so always check which process a module is meant for. + +### Backend API (`/api/`) + +- **App factory**: `api/apps/__init__.py` — creates the Quart app, configures auth (`login_required` decorator, JWT + API token + session fallback), and dynamically discovers/registers blueprints +- **Two API coexisting patterns**: + - **RESTful APIs** in `api/apps/restful_apis/` — newer pattern with Pydantic request validation, service layer in `api/apps/services/`, routes registered under `/api/v1` + - **Legacy APIs** in `api/apps/*_app.py` — older pattern using `@validate_request()`, routes registered under `/v1/` + - **SDK APIs** in `api/apps/sdk/` — registered under `/v1/` +- **Services**: `api/db/services/` — business logic wrapping Peewee model operations. `api/apps/services/` — service layer for the RESTful APIs +- **Models**: `api/db/db_models.py` — Peewee ORM models with pooled MySQL/PostgreSQL connections, custom `JSONField`/`ListField` types, retry logic on connection loss ### Core Processing (`/rag/`) -- **Document Processing**: `deepdoc/` - PDF parsing, OCR, layout analysis -- **LLM Integration**: `rag/llm/` - Model abstractions for chat, embedding, reranking -- **RAG Pipeline**: `rag/flow/` - Chunking, parsing, tokenization -- **Graph RAG**: `rag/graphrag/` - Knowledge graph construction and querying +- **Document ingestion pipeline**: `rag/flow/pipeline.py` — `Pipeline` (extends `agent.canvas.Graph`) orchestrates the ingestion DAG. Components: File (fetches binary from storage), Parser (dispatches to `deepdoc.parser` based on file type), TokenChunker/TitleChunker (splits into chunks), Tokenizer (computes full-text tokens + embedding vectors), Extractor (LLM-based extraction). Data flows via Pydantic `*FromUpstream` schemas. +- **Document parsing**: `deepdoc/` — PDF parsing (vision-based OCR, layout analysis, table structure recognition) and format-specific parsers (DOCX, XLSX, PPT, Markdown, HTML, images). All parsers normalize to a common structure (list of bbox dicts for PDFs, `{text, doc_type_kwd}` for others). +- **LLM Integration**: `rag/llm/` — factory pattern with runtime class discovery. `chat_model.py` (30+ providers via OpenAI SDK and LiteLLM wrappers), `embedding_model.py`, `rerank_model.py`, `cv_model.py` (image-to-text), `sequence2txt_model.py` (ASR), `tts_model.py`. Use `LLMBundle` (from `api.db.services.llm_service`) as the unified interface. +- **Graph RAG**: `rag/graphrag/` — multi-phase pipeline: per-document subgraph extraction (LLM or spaCy NER), Leiden community detection, entity resolution, community summarization. Entities/relations/reports are indexed as chunks alongside regular text chunks, differentiated by `knowledge_graph_kwd`. +- **Search**: `rag/nlp/search.py` — `Dealer` class combines vector similarity + BM25 + re-ranking. `KGSearch` extends it for graph-aware retrieval (entity resolution, n-hop enrichment). ### Agent System (`/agent/`) -- **Components**: Modular workflow components (LLM, retrieval, categorize, etc.) -- **Templates**: Pre-built agent workflows in `agent/templates/` -- **Tools**: External API integrations (Tavily, Wikipedia, SQL execution, etc.) +- **Execution engine**: `agent/canvas.py` — `Canvas` (extends `Graph`) executes the DAG. Components are run in topological order via `_run_batch`, each receiving upstream outputs as kwargs. Control-flow components (`Categorize`, `Switch`, `Iteration`, `Loop`) dynamically modify the execution path. +- **Component base**: `agent/component/base.py` — `ComponentBase` with `invoke(**kwargs)` / `invoke_async(**kwargs)` lifecycle. Variable references (`{component_id@output_var}` or `{sys.query}`) are resolved from the canvas graph at runtime. +- **Components**: Modular workflow components in `agent/component/` — Begin, LLM, Agent (tool-calling LLM), Categorize, Switch, Iteration, Loop, Message, Invoke (HTTP), and data manipulation nodes. Auto-discovered by `__init__.py`. +- **Templates**: Pre-built agent workflows as JSON DSL files in `agent/templates/`. Each contains a complete `components` DAG, `path`, and `globals`. +- **Tools**: `agent/tools/` — Retrieval, web search (DuckDuckGo, Google, Tavily, SearXNG), academic search (ArXiv, PubMed, Google Scholar, Wikipedia), code execution, SQL execution, email, GitHub, finance data, translation, weather. Tools implement `ToolBase` (extends `ComponentBase`) and produce OpenAI-compatible function descriptors. +- **Plugins**: `agent/plugin/` — plugin system using `pluginlib` for loading external LLM tool plugins from `embedded_plugins/`. ### Frontend (`/web/`) - React/TypeScript with vitejs framework -- shadcn/ui components -- State management with Zustand -- Tailwind CSS for styling +- shadcn/ui components (Radix UI primitives + Tailwind CSS) +- `@tanstack/react-query` for server state (cache keys, mutations, invalidation) +- Zustand for local state (primarily agent canvas graph store) +- `react-router` v7 with lazy-loaded pages +- `react-i18next` for i18n (17 languages) +- Axios for HTTP with a layered pattern: endpoint definitions (`utils/api.ts`) → HTTP client (`utils/next-request.ts`) → service layer (`services/`) → query hooks (`hooks/use-*-request.ts`) → components +- `@xyflow/react` for the agent workflow canvas +- `react-hook-form` + `zod` for form validation +- Two API proxy prefixes: `webAPI = '/v1'` (legacy) and `restAPIv1 = '/api/v1'` (RESTful) ## Common Development Commands diff --git a/rag/svr/task_executor.py b/rag/svr/task_executor.py index 3a4290d579..bf162dd4b8 100644 --- a/rag/svr/task_executor.py +++ b/rag/svr/task_executor.py @@ -1728,6 +1728,7 @@ async def handle_task(): 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', '')}") + set_recording_context(NullRecordingContext()) await TaskManager.run_refactored_task(task, chat_limiter, minio_limiter, chunk_limiter, embed_limiter,kg_limiter, set_progress, has_canceled) else: # original version diff --git a/rag/svr/task_executor_refactor/chunk_post_processor.py b/rag/svr/task_executor_refactor/chunk_post_processor.py index 5677c9bf53..d9f0f11bf9 100644 --- a/rag/svr/task_executor_refactor/chunk_post_processor.py +++ b/rag/svr/task_executor_refactor/chunk_post_processor.py @@ -41,7 +41,7 @@ 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_from_provider_instance from rag.prompts.generator import gen_metadata, keyword_extraction, question_proposal, content_tagging -from rag.graphrag.utils import get_llm_cache, set_llm_cache +from rag.graphrag.utils import get_llm_cache, set_llm_cache, get_tags_from_cache, set_tags_to_cache async def extract_keywords(docs: List[Dict], ctx: TaskContext) -> None: @@ -243,10 +243,14 @@ async def apply_tags(docs: List[Dict], ctx: TaskContext) -> None: S = 1000 st = timer() examples = [] - all_tags = settings.retriever.all_tags_in_portion(tenant_id, kb_ids, S) + all_tags = get_tags_from_cache(kb_ids) + if not all_tags: + all_tags = settings.retriever.all_tags_in_portion(tenant_id, kb_ids, S) + set_tags_to_cache(kb_ids, all_tags) + else: + all_tags = json.loads(all_tags) chat_model_config = get_model_config_from_provider_instance(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): diff --git a/rag/svr/task_executor_refactor/chunk_service.py b/rag/svr/task_executor_refactor/chunk_service.py index 060a99ed5f..d0fdd836f5 100644 --- a/rag/svr/task_executor_refactor/chunk_service.py +++ b/rag/svr/task_executor_refactor/chunk_service.py @@ -34,6 +34,7 @@ from typing import Any, Dict, List import xxhash from common import settings +from common.connection_utils import timeout from common.constants import PAGERANK_FLD, TAG_FLD from common.misc_utils import thread_pool_exec from common.float_utils import normalize_overlapped_percent @@ -85,6 +86,7 @@ class ChunkService: """ self._task_context = ctx + @timeout(60 * 80, 1) async def build_chunks( self, storage_binary: bytes, @@ -190,6 +192,7 @@ class ChunkService: st = timer() + @timeout(60) async def upload_to_minio(document, chunk): try: d = copy.deepcopy(document) diff --git a/rag/svr/task_executor_refactor/dataflow_service.py b/rag/svr/task_executor_refactor/dataflow_service.py index d67c6cb266..9cf345a300 100644 --- a/rag/svr/task_executor_refactor/dataflow_service.py +++ b/rag/svr/task_executor_refactor/dataflow_service.py @@ -39,6 +39,7 @@ 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_from_provider_instance +from common.connection_utils import timeout from common.constants import LLMType, PipelineTaskType from common.metadata_utils import update_metadata_to from common.misc_utils import thread_pool_exec @@ -102,9 +103,10 @@ class DataflowService: task_dataset_id = ctx.kb_id # Load DSL - dsl = await self._load_dsl(dataflow_id) + dsl, corrected_id = await self._load_dsl(dataflow_id) if dsl is None: return + dataflow_id = corrected_id # Run pipeline pipeline = Pipeline( @@ -193,17 +195,23 @@ class DataflowService: await self._billing_hook.on_pipeline_error() raise - async def _load_dsl(self, dataflow_id: str) -> Optional[str]: - """Load dataflow DSL from service.""" + async def _load_dsl(self, dataflow_id: str) -> tuple: + """Load dataflow DSL from service. + + Returns: + Tuple of (dsl, corrected_dataflow_id). + When task_type is not 'dataflow', the dataflow_id is corrected + from the pipeline log's pipeline_id. + """ 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 + return cvs.dsl, dataflow_id else: e, pipeline_log = PipelineOperationLogService.get_by_id(dataflow_id) assert e, "Pipeline log not found." - return pipeline_log.dsl + return pipeline_log.dsl, pipeline_log.pipeline_id @staticmethod def _get_output_type(chunks: Dict) -> str: @@ -235,6 +243,7 @@ class DataflowService: return [{"text": [chunks["html"]]}] if chunks["html"] else [] return [] + @timeout(60) async def _embed_chunks( self, chunks: List[Dict], token_consumption: int ) -> Tuple[Optional[List[Dict]], int]: diff --git a/rag/svr/task_executor_refactor/embedding_service.py b/rag/svr/task_executor_refactor/embedding_service.py index 35ddde7002..a22c733fc3 100644 --- a/rag/svr/task_executor_refactor/embedding_service.py +++ b/rag/svr/task_executor_refactor/embedding_service.py @@ -19,11 +19,12 @@ 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 common.misc_utils import thread_pool_exec +from common.token_utils import truncate from rag.svr.task_executor_refactor.embedding_utils import EmbeddingUtils from rag.svr.task_executor_refactor.task_context import TaskContext @@ -54,7 +55,7 @@ class EmbeddingService: self._embedding_batch_size = embedding_batch_size or settings.EMBEDDING_BATCH_SIZE - def embed_chunks( + async def embed_chunks( self, docs: List[Dict[str, Any]], embedding_model, @@ -79,7 +80,8 @@ class EmbeddingService: # 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) + async with self._task_context.embed_limiter: + vts, c = await thread_pool_exec(embedding_model.encode, titles[0:1]) tts = np.tile(vts[0], (len(contents), 1)) tk_count += c else: @@ -89,7 +91,12 @@ class EmbeddingService: 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) + async with self._task_context.embed_limiter: + vts, c = await thread_pool_exec( + self._batch_encode_wrapper, + [truncate(t, embedding_model.max_length - 10) for t in batch], + embedding_model, + ) vects_batches.append(vts) tk_count += c if self._task_context.progress_cb: @@ -109,19 +116,7 @@ class EmbeddingService: 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()) + @staticmethod + def _batch_encode_wrapper(txts: List[str], embedding_model) -> Tuple[np.ndarray, int]: + """Synchronous wrapper for batch encoding — used with thread_pool_exec.""" + return embedding_model.encode(txts) diff --git a/rag/svr/task_executor_refactor/raptor_service.py b/rag/svr/task_executor_refactor/raptor_service.py index ba80bac3a5..c1695b9958 100644 --- a/rag/svr/task_executor_refactor/raptor_service.py +++ b/rag/svr/task_executor_refactor/raptor_service.py @@ -31,6 +31,7 @@ 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.connection_utils import timeout from common.constants import PAGERANK_FLD from common.misc_utils import thread_pool_exec from common.token_utils import num_tokens_from_string @@ -68,6 +69,7 @@ class RaptorService: """ self._task_context = ctx + @timeout(3600) async def run_raptor_for_kb( self, kb_parser_config: Dict, @@ -166,8 +168,8 @@ class RaptorService: 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)) + 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: @@ -370,7 +372,8 @@ class RaptorService: 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" + assert chunks, "_generate_raptor must not be called with empty chunks" + vctr_nm = "q_%d_vec" % len(chunks[0][1]) raptor = Raptor( raptor_config.get("max_cluster", 64), diff --git a/rag/svr/task_executor_refactor/task_context.py b/rag/svr/task_executor_refactor/task_context.py index 2a15d6b50c..8d9f41c4db 100644 --- a/rag/svr/task_executor_refactor/task_context.py +++ b/rag/svr/task_executor_refactor/task_context.py @@ -224,7 +224,7 @@ class TaskContext: "parser_id": "", "parser_config": {}, "kb_parser_config": {}, - "language": "en", + "language": "Chinese", "llm_id": "", "embd_id": "", "from_page": 0, diff --git a/rag/svr/task_executor_refactor/task_handler.py b/rag/svr/task_executor_refactor/task_handler.py index ebcd92c039..a22d407de4 100644 --- a/rag/svr/task_executor_refactor/task_handler.py +++ b/rag/svr/task_executor_refactor/task_handler.py @@ -39,6 +39,7 @@ 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.connection_utils import timeout 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 @@ -111,6 +112,7 @@ class TaskHandler: logging.exception( f"Remove doc({task_doc_id}) from docStore failed when task({task_id}) canceled, exception: {e}") + @timeout(60 * 60 * 3, 1) async def handle(self) -> None: """Handle a document processing task.""" ctx = self._task_context @@ -125,14 +127,6 @@ class TaskHandler: 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 @@ -140,15 +134,25 @@ class TaskHandler: 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: + # Language defaults to "Chinese" via TaskContext._DEFAULTS — safe to bind model directly. + # Bind embedding model (matching original do_handle_task order: bind + init_kb before routing) + result = await self._bind_embedding_model() + if result is None: return + embedding_model, vector_size = result with embedding_model: - vector_size = self._get_vector_size(embedding_model) self._init_kb(vector_size) + # Handle dataflow tasks (after init_kb, matching original behavior) + 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 + # Route to appropriate handler if task_type == "raptor": await self._run_raptor(embedding_model, vector_size) @@ -166,12 +170,6 @@ class TaskHandler: 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 @@ -203,8 +201,12 @@ class TaskHandler: ctx = self._task_context ctx.progress_cb(1, "Clone task placeholder") - async def _bind_embedding_model(self) -> Optional[LLMBundle]: - """Bind embedding model to task.""" + async def _bind_embedding_model(self) -> Optional[tuple]: + """Bind embedding model to task. + + Returns: + Tuple of (embedding_model, vector_size) on success, or None on failure. + """ ctx = self._task_context task_tenant_id = ctx.tenant_id task_embedding_id = ctx.embd_id @@ -221,7 +223,7 @@ class TaskHandler: ) embedding_model = LLMBundle(task_tenant_id, embd_model_config, lang=task_language) vts, _ = embedding_model.encode(["ok"]) - return embedding_model + return embedding_model, len(vts[0]) except Exception as e: error_message = f'Fail to bind embedding model: {str(e)}' ctx.progress_cb(-1, msg=error_message) @@ -340,8 +342,24 @@ class TaskHandler: kb_parser_config.update({ "graphrag": { "use_graphrag": True, - "entity_types": ["organization", "person", "geo", "event", "category"], + "entity_types": [ + "organization", + "person", + "geo", + "event", + "category", + ], "method": "light", + "batch_chunk_token_size": 4096, + "retry_attempts": 2, + "retry_backoff_seconds": 2.0, + "retry_backoff_max_seconds": 60.0, + "build_subgraph_timeout_per_chunk_seconds": 300, + "build_subgraph_min_timeout_seconds": 600, + "merge_timeout_seconds": 180, + "resolution_timeout_seconds": 1800, + "community_timeout_seconds": 1800, + "lock_acquire_timeout_seconds": 600, } }) if ctx.write_interceptor: @@ -400,6 +418,10 @@ class TaskHandler: # Get storage binary bucket, name = File2DocumentService.get_storage_address(doc_id=ctx.doc_id) binary = await self._get_storage_binary(bucket, name) + if binary is None: + raise FileNotFoundError( + f"Can not find file <{ctx.name}> from minio. Could you try it again." + ) chunks = await chunk_service.build_chunks(binary) ctx.recording_context.record("chunks", chunks) @@ -418,7 +440,7 @@ class TaskHandler: start_ts = timer() embedding_service = EmbeddingService(ctx=ctx) try: - token_count, vector_size = embedding_service.embed_chunks( + token_count, vector_size = await embedding_service.embed_chunks( chunks, embedding_model, ctx.parser_config ) except TaskCanceledException: diff --git a/test/unit_test/rag/svr/task_executor_refactor/conftest.py b/test/unit_test/rag/svr/task_executor_refactor/conftest.py index 64da81a0db..8f57701d04 100644 --- a/test/unit_test/rag/svr/task_executor_refactor/conftest.py +++ b/test/unit_test/rag/svr/task_executor_refactor/conftest.py @@ -272,6 +272,7 @@ class MockChatModel: def __init__(self): self.llm_name = "mock_chat" + self.max_length = 4096 def __enter__(self): return self @@ -469,26 +470,293 @@ def mock_chunk_service_factory(): # ============================================================================= -# RaptorService Fixtures +# Unified Mock TaskContext Factory +# ============================================================================= + +def make_task_context(**overrides): + """Build a MagicMock TaskContext with sensible defaults for all services. + + Every test file that needs a mock ``TaskContext`` should use this factory + with keyword-only overrides instead of defining its own ``_create_mock_context``. + + Usage:: + + ctx = make_task_context(parser_id="table", kb_parser_config={"tag_kb_ids": ["kb_1"]}) + """ + defaults = { + "id": "task_1", + "tenant_id": "tenant_1", + "kb_id": "kb_1", + "doc_id": "doc_1", + "name": "test.pdf", + "location": "/path/to/test.pdf", + "language": "en", + "parser_id": "naive", + "parser_config": {}, + "kb_parser_config": {}, + "llm_id": "llm_1", + "embd_id": "embd_1", + "from_page": 0, + "to_page": -1, + "size": 1000, + "pagerank": 0, + "task_type": "standard", + "dataflow_id": "", + "doc_ids": [], + "file": None, + "memory_id": "", + "source_id": "", + "message_dict": {}, + } + ctx = MagicMock() + for k, v in defaults.items(): + setattr(ctx, k, v if k not in overrides else overrides.pop(k)) + + # Callbacks + ctx.progress_cb = MagicMock() + ctx.has_canceled_func = MagicMock(return_value=False) + ctx.recording_context = MagicMock() + ctx.write_interceptor = None + + # Raw task dict — derive from context attributes + ctx.raw_task = MagicMock() + + # Limiters — all use AsyncMockLimiter so services that acquire them work + limiter = AsyncMockLimiter() + ctx.chunk_limiter = limiter + ctx.chat_limiter = limiter + ctx.embed_limiter = limiter + ctx.kg_limiter = limiter + ctx.minio_limiter = limiter + + # Apply remaining overrides + for k, v in overrides.items(): + setattr(ctx, k, v) + + return ctx + + +# ============================================================================= +# RaptorService Fixtures (kept for backward compatibility) # ============================================================================= 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 + return make_task_context() @pytest.fixture def mock_raptor_context(): """Provide a mock TaskContext for RaptorService tests.""" - return create_mock_raptor_context() + return make_task_context() + + +# ============================================================================= +# Embedding Binding Patch Helper +# ============================================================================= + +class patch_embedding_binding: + """Context manager that patches embedding model binding at the external boundary. + + Patches ``LLMBundle``, ``get_model_config_from_provider_instance``, and + ``get_tenant_default_model_by_type`` so that ``TaskHandler._bind_embedding_model`` + executes its real logic without making actual API calls. + + Usage:: + + with patch_embedding_binding(vector_size=128): + handler = TaskHandler(ctx) + await handler.handle() + """ + + def __init__(self, vector_size: int = 128): + self._vector_size = vector_size + self._patches = [] + + def __enter__(self): + mock_model = MagicMock() + mock_model.encode = MagicMock( + return_value=( + np.random.rand(1, self._vector_size).astype(np.float32), + 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) + + self._patches = [ + patch( + "rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance", + 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(), + ), + ] + for p in self._patches: + p.__enter__() + return self + + def __exit__(self, *args): + for p in reversed(self._patches): + p.__exit__(*args) + + +# ============================================================================= +# Common mock callbacks +# ============================================================================= + +async def mock_thread_return_binary(func, *args, **kwargs): + """Reusable mock for thread_pool_exec — returns fake binary.""" + return b"fake pdf binary" + + +async def mock_thread_return_none(func, *args, **kwargs): + """Reusable mock for thread_pool_exec — returns None.""" + return None + + +# ============================================================================= +# Patch helpers for integration tests +# ============================================================================= + +def patch_get_storage_binary(): + """Patch TaskHandler._get_storage_binary to return fake binary.""" + return patch("rag.svr.task_executor_refactor.task_handler.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) + + +# ============================================================================= +# Shared Task Dictionary Factory +# ============================================================================= + +def make_task_dict(**overrides): + """Build a task dict with sensible defaults for integration tests. + + All ``_create_standard_task_dict`` / ``_create_raptor_task_dict`` / etc. + helpers in integration tests should be replaced with this single factory. + + Usage:: + + task_dict = make_task_dict(task_type="raptor", doc_ids=["doc1"]) + """ + 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, + **overrides, + } + + +# ============================================================================= +# Shared Pipeline Mock Block for Integration Tests +# ============================================================================= + +class patch_pipeline_mocks: + """Context manager bundling common integration-test mock blocks. + + Patches external boundaries so ``TaskHandler.handle()`` executes without + actual API calls. Use ``mode="raptor"`` or ``mode="graphrag"``. + + Usage:: + + with patch_pipeline_mocks() as m: + m.get_model_config_from_provider_instance.return_value = MagicMock() + handler = TaskHandler(ctx) + await handler.handle() + """ + + _MODULES = { + "task_handler": "rag.svr.task_executor_refactor.task_handler", + "chunk_service": "rag.svr.task_executor_refactor.chunk_service", + } + + # (module_key, attr_name, use_AsyncMock) + _COMMON = [ + ("task_handler", "get_model_config_from_provider_instance", False), + ("task_handler", "LLMBundle", False), + ("task_handler", "get_tenant_default_model_by_type", False), + ("task_handler", "search.index_name", False), + ("task_handler", "thread_pool_exec", False), + ("task_handler", "DocumentService", False), + ] + + _STANDARD = [ + ("task_handler", "File2DocumentService", False), + ("chunk_service", "thread_pool_exec", False), + ("task_handler", "ChunkService", False), + ] + + _RAPTOR = [ + ("task_handler", "KnowledgebaseService", False), + ("task_handler", "RaptorService", False), + ("task_handler", "ChunkService", False), + ] + + _GRAPH_RAG = [ + ("task_handler", "KnowledgebaseService", False), + ("task_handler", "run_graphrag_for_kb", True), + ] + + def __init__(self, mode: str = "standard"): + self._mode = mode + self._stack = None + + def __enter__(self): + import contextlib + from unittest.mock import patch, MagicMock, AsyncMock + + prefixes = list(self._COMMON) + if self._mode == "standard": + prefixes += self._STANDARD + elif self._mode == "raptor": + prefixes += self._RAPTOR + elif self._mode == "graphrag": + prefixes += self._GRAPH_RAG + + mocks = MagicMock() + ctx_managers = [] + for mod_key, attr, use_async in prefixes: + target = f"{self._MODULES[mod_key]}.{attr}" + if use_async: + cm = patch(target, new_callable=AsyncMock) + else: + cm = patch(target) + + mock_handle = cm.__enter__() + setattr(mocks, attr.replace(".", "_"), mock_handle) + ctx_managers.append(cm) + + self._stack = contextlib.ExitStack() + self._ctx_managers = ctx_managers + return mocks + + def __exit__(self, *args): + for cm in reversed(self._ctx_managers): + cm.__exit__(*args) 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 index 1c4f844e8f..bc2eed76de 100644 --- 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 @@ -18,12 +18,13 @@ Unit tests for ChunkBuilder module. """ import pytest -from unittest.mock import MagicMock, patch, AsyncMock +from unittest.mock import MagicMock, patch from rag.svr.task_executor_refactor.chunk_builder import ( get_parser, run_chunking, extract_outline, ) +from test.unit_test.rag.svr.task_executor_refactor.conftest import make_task_context class TestGetParser: @@ -49,151 +50,99 @@ class TestGetParser: 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() + ctx = make_task_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 = make_task_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: + with patch("rag.svr.task_executor_refactor.chunk_builder.thread_pool_exec") as mock_thread, \ + patch("rag.svr.task_executor_refactor.chunk_builder.merge_table_parser_config_from_kb") as mock_merge: 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) + 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() + ctx = make_task_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() + @staticmethod + def _ctx(recording_ctx=None, **overrides): + ctx = make_task_context(**overrides) + ctx.recording_context = recording_ctx or 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() + ctx = self._ctx() 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 + ctx.recording_context.record.assert_called_with("outline_data", outline_data) 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() - + ctx = self._ctx() 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) + ctx.recording_context.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 - + ctx = self._ctx() await extract_outline([], ctx) - - mock_rec_ctx.record.assert_called_with("outline_data", None) + ctx.recording_context.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() + ctx = self._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" ) @@ -201,19 +150,13 @@ class TestExtractOutline: @pytest.mark.asyncio async def test_extract_outline_persistence_exception(self): """Test outline extraction handles persistence exception.""" - ctx = self._create_mock_context() + ctx = self._ctx() 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) + ctx.recording_context.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 index 015016c4a5..a5e5427d86 100644 --- 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 @@ -15,10 +15,16 @@ """ Unit tests for ChunkPostProcessor module. + +Mock strategy: the LLM prompt functions (``keyword_extraction``, ``question_proposal``, +``gen_metadata``, ``content_tagging``) are mocked since they make actual LLM API +calls. ``get_llm_cache`` / ``set_llm_cache`` run as real code, so cache +population and retrieval are exercised. ``rag_tokenizer`` is mocked because +it requires NLTK data in the test environment. """ import pytest -from unittest.mock import MagicMock, patch, AsyncMock +from unittest.mock import MagicMock, patch from rag.svr.task_executor_refactor.chunk_post_processor import ( extract_keywords, generate_questions, @@ -27,394 +33,314 @@ from rag.svr.task_executor_refactor.chunk_post_processor import ( count_with_key, build_metadata_config, ) +from test.unit_test.rag.svr.task_executor_refactor.conftest import ( + make_task_context, + MockChatModel, +) -class TestExtractKeywords: +class _BasePostProcessorTest: + """Shared helpers for post-processor test classes.""" + + @staticmethod + def _mock_llm_binding(chat_model_cls=MockChatModel): + """Patch model config lookup + LLMBundle to return a MockChatModel.""" + p1 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance", + return_value=MagicMock()) + p2 = patch("rag.svr.task_executor_refactor.chunk_post_processor.LLMBundle", + return_value=chat_model_cls()) + return p1, p2 + + @staticmethod + def _patch_prompt_func(func_path: str, return_value): + """Patch a prompt-level LLM function (the actual API call).""" + return patch(func_path, return_value=return_value) + + +class TestExtractKeywords(_BasePostProcessorTest): """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() + """Test successful keyword extraction — cache miss → LLM prompt runs.""" + ctx = make_task_context(parser_config={"auto_keywords": 5}) 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_from_provider_instance") 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] + p1, p2 = self._mock_llm_binding() + p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache", + return_value=None) # cache miss + p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache") # Redis stub + p5 = self._patch_prompt_func( + "rag.svr.task_executor_refactor.chunk_post_processor.keyword_extraction", + return_value="keyword1, keyword2", + ) + p6 = patch("rag.svr.task_executor_refactor.chunk_post_processor.rag_tokenizer") + with p1, p2, p3, p4, p5, p6 as mock_tok: + mock_tok.tokenize.return_value = "keyword1 keyword2" + await extract_keywords(docs, ctx) + 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) + ctx = make_task_context(parser_config={"auto_keywords": 5}, + 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_from_provider_instance") 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] + p1, p2 = self._mock_llm_binding() + with p1, p2: + await extract_keywords(docs, ctx) + 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() + ctx = make_task_context(parser_config={"auto_keywords": 5}) docs = [] - with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance") 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() + p1, p2 = self._mock_llm_binding() + with p1, p2: + await extract_keywords(docs, ctx) + ctx.progress_cb.assert_called() -class TestGenerateQuestions: +class TestGenerateQuestions(_BasePostProcessorTest): """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"}, - ] + """Test successful question generation — cache miss → LLM prompt runs.""" + ctx = make_task_context(parser_config={"auto_questions": 3}) + docs = [{"content_with_weight": "This is test content one"}] - with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance") 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] + p1, p2 = self._mock_llm_binding() + p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache", + return_value=None) + p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache") + p5 = self._patch_prompt_func( + "rag.svr.task_executor_refactor.chunk_post_processor.question_proposal", + return_value="Question 1\nQuestion 2", + ) + p6 = patch("rag.svr.task_executor_refactor.chunk_post_processor.rag_tokenizer") + with p1, p2, p3, p4, p5, p6 as mock_tok: + mock_tok.tokenize.return_value = "Question 1 Question 2" + await generate_questions(docs, ctx) + 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) + ctx = make_task_context(parser_config={"auto_questions": 3}, + 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_from_provider_instance") 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] + p1, p2 = self._mock_llm_binding() + with p1, p2: + await generate_questions(docs, ctx) + assert "question_kwd" not in docs[0] -class TestGenerateMetadata: +class TestGenerateMetadata(_BasePostProcessorTest): """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"}}, - ] + """Test successful metadata generation — cache miss → LLM prompt runs.""" + ctx = make_task_context( + parser_config={ + "enable_metadata": True, + "metadata": [{"name": "category", "type": "string"}], + "built_in_metadata": ["author", "date"], + }, + ) + docs = [{"content_with_weight": "This is test content"}] - with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance") 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() + p1, p2 = self._mock_llm_binding() + p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache", + return_value=None) + p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache") + p5 = self._patch_prompt_func( + "rag.svr.task_executor_refactor.chunk_post_processor.gen_metadata", + return_value={"category": "test"}, + ) + p6 = patch("rag.svr.task_executor_refactor.chunk_post_processor.DocMetadataService") + with p1, p2, p3, p4, p5, p6 as mock_meta: + mock_meta.get_document_metadata.return_value = {} + mock_meta.update_document_metadata = MagicMock() + await generate_metadata(docs, ctx) + 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"}}, - ] + ctx = make_task_context( + parser_config={ + "enable_metadata": True, + "metadata": [{"name": "category", "type": "string"}], + "built_in_metadata": ["author", "date"], + }, + write_interceptor=MagicMock(), + ) + docs = [{"content_with_weight": "This is test content"}] - with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance") 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" - ) + p1, p2 = self._mock_llm_binding() + p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache", + return_value=None) + p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache") + p5 = self._patch_prompt_func( + "rag.svr.task_executor_refactor.chunk_post_processor.gen_metadata", + return_value={"category": "test"}, + ) + p6 = patch("rag.svr.task_executor_refactor.chunk_post_processor.DocMetadataService") + with p1, p2, p3, p4, p5, p6: + await generate_metadata(docs, ctx) + ctx.write_interceptor.intercept.assert_called_once_with( + "DocMetadataService.update_document_metadata" + ) -class TestApplyTags: +class TestApplyTags(_BasePostProcessorTest): """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"}, - ] + """Test successful tag application with tag cache miss.""" + ctx = make_task_context( + kb_parser_config={"tag_kb_ids": ["kb_1"], "topn_tags": 3}, + ) + docs = [{"content_with_weight": "This is test content"}] - with patch("rag.svr.task_executor_refactor.chunk_post_processor.get_model_config_from_provider_instance") 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 + p1, p2 = self._mock_llm_binding() + p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.settings") + p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache", + return_value='{"tag1": 1}') # cache hit → skip LLM + p5 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache") + p6 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_tags_from_cache", + return_value=None) + p7 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_tags_to_cache") + with p1, p2, p3 as mock_settings, p4, p5, p6 as mock_get_tags, p7 as mock_set_tags: + mock_settings.retriever.all_tags_in_portion.return_value = {"tag1": 10, "tag2": 5} + mock_settings.retriever.tag_content.return_value = True + await apply_tags(docs, ctx) + assert len(docs) == 1 + mock_get_tags.assert_called_once() + mock_set_tags.assert_called_once() @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"}, - ] + ctx = make_task_context( + kb_parser_config={"tag_kb_ids": ["kb_1"], "topn_tags": 3}, + 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_from_provider_instance") as mock_config: - mock_config.return_value = MagicMock() + p1, p2 = self._mock_llm_binding() + p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.settings") + p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_tags_from_cache", + return_value=None) + p5 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_tags_to_cache") + with p1, p2, p3 as mock_settings, p4, p5: + mock_settings.retriever.all_tags_in_portion.return_value = {"tag1": 10} + await apply_tags(docs, ctx) - 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) + @pytest.mark.asyncio + async def test_apply_tags_tag_cache_miss(self): + """Test apply_tags when get_tags_from_cache returns None (cache miss).""" + ctx = make_task_context( + kb_parser_config={"tag_kb_ids": ["kb_1"], "topn_tags": 3}, + ) + docs = [{"content_with_weight": "This is test content"}] - 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} + p1, p2 = self._mock_llm_binding() + p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.settings") + p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache", + return_value='{"tag1": 1}') # cache hit → skip LLM + p5 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache") + p6 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_tags_from_cache", + return_value=None) + p7 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_tags_to_cache") + with p1, p2, p3 as mock_settings, p4, p5, p6 as mock_get_tags, p7 as mock_set_tags: + mock_settings.retriever.all_tags_in_portion.return_value = {"tag1": 10, "tag2": 5} + mock_settings.retriever.tag_content.return_value = True + await apply_tags(docs, ctx) + mock_get_tags.assert_called_once_with(["kb_1"]) + mock_set_tags.assert_called_once() + mock_settings.retriever.all_tags_in_portion.assert_called_once() - await apply_tags(docs, ctx) + @pytest.mark.asyncio + async def test_apply_tags_tag_cache_hit(self): + """Test apply_tags when get_tags_from_cache returns valid data (cache hit).""" + ctx = make_task_context( + kb_parser_config={"tag_kb_ids": ["kb_1"], "topn_tags": 3}, + ) + docs = [{"content_with_weight": "This is test content"}] - # Should return early due to cancellation + p1, p2 = self._mock_llm_binding() + p3 = patch("rag.svr.task_executor_refactor.chunk_post_processor.settings") + p4 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_llm_cache", + return_value='{"tag1": 1}') # cache hit → skip LLM + p5 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_llm_cache") + p6 = patch("rag.svr.task_executor_refactor.chunk_post_processor.get_tags_from_cache", + return_value='{"cached_tag": 10}') + p7 = patch("rag.svr.task_executor_refactor.chunk_post_processor.set_tags_to_cache") + with p1, p2, p3 as mock_settings, p4, p5, p6 as mock_get_tags, p7 as mock_set_tags: + mock_settings.retriever.tag_content.return_value = True + await apply_tags(docs, ctx) + mock_get_tags.assert_called_once_with(["kb_1"]) + mock_set_tags.assert_not_called() + mock_settings.retriever.all_tags_in_portion.assert_not_called() 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 + assert count_with_key(docs, "tag") == 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 + assert count_with_key(docs, "tag") == 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 + assert count_with_key(docs, "tag") == 0 def test_count_with_key_empty_docs(self): - """Test counting with empty docs list.""" - result = count_with_key([], "tag") - assert result == 0 + assert count_with_key([], "tag") == 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 + assert count_with_key(docs, "tag") == 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 + assert count_with_key(docs, "tag") == 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": {}} + assert build_metadata_config(parser_config) == {"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": [], @@ -423,38 +349,32 @@ class TestBuildMetadataConfig: 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"}] + assert build_metadata_config(parser_config) == [{"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"}] + assert build_metadata_config(parser_config) == [{"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"}] + assert build_metadata_config(parser_config) == [{"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 + assert build_metadata_config(parser_config) == build_metadata_config(parser_config) def test_missing_metadata_key(self): - """When parser_config has no 'metadata' key, built_in_metadata alone is returned.""" - parser_config = {"built_in_metadata": []} + assert build_metadata_config({"built_in_metadata": []}) == [] + + def test_metadata_is_none(self): + """When metadata is None, built_in_metadata alone is returned.""" + parser_config = {"metadata": None, "built_in_metadata": [{"key": "author"}]} result = build_metadata_config(parser_config) - assert result == [] + assert result == [{"key": "author"}] 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 index 60d937ab29..b8ce477dcb 100644 --- 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 @@ -29,6 +29,7 @@ This test file now focuses on ChunkService-specific functionality: import pytest from unittest.mock import MagicMock, patch, AsyncMock from rag.svr.task_executor_refactor.chunk_service import ChunkService +from test.unit_test.rag.svr.task_executor_refactor.conftest import make_task_context class TestChunkServiceInit: @@ -44,40 +45,11 @@ class TestChunkServiceInit: 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 + ctx = make_task_context(size=1000000000) # Very large size service = ChunkService(ctx=ctx) @@ -95,7 +67,7 @@ class TestChunkServiceBuildChunks: @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) + ctx = make_task_context(size=1000) service = ChunkService(ctx=ctx) @@ -123,138 +95,43 @@ class TestChunkServiceBuildChunks: 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}) - + @pytest.mark.parametrize("task_kwargs,func_path,func_name", [ + ({"parser_config": {"auto_keywords": 5}}, "extract_keywords", "extract_keywords"), + ({"parser_config": {"auto_questions": 3}}, "generate_questions", "generate_questions"), + ({"kb_parser_config": {"tag_kb_ids": ["kb_1"]}}, "apply_tags", "apply_tags"), + ({"parser_config": {"enable_metadata": True, "metadata": [{"name": "category", "type": "string"}]}}, + "generate_metadata", "generate_metadata"), + ]) + async def test_build_chunks_with_post_processing(self, task_kwargs, func_path, func_name): + """Test build_chunks triggers post-processing when configured.""" + ctx = make_task_context(**task_kwargs) service = ChunkService(ctx=ctx) - with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings: + mock_rec_ctx = MagicMock() + ctx.recording_context = mock_rec_ctx + + with patch("rag.svr.task_executor_refactor.chunk_service.settings") as mock_settings, \ + patch("rag.svr.task_executor_refactor.chunk_service.get_parser") as mock_get_parser, \ + patch("rag.svr.task_executor_refactor.chunk_service.run_chunking", new_callable=AsyncMock) as mock_run_chunking, \ + patch("rag.svr.task_executor_refactor.chunk_service.extract_outline", new_callable=AsyncMock), \ + patch.object(service, '_prepare_docs_and_upload', new_callable=AsyncMock) as mock_prepare, \ + patch(f"rag.svr.task_executor_refactor.chunk_service.{func_path}", new_callable=AsyncMock) as mock_fn: 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() + mock_get_parser.return_value = MagicMock() + mock_run_chunking.return_value = [] + mock_prepare.return_value = [{"id": "chunk_1", "content_with_weight": "test"}] + await service.build_chunks(b"test binary") + mock_fn.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() + ctx = make_task_context() service = ChunkService(ctx=ctx) cks = [{"content_with_weight": "test chunk"}] @@ -274,7 +151,7 @@ class TestChunkServicePrepareDocsAndUpload: @pytest.mark.asyncio async def test_prepare_docs_and_upload_with_pagerank(self): """Test document preparation with pagerank.""" - ctx = self._create_mock_context() + ctx = make_task_context() ctx.pagerank = 5 service = ChunkService(ctx=ctx) @@ -293,23 +170,11 @@ class TestChunkServicePrepareDocsAndUpload: 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() + ctx = make_task_context() service = ChunkService(ctx=ctx) chunks = [ @@ -338,7 +203,7 @@ class TestChunkServiceInsertChunks: @pytest.mark.asyncio async def test_insert_chunks_canceled(self): """Test chunk insertion when task is canceled.""" - ctx = self._create_mock_context() + ctx = make_task_context() ctx.has_canceled_func = MagicMock(return_value=True) service = ChunkService(ctx=ctx) @@ -363,7 +228,7 @@ class TestChunkServiceInsertChunks: @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() + ctx = make_task_context() service = ChunkService(ctx=ctx) chunks = [{"id": "chunk_1", "content_with_weight": "test1"}] 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 index 207795d798..307015e5d2 100644 --- a/test/unit_test/rag/svr/task_executor_refactor/test_comparator.py +++ b/test/unit_test/rag/svr/task_executor_refactor/test_comparator.py @@ -138,25 +138,18 @@ class TestComparisonReport: assert "No comparison details" in md -class TestContextComparatorInit: - """Tests for ContextComparator initialization.""" +class TestContextComparatorCompareValue: + """Tests for ContextComparator.compare_value method and initialization.""" def test_init_default_tolerance(self): """Test initialization with default tolerance.""" - comparator = ContextComparator() - assert comparator.float_tolerance == 1e-6 + assert ContextComparator().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.""" + assert ContextComparator(float_tolerance=0.01).float_tolerance == 0.01 def setup_method(self): - """Set up test fixtures.""" self.comparator = ContextComparator() def test_compare_none_values(self): 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 index 1ec8a75063..6be9261356 100644 --- 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 @@ -66,10 +66,11 @@ class TestDataflowServiceRunDataflow: await service.run_dataflow() @pytest.mark.asyncio + @pytest.mark.parametrize("output_key", ["chunks", "json"]) @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.""" + async def test_run_dataflow_with_output_type(self, mock_canvas, mock_pipeline_class, task_context, output_key): + """Test run_dataflow processes output end-to-end (chunks / json).""" task_context._task["task_type"] = "dataflow" task_context._task["dataflow_id"] = "dataflow_test" task_context._task["tenant_id"] = "tenant_test" @@ -79,59 +80,22 @@ class TestDataflowServiceRunDataflow: 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, - } + data = {output_key: [{"text": "content", "content_with_weight": "content"}]} + data["embedding_token_consumption"] = 5 mock_pipeline = MagicMock() - mock_pipeline.run = AsyncMock(return_value=chunks) + mock_pipeline.run = AsyncMock(return_value=data) 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() + with patch.object(DataflowService, '_embed_chunks', new_callable=AsyncMock, + return_value=(data[output_key], 5)), \ + patch.object(DataflowService, '_insert_chunks', new_callable=AsyncMock, return_value=True), \ + patch.object(DataflowService, '_update_document_metadata'), \ + patch.object(DataflowService, '_record_pipeline_log'), \ + patch("api.db.services.document_service.DocumentService.increment_chunk_num"): - # 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() + service = DataflowService(ctx=task_context) + await service.run_dataflow() + DataflowService._insert_chunks.assert_called_once() @pytest.mark.asyncio @patch("rag.svr.task_executor_refactor.dataflow_service.Pipeline") @@ -154,13 +118,11 @@ class TestDataflowServiceRunDataflow: 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() + with patch.object(DataflowService, '_embed_chunks', new_callable=AsyncMock, return_value=(None, 0)), \ + patch.object(DataflowService, '_record_pipeline_log'): + service = DataflowService(ctx=task_context) + await service.run_dataflow() + service._record_pipeline_log.assert_called() @pytest.mark.asyncio @patch("rag.svr.task_executor_refactor.dataflow_service.Pipeline") @@ -190,16 +152,16 @@ class TestDataflowServiceRunDataflow: 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() + with patch.object(DataflowService, '_embed_chunks', new_callable=AsyncMock, return_value=(chunks["chunks"], 1)), \ + patch.object(DataflowService, '_insert_chunks', new_callable=AsyncMock, return_value=True), \ + patch.object(DataflowService, '_update_document_metadata'), \ + patch.object(DataflowService, '_record_pipeline_log'), \ + patch("api.db.services.document_service.DocumentService.increment_chunk_num"): - billing_hook.on_pipeline_success.assert_called_once() - billing_hook.on_pipeline_error.assert_not_called() + 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") @@ -378,4 +340,46 @@ class TestDataflowServiceInit: 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 + assert service._billing_hook is hook + + +class TestDataflowServiceLoadDsl: + """Tests for _load_dsl with dataflow_id correction.""" + + @pytest.mark.asyncio + async def test_load_dsl_for_dataflow_task_type_returns_unchanged_id(self): + """When task_type == 'dataflow', dataflow_id is returned unchanged.""" + ctx = MagicMock() + ctx.task_type = "dataflow" + dataflow_id = "original_dataflow_id" + + with patch("rag.svr.task_executor_refactor.dataflow_service.UserCanvasService") as mock_canvas: + mock_canvas.get_by_id.return_value = (True, MagicMock(dsl='{"id": "test"}')) + service = DataflowService(ctx=ctx) + + dsl, corrected_id = await service._load_dsl(dataflow_id) + + assert dsl == '{"id": "test"}' + assert corrected_id == "original_dataflow_id" + mock_canvas.get_by_id.assert_called_once_with(dataflow_id) + + @pytest.mark.asyncio + async def test_load_dsl_for_pipeline_log_task_type_returns_corrected_id(self): + """When task_type != 'dataflow', dataflow_id comes from pipeline_log.pipeline_id.""" + ctx = MagicMock() + ctx.task_type = "raptor" + dataflow_id = "pipeline_log_id" + + with patch("rag.svr.task_executor_refactor.dataflow_service.PipelineOperationLogService") as mock_log: + mock_log_instance = MagicMock() + mock_log_instance.dsl = '{"id": "test_pipeline"}' + mock_log_instance.pipeline_id = "corrected_pipeline_id" + mock_log.get_by_id.return_value = (True, mock_log_instance) + + service = DataflowService(ctx=ctx) + + dsl, corrected_id = await service._load_dsl(dataflow_id) + + assert dsl == '{"id": "test_pipeline"}' + assert corrected_id == "corrected_pipeline_id" + mock_log.get_by_id.assert_called_once_with(dataflow_id) \ 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 index b6bfccb11c..76ff5c9d19 100644 --- 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 @@ -17,13 +17,13 @@ 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. +reaching into private orchestration methods. The new async implementation uses +thread_pool_exec for model.encode calls; tests mock that boundary. """ import numpy as np -from unittest.mock import MagicMock, patch +import pytest +from unittest.mock import AsyncMock, MagicMock, patch from rag.svr.task_executor_refactor.embedding_service import EmbeddingService @@ -56,17 +56,18 @@ class TestEmbeddingServiceInit: 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. + The async implementation uses thread_pool_exec for model.encode calls. + Tests mock thread_pool_exec at the module level to control returned vectors. """ - @patch.object(EmbeddingService, '_run_encode') - def test_embed_chunks_basic(self, mock_run_encode): + @pytest.mark.asyncio + @patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock) + async def test_embed_chunks_basic(self, mock_thread_pool): """Test basic chunk embedding.""" - mock_run_encode.return_value = (np.array([[1.0, 2.0]]), 10) + mock_thread_pool.return_value = (np.array([[1.0, 2.0]]), 10) ctx = MagicMock() ctx.progress_cb = None + ctx.embed_limiter = AsyncMockLimiter() service = EmbeddingService(ctx=ctx, embedding_batch_size=10) model = MagicMock() model.max_length = 100 @@ -74,23 +75,20 @@ class TestEmbeddingServiceEmbedChunks: docs = [ {"docnm_kwd": "Title1", "content_with_weight": "Content1"}, ] - tk_count, vector_size = service.embed_chunks(docs, model) + tk_count, vector_size = await 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) + @pytest.mark.asyncio + @patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock) + async def test_embed_chunks_uses_embedding_utils(self, mock_thread_pool): + """Test that embed_chunks uses thread_pool_exec for encoding.""" + mock_thread_pool.return_value = (np.array([[1.0, 2.0]]), 10) ctx = MagicMock() ctx.progress_cb = None + ctx.embed_limiter = AsyncMockLimiter() service = EmbeddingService(ctx=ctx, embedding_batch_size=10) model = MagicMock() model.max_length = 100 @@ -98,16 +96,19 @@ class TestEmbeddingServiceEmbedChunks: docs = [ {"docnm_kwd": "Title1", "content_with_weight": "Content1"}, ] - service.embed_chunks(docs, model) + await service.embed_chunks(docs, model) - mock_run_encode.assert_called() + # thread_pool_exec should be called at least once for encoding + mock_thread_pool.assert_called() - @patch.object(EmbeddingService, '_run_encode') - def test_embed_chunks_with_title_content_combination(self, mock_run_encode): + @pytest.mark.asyncio + @patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock) + async def test_embed_chunks_with_title_content_combination(self, mock_thread_pool): """Test that title and content vectors are combined.""" - mock_run_encode.return_value = (np.array([[1.0, 2.0]]), 10) + mock_thread_pool.return_value = (np.array([[1.0, 2.0]]), 10) ctx = MagicMock() ctx.progress_cb = None + ctx.embed_limiter = AsyncMockLimiter() service = EmbeddingService(ctx=ctx, embedding_batch_size=10) model = MagicMock() model.max_length = 100 @@ -115,21 +116,23 @@ class TestEmbeddingServiceEmbedChunks: docs = [ {"docnm_kwd": "Title1", "content_with_weight": "Content1"}, ] - _, vector_size = service.embed_chunks(docs, model, parser_config={"filename_embd_weight": 0.5}) + _, vector_size = await 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): + @pytest.mark.asyncio + @patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock) + async def test_embed_chunks_handles_long_text(self, mock_thread_pool): """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) + mock_thread_pool.return_value = (np.array([[1.0, 2.0]]), 10) ctx = MagicMock() ctx.progress_cb = None + ctx.embed_limiter = AsyncMockLimiter() service = EmbeddingService(ctx=ctx, embedding_batch_size=10) model = MagicMock() model.max_length = 100 @@ -137,9 +140,125 @@ class TestEmbeddingServiceEmbedChunks: docs = [ {"docnm_kwd": "Title1", "content_with_weight": "a" * 200}, ] - tk_count, vector_size = service.embed_chunks(docs, model) + tk_count, vector_size = await 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] + + + @pytest.mark.asyncio + @patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock) + async def test_embed_chunks_empty_docs(self, mock_thread_pool): + """Test embedding with empty docs list returns zero results.""" + ctx = MagicMock() + ctx.progress_cb = None + ctx.embed_limiter = AsyncMockLimiter() + service = EmbeddingService(ctx=ctx) + model = MagicMock() + model.max_length = 100 + + tk_count, vector_size = await service.embed_chunks([], model) + + assert tk_count == 0 + assert vector_size == 0 + + @pytest.mark.asyncio + @patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock) + async def test_embed_chunks_no_title(self, mock_thread_pool): + """Test embedding when chunks have no title — content vectors used directly.""" + mock_thread_pool.return_value = (np.array([[3.0, 4.0]]), 5) + ctx = MagicMock() + ctx.progress_cb = None + ctx.embed_limiter = AsyncMockLimiter() + service = EmbeddingService(ctx=ctx, embedding_batch_size=10) + model = MagicMock() + model.max_length = 100 + + docs = [{"content_with_weight": "Content only, no title"}] + tk_count, vector_size = await service.embed_chunks(docs, model) + + # With no title, only content is encoded (1 call); vector_size comes from content vec + assert vector_size == 2 + assert "q_2_vec" in docs[0] + + @pytest.mark.asyncio + @patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock) + async def test_embed_chunks_title_weight_zero(self, mock_thread_pool): + """Test embedding with filename_embd_weight=0.0 — no title contribution.""" + mock_thread_pool.return_value = (np.array([[1.0, 2.0]]), 5) + ctx = MagicMock() + ctx.progress_cb = None + ctx.embed_limiter = AsyncMockLimiter() + service = EmbeddingService(ctx=ctx, embedding_batch_size=10) + model = MagicMock() + model.max_length = 100 + + docs = [{"docnm_kwd": "Title1", "content_with_weight": "Content1"}] + _, vector_size = await service.embed_chunks(docs, model, + parser_config={"filename_embd_weight": 0.0}) + + assert vector_size == 2 + + @pytest.mark.asyncio + @patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock) + async def test_embed_chunks_title_weight_one(self, mock_thread_pool): + """Test embedding with filename_embd_weight=1.0 — full title contribution.""" + mock_thread_pool.return_value = (np.array([[1.0, 2.0]]), 5) + ctx = MagicMock() + ctx.progress_cb = None + ctx.embed_limiter = AsyncMockLimiter() + service = EmbeddingService(ctx=ctx, embedding_batch_size=10) + model = MagicMock() + model.max_length = 100 + + docs = [{"docnm_kwd": "Title1", "content_with_weight": "Content1"}] + _, vector_size = await service.embed_chunks(docs, model, + parser_config={"filename_embd_weight": 1.0}) + + assert vector_size == 2 + + @pytest.mark.asyncio + @patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock) + async def test_embed_chunks_encode_failure_propagates(self, mock_thread_pool): + """Test that model.encode exceptions are propagated to caller.""" + mock_thread_pool.side_effect = RuntimeError("embedding service unavailable") + ctx = MagicMock() + ctx.progress_cb = None + ctx.embed_limiter = AsyncMockLimiter() + service = EmbeddingService(ctx=ctx, embedding_batch_size=10) + model = MagicMock() + model.max_length = 100 + + docs = [{"docnm_kwd": "Title1", "content_with_weight": "Content1"}] + with pytest.raises(RuntimeError, match="embedding service unavailable"): + await service.embed_chunks(docs, model) + + @pytest.mark.asyncio + @patch("rag.svr.task_executor_refactor.embedding_service.thread_pool_exec", new_callable=AsyncMock) + async def test_embed_chunks_multiple_batches(self, mock_thread_pool): + """Test embedding with more chunks than batch size — multiple encode calls.""" + # Each call returns vectors matching input count + def side_effect(func, texts, *args, **kw): + n = len(texts) if isinstance(texts, list) else 1 + return np.random.rand(n, 2).astype(np.float32), 10 * n + + mock_thread_pool.side_effect = side_effect + ctx = MagicMock() + ctx.progress_cb = MagicMock() + ctx.embed_limiter = AsyncMockLimiter() + # batch_size=2, 5 chunks with titles → 1 title + ceil(5/2)=3 content = 4 calls + service = EmbeddingService(ctx=ctx, embedding_batch_size=2) + model = MagicMock() + model.max_length = 100 + + docs = [{"docnm_kwd": f"Title{i}", "content_with_weight": f"Content{i}"} for i in range(5)] + _, vector_size = await service.embed_chunks(docs, model) + + assert mock_thread_pool.call_count == 4 + assert vector_size > 0 + + +# Reuse from conftest +from test.unit_test.rag.svr.task_executor_refactor.conftest import AsyncMockLimiter 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 index ce808b2e52..7e42187693 100644 --- 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 @@ -60,6 +60,58 @@ class TestPostProcessorProcessTableParserMetadata: mock_merge.assert_called_once() + @pytest.mark.asyncio + async def test_processes_table_parser_manual_mode(self): + """Test that table parser in manual mode aggregates and persists metadata.""" + ctx = MagicMock() + ctx.parser_id = "table" + ctx.raw_task = {"parser_config": {}} + ctx.write_interceptor = None + chunks = [{"col_key": "val"}] + + with patch("rag.svr.task_executor_refactor.post_processor.merge_table_parser_config_from_kb") as mock_merge, \ + patch("rag.svr.task_executor_refactor.post_processor.aggregate_table_manual_doc_metadata") as mock_agg, \ + patch("rag.svr.task_executor_refactor.post_processor.table_parser_strip_doc_metadata_keys") as mock_strip, \ + patch("rag.svr.task_executor_refactor.post_processor.update_metadata_to") as mock_update, \ + patch("rag.svr.task_executor_refactor.post_processor.DocMetadataService") as mock_meta: + + mock_merge.return_value = {"table_column_mode": "manual"} + mock_agg.return_value = {"col_key": ["val1", "val2"]} + mock_strip.return_value = set() + mock_meta.get_document_metadata.return_value = {} + mock_update.return_value = {"col_key": ["val1", "val2"]} + + service = PostProcessor(ctx=ctx) + await service.process_table_parser_metadata("doc_1", chunks) + + mock_agg.assert_called_once_with(chunks, ctx.raw_task) + mock_meta.update_document_metadata.assert_called_once() + + @pytest.mark.asyncio + async def test_processes_table_parser_with_write_interceptor(self): + """Test table parser with write interceptor bypasses DB.""" + ctx = MagicMock() + ctx.parser_id = "table" + ctx.raw_task = {} + ctx.write_interceptor = MagicMock() + + with patch("rag.svr.task_executor_refactor.post_processor.merge_table_parser_config_from_kb") as mock_merge, \ + patch("rag.svr.task_executor_refactor.post_processor.aggregate_table_manual_doc_metadata") as mock_agg, \ + patch("rag.svr.task_executor_refactor.post_processor.table_parser_strip_doc_metadata_keys") as mock_strip, \ + patch("rag.svr.task_executor_refactor.post_processor.DocMetadataService") as mock_meta: + + mock_merge.return_value = {"table_column_mode": "manual"} + mock_agg.return_value = {"key": ["v"]} + mock_strip.return_value = set() + mock_meta.get_document_metadata.return_value = {} + + service = PostProcessor(ctx=ctx) + await service.process_table_parser_metadata("doc_1", []) + + ctx.write_interceptor.intercept.assert_called_once_with( + "DocMetadataService.update_document_metadata" + ) + class TestPostProcessorInsertTocChunk: """Tests for insert_toc_chunk method.""" 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 index 2399d13ed6..9f3aa60b5b 100644 --- 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 @@ -33,6 +33,7 @@ import pytest from unittest.mock import AsyncMock, MagicMock, patch from rag.svr.task_executor_refactor.raptor_service import RaptorService +from test.unit_test.rag.svr.task_executor_refactor.conftest import make_task_context # ============================================================================= @@ -216,7 +217,8 @@ class TestRaptorServiceRunRaptorForKb: # ---- Basic dispatch (file-level scope) ---- - def test_run_raptor_for_kb_file_scope_delegates_to_file_level( + @pytest.mark.asyncio + async 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.""" @@ -233,20 +235,7 @@ class TestRaptorServiceRunRaptorForKb: 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() + chunks, tk_count, cleanup = await svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, vector_size, doc_ids) mock_file.assert_called_once() mock_dataset.assert_not_called() @@ -255,7 +244,8 @@ class TestRaptorServiceRunRaptorForKb: # ---- Basic dispatch (dataset-level scope) ---- - def test_run_raptor_for_kb_dataset_scope_delegates_to_dataset_level( + @pytest.mark.asyncio + async 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.""" @@ -271,15 +261,7 @@ class TestRaptorServiceRunRaptorForKb: 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() + chunks, tk_count, cleanup = await svc.run_raptor_for_kb(raptor_config_dataset_scope, chat_mdl, embd_mdl, vector_size, doc_ids) mock_dataset.assert_called_once() mock_file.assert_not_called() @@ -288,7 +270,8 @@ class TestRaptorServiceRunRaptorForKb: # ---- Empty / no documents ---- - def test_run_raptor_for_kb_empty_doc_ids(self, mock_raptor_context, raptor_config_file_scope): + @pytest.mark.asyncio + async 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() @@ -299,15 +282,7 @@ class TestRaptorServiceRunRaptorForKb: 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() + chunks, tk_count, cleanup = await svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, 128, []) assert chunks == [] assert tk_count == 0 @@ -315,15 +290,11 @@ class TestRaptorServiceRunRaptorForKb: # ---- Cleanup scheduling through the public API ---- - def test_run_raptor_for_kb_returns_cleanup_list( + @pytest.mark.asyncio + async 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. - """ + """Cleanup list from internal method is propagated to caller.""" svc = RaptorService(mock_raptor_context) doc_ids = ["doc_1"] chat_mdl = MagicMock() @@ -336,28 +307,19 @@ class TestRaptorServiceRunRaptorForKb: }), 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() + chunks, tk_count, cleanup = await svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, 128, doc_ids) assert cleanup == expected_cleanup # ---- Dispatch with missing raptor config key ---- - def test_run_raptor_for_kb_defaults_to_file_scope_when_no_raptor_key( + @pytest.mark.asyncio + async 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.""" @@ -373,22 +335,15 @@ class TestRaptorServiceRunRaptorForKb: 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() + await svc.run_raptor_for_kb(config, chat_mdl, embd_mdl, 128, doc_ids) 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( + @pytest.mark.asyncio + async 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.""" @@ -403,15 +358,7 @@ class TestRaptorServiceRunRaptorForKb: }), 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() + await svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, vector_size, doc_ids) # Verify _run_file_level_raptor received vctr_nm with the correct vector size # Positional args: 0=raptor_config, 1=tree_builder, 2=clustering_method, @@ -421,7 +368,8 @@ class TestRaptorServiceRunRaptorForKb: # ---- Document info collection through public API ---- - def test_run_raptor_for_kb_collects_doc_info( + @pytest.mark.asyncio + async 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.""" @@ -436,17 +384,66 @@ class TestRaptorServiceRunRaptorForKb: 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() + await svc.run_raptor_for_kb(raptor_config_file_scope, chat_mdl, embd_mdl, 128, doc_ids) 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 + + +class TestRaptorServiceFileLevelRaptorCheckpoint: + """Tests for _run_file_level_raptor checkpoint behavior. + + Verifies the fix that moves progress_cb and continue to the outer if + block so progress is reported even when existing_methods == {tree_builder}. + """ + + @pytest.mark.asyncio + async def test_file_level_raptor_existing_methods_exact_match_updates_progress(self): + """When existing_methods == {tree_builder}, progress_cb is still called.""" + ctx = make_task_context() + svc = RaptorService(ctx) + + doc_ids = ["doc_1"] + doc_info_by_id = { + "doc_1": {"name": "a.pdf", "type": "pdf", "parser_id": "naive", "parser_config": {}} + } + raptor_config = { + "scope": "file", + "max_cluster": 64, "prompt": "test prompt", + "max_token": 256, "threshold": 0.1, "random_seed": 0, + "clustering_method": "gmm", "tree_builder": "raptor", + "ext": {}, + } + + with patch.object(svc, "_get_raptor_chunk_methods", new_callable=AsyncMock) as mock_methods, \ + patch.object(svc, "_should_skip_raptor", return_value=False): + + mock_methods.return_value = {"raptor"} + + result = await svc._run_file_level_raptor( + raptor_config=raptor_config, tree_builder="raptor", + clustering_method="gmm", chat_mdl=MagicMock(), + embd_mdl=MagicMock(), vctr_nm="q_128_vec", + doc_ids=doc_ids, doc_info_by_id=doc_info_by_id, + max_errors=3, res=[], tk_count=0, + cleanup_raptor_chunks=[], + ) + + msg_calls = [ + call.kwargs.get("msg", "") + for call in ctx.progress_cb.call_args_list + if call.kwargs.get("msg") is not None + ] + assert any("already has" in m for m in msg_calls), \ + f"Expected 'already has' progress message, got: {msg_calls}" + + prog_calls = [ + call.kwargs.get("prog") + for call in ctx.progress_cb.call_args_list + if call.kwargs.get("prog") is not None + ] + assert len(prog_calls) > 0, \ + "Expected progress_cb to be called with prog update" + assert result[0] == [] 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 index b7af25499a..1b7017ae14 100644 --- 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 @@ -216,10 +216,10 @@ class TestTaskContextLanguageAndModelProperties: """Tests for language and model properties.""" def test_language_default(self): - """Test language property defaults to 'en'.""" + """Test language property defaults to 'Chinese'.""" task = {"id": "task_1", "tenant_id": "tenant_1"} ctx = _make_ctx(task=task) - assert ctx.language == "en" + assert ctx.language == "Chinese" def test_language(self): """Test language property.""" 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 index fed3eb06af..23997c3bff 100644 --- 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 @@ -16,21 +16,27 @@ """ 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. +Mock strategy: external boundaries (LLMBundle, model config services, settings) +are mocked so that ``handle()`` and ``_bind_embedding_model`` execute their +real logic. Heavy orchestration methods (``_run_standard_chunking``, +``_run_raptor``, ``_run_graphrag``) are mocked since they are tested +exhaustively in the integration test suite. -Stable pure helpers (_build_toc, _get_vector_size) are tested directly since they -are side-effect-free data transformations. +Stable pure helpers (_build_toc) are tested directly. """ import pytest -import numpy as np from unittest.mock import MagicMock, AsyncMock, patch from rag.svr.task_executor_refactor.task_handler import TaskHandler +# Reuse shared helpers from conftest +from test.unit_test.rag.svr.task_executor_refactor.conftest import ( + patch_embedding_binding, + create_mock_settings, + make_task_context, +) + class TestTaskHandlerHandleTask: """Tests for the public handle_task() entry point.""" @@ -68,63 +74,86 @@ class TestTaskHandlerHandleTask: 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 + @pytest.mark.asyncio + async def test_handle_task_cleanup_skips_when_index_missing(self): + """Cancel cleanup should not call delete when the index doesn't exist.""" + from common import settings + mock_doc_store = MagicMock() + mock_doc_store.index_exist = MagicMock(return_value=False) + mock_doc_store.delete = MagicMock() + 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")) + with pytest.raises(Exception, match="test error"): + await handler.handle_task() + mock_doc_store.delete.assert_not_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. + External boundaries (LLMBundle, model config services, settings) are mocked + so that ``_bind_embedding_model`` and ``_init_kb`` execute their real logic + through ``handle()``. Only the heavy orchestration methods + (``_run_standard_chunking``, ``_run_raptor``, ``_run_graphrag``) are mocked. """ + # ── Context factory: make_task_context from conftest — see import above + @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) + """Test handle returns after dispatching memory task — no further processing.""" + ctx = make_task_context(task_type="memory") + ctx.raw_task = {"memory_id": "mem_1", "id": "task_1"} + + with patch("rag.svr.task_executor_refactor.task_handler.handle_save_to_memory_task", + new_callable=AsyncMock) as mock_handle: - 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() + handler._run_dataflow = AsyncMock() await handler.handle() + mock_handle.assert_called_once_with(ctx.raw_task) + # After memory task, should return immediately — no further routing + handler._run_standard_chunking.assert_not_called() + handler._run_dataflow.assert_not_called() @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) + """Test handle dispatches dataflow tasks (after embedding binding + init_kb).""" + ctx = make_task_context(task_type="dataflow", doc_id="doc_1") - handler = TaskHandler(ctx=ctx) - handler._run_dataflow = AsyncMock() - await handler.handle() - handler._run_dataflow.assert_called_once() + with patch_embedding_binding(), \ + patch("rag.svr.task_executor_refactor.task_handler.settings", create_mock_settings()), \ + patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_idx"): + + 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() + ctx = make_task_context(has_canceled_func=MagicMock(return_value=True)) handler = TaskHandler(ctx=ctx) await handler.handle() @@ -132,102 +161,87 @@ class TestTaskHandlerHandle: @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 + """Test handle routes to standard chunking. - 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() + ``_bind_embedding_model`` and ``_init_kb`` run their real code; + only the external boundary (LLM API, settings) is mocked. + """ + ctx = make_task_context() - await handler.handle() - handler._run_standard_chunking.assert_called_once() + with patch_embedding_binding(), \ + patch("rag.svr.task_executor_refactor.task_handler.settings", create_mock_settings()), \ + patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_idx"): + + handler = TaskHandler(ctx=ctx) + 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() + """Test handle routes to RAPTOR with real embedding binding.""" + ctx = make_task_context(task_type="raptor") - 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() + with patch_embedding_binding(), \ + patch("rag.svr.task_executor_refactor.task_handler.settings", create_mock_settings()), \ + patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_idx"): - await handler.handle() - handler._run_raptor.assert_called_once() + handler = TaskHandler(ctx=ctx) + 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() + """Test handle routes to GraphRAG with real embedding binding.""" + ctx = make_task_context(task_type="graphrag") - 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() + with patch_embedding_binding(), \ + patch("rag.svr.task_executor_refactor.task_handler.settings", create_mock_settings()), \ + patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_idx"): - await handler.handle() - handler._run_graphrag.assert_called_once() + handler = TaskHandler(ctx=ctx) + 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) + """Test handle returns early when embedding model binding fails. - handler = TaskHandler(ctx=ctx) - handler._bind_embedding_model = AsyncMock(return_value=None) + ``LLMBundle`` is patched to raise, so ``_bind_embedding_model`` + itself raises — no need to mock the private method. + """ + ctx = make_task_context() - await handler.handle() - # Should not call _run_standard_chunking when model is None - assert not hasattr(handler, '_run_standard_chunking_called') + with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_cfg, \ + patch("rag.svr.task_executor_refactor.task_handler.get_tenant_default_model_by_type") as mock_default, \ + patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle: + mock_cfg.return_value = MagicMock() + mock_default.return_value = MagicMock() + mock_bundle.side_effect = RuntimeError("embedding service unavailable") -class TestTaskHandlerGetVectorSize: - """Tests for _get_vector_size — stable pure helper.""" + handler = TaskHandler(ctx=ctx) + with pytest.raises(RuntimeError, match="embedding service unavailable"): + await handler.handle() - 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 + @pytest.mark.asyncio + async def test_handle_storage_binary_none_raises_file_not_found(self): + """Verify that None binary from storage raises FileNotFoundError.""" + ctx = make_task_context() + + with patch_embedding_binding(), \ + patch("rag.svr.task_executor_refactor.task_handler.settings", create_mock_settings()), \ + patch("rag.svr.task_executor_refactor.task_handler.search.index_name", return_value="test_idx"), \ + patch("rag.svr.task_executor_refactor.task_handler.File2DocumentService.get_storage_address", + return_value=("bucket_test", "name_test")), \ + patch.object(TaskHandler, "_get_storage_binary", new_callable=AsyncMock, return_value=None): + + handler = TaskHandler(ctx=ctx) + # Do NOT mock _run_standard_chunking — we want real code path for the check + with pytest.raises(FileNotFoundError, match="Can not find file from minio"): + await handler.handle() class TestTaskHandlerBuildToc: @@ -243,18 +257,19 @@ class TestTaskHandlerBuildToc: 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_from_provider_instance") as mock_cfg: + with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_cfg, \ + patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \ + patch("rag.svr.task_executor_refactor.task_handler.asyncio.run", side_effect=mock_asyncio_run): + 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 + mock_msg = MagicMock() + mock_bundle.return_value.__enter__.return_value = mock_msg + + 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.""" @@ -267,21 +282,22 @@ class TestTaskHandlerBuildToc: 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_from_provider_instance") as mock_cfg: + with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_cfg, \ + patch("rag.svr.task_executor_refactor.task_handler.LLMBundle") as mock_bundle, \ + patch("rag.svr.task_executor_refactor.task_handler.asyncio.run", side_effect=mock_asyncio_run): + 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 + mock_msg = MagicMock() + mock_bundle.return_value.__enter__.return_value = mock_msg + + 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: @@ -297,4 +313,4 @@ class TestTaskHandlerInit: 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 + assert handler._billing_hook is None 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 index dc0fa0ace4..9c710e6ec2 100644 --- 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 @@ -19,7 +19,6 @@ Integration tests for TaskHandler orchestration. import asyncio import gc -import uuid from typing import Any, Dict from unittest.mock import MagicMock, AsyncMock, patch @@ -37,6 +36,11 @@ from test.unit_test.rag.svr.task_executor_refactor.conftest import ( create_default_chunks, create_mock_settings, create_mock_chunk_service, + make_task_dict, + patch_get_storage_binary, + patch_task_handler_settings, + mock_thread_return_binary, + mock_thread_return_none, ) @@ -82,48 +86,13 @@ def create_task_context( 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() + task_dict = make_task_dict() ctx = create_task_context(task_dict) mock_embedding = create_mock_embedding_model(vector_size=128) mock_settings = create_mock_settings() @@ -152,11 +121,8 @@ class TestStandardChunkingPipelineIntegration: 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 + mock_thread_exec.side_effect = mock_thread_return_binary + mock_chunk_thread_exec.side_effect = mock_thread_return_binary handler = TaskHandler(ctx=ctx) await handler.handle() @@ -168,7 +134,7 @@ class TestStandardChunkingPipelineIntegration: @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() + task_dict = make_task_dict() ctx = create_task_context(task_dict) mock_embedding = create_mock_embedding_model(vector_size=128) mock_settings = create_mock_settings() @@ -197,11 +163,8 @@ class TestStandardChunkingPipelineIntegration: 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 + mock_thread_exec.side_effect = mock_thread_return_binary + mock_chunk_thread_exec.side_effect = mock_thread_return_binary handler = TaskHandler(ctx=ctx) await handler.handle() @@ -213,7 +176,7 @@ class TestStandardChunkingPipelineIntegration: @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() + task_dict = make_task_dict() ctx = create_task_context(task_dict) mock_embedding = create_mock_embedding_model(vector_size=128) mock_settings = create_mock_settings() @@ -245,11 +208,8 @@ class TestStandardChunkingPipelineIntegration: 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 + mock_thread_exec.side_effect = mock_thread_return_binary + mock_chunk_thread_exec.side_effect = mock_thread_return_binary handler = TaskHandler(ctx=ctx) await handler.handle() @@ -262,7 +222,7 @@ class TestStandardChunkingPipelineIntegration: @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() + task_dict = make_task_dict() ctx = create_task_context(task_dict) mock_embedding = create_mock_embedding_model(vector_size=128) mock_settings = create_mock_settings() @@ -291,11 +251,8 @@ class TestStandardChunkingPipelineIntegration: 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 + mock_thread_exec.side_effect = mock_thread_return_binary + mock_chunk_thread_exec.side_effect = mock_thread_return_binary handler = TaskHandler(ctx=ctx) await handler.handle() @@ -311,7 +268,7 @@ class TestStandardChunkingPipelineIntegration: @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() + task_dict = make_task_dict() ctx = create_task_context(task_dict) mock_embedding = create_mock_embedding_model(vector_size=128) mock_settings = create_mock_settings() @@ -340,11 +297,8 @@ class TestStandardChunkingPipelineIntegration: 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 + mock_thread_exec.side_effect = mock_thread_return_binary + mock_chunk_thread_exec.side_effect = mock_thread_return_binary handler = TaskHandler(ctx=ctx) await handler.handle() @@ -357,31 +311,10 @@ class TestStandardChunkingPipelineIntegration: 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() + task_dict = make_task_dict() ctx = create_task_context(task_dict, is_canceled=True) mock_settings = create_mock_settings() @@ -411,7 +344,7 @@ class TestTaskCancellationCleanupIntegration: @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() + task_dict = make_task_dict() ctx = create_task_context(task_dict, is_canceled=True) mock_settings = create_mock_settings() @@ -450,7 +383,7 @@ class TestTaskCancellationCleanupIntegration: @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() + task_dict = make_task_dict() ctx = create_task_context(task_dict, is_canceled=True) mock_settings = create_mock_settings() @@ -481,30 +414,10 @@ class TestTaskCancellationCleanupIntegration: 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() + task_dict = make_task_dict(doc_id=GRAPH_RAPTOR_FAKE_DOC_ID, doc_ids=["doc1", "doc2"], task_type="raptor", parser_config={"raptor": {"use_raptor": False}}, kb_parser_config={"raptor": {"use_raptor": False}}) ctx = create_task_context(task_dict) mock_embedding = create_mock_embedding_model(vector_size=128) mock_settings = create_mock_settings() @@ -534,10 +447,7 @@ class TestRaptorPipelineIntegration: 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 + mock_thread_exec.side_effect = mock_thread_return_none handler = TaskHandler(ctx=ctx) await handler.handle() @@ -549,7 +459,7 @@ class TestRaptorPipelineIntegration: @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() + task_dict = make_task_dict(doc_id=GRAPH_RAPTOR_FAKE_DOC_ID, doc_ids=["doc1", "doc2"], task_type="raptor", parser_config={"raptor": {"use_raptor": False}}, kb_parser_config={"raptor": {"use_raptor": False}}) ctx = create_task_context(task_dict) mock_embedding = create_mock_embedding_model(vector_size=128) mock_settings = create_mock_settings() @@ -579,10 +489,7 @@ class TestRaptorPipelineIntegration: 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 + mock_thread_exec.side_effect = mock_thread_return_none handler = TaskHandler(ctx=ctx) await handler.handle() @@ -598,31 +505,10 @@ class TestRaptorPipelineIntegration: 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() + task_dict = make_task_dict() ctx = create_task_context(task_dict) with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \ @@ -639,7 +525,7 @@ class TestEmbeddingModelBindingFailureIntegration: @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() + task_dict = make_task_dict() ctx = create_task_context(task_dict) with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") as mock_get_config, \ @@ -659,49 +545,26 @@ class TestEmbeddingModelBindingFailureIntegration: 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() + task_dict = make_task_dict(doc_id=CANVAS_DEBUG_DOC_ID, task_type="dataflow") ctx = create_task_context(task_dict) + mock_embedding = create_mock_embedding_model(vector_size=128) - 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 + with patch("rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance") 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", return_value="test_idx"), \ + patch("rag.svr.task_executor_refactor.task_handler.settings") as mock_settings, \ + patch("rag.svr.task_executor_refactor.task_handler.DataflowService") as mock_dataflow_service: - handler = TaskHandler(ctx=ctx) - await handler.handle() + mock_get_config.return_value = MagicMock() + mock_get_default.return_value = MagicMock() + mock_bundle.return_value = mock_embedding + mock_settings.docStoreConn = MagicMock() + mock_settings.docStoreConn.create_idx = MagicMock() - 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 @@ -716,37 +579,11 @@ class TestDataflowPipelineIntegration: 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() + task_dict = make_task_dict(parser_config={"auto_keywords": 0, "auto_questions": 0, "enable_metadata": False, "toc_extraction": True}) ctx = create_task_context(task_dict) mock_embedding = create_mock_embedding_model(vector_size=128) mock_settings = create_mock_settings() @@ -779,11 +616,8 @@ class TestTocAsyncFlowIntegration: 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 + mock_thread_exec.side_effect = mock_thread_return_binary + mock_chunk_thread_exec.side_effect = mock_thread_return_binary handler = TaskHandler(ctx=ctx) await handler.handle() @@ -811,7 +645,7 @@ class TestTocAsyncFlowIntegration: by pytest-asyncio. """ - task_dict = self._create_toc_enabled_task_dict() + task_dict = make_task_dict(parser_config={"auto_keywords": 0, "auto_questions": 0, "enable_metadata": False, "toc_extraction": True}) task_dict["parser_config"]["toc_extraction"] = False ctx = create_task_context(task_dict) mock_embedding = create_mock_embedding_model(vector_size=128) @@ -842,11 +676,8 @@ class TestTocAsyncFlowIntegration: 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 + mock_thread_exec.side_effect = mock_thread_return_binary + mock_chunk_thread_exec.side_effect = mock_thread_return_binary handler = TaskHandler(ctx=ctx) await handler.handle() @@ -873,35 +704,10 @@ class TestTocAsyncFlowIntegration: 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() + task_dict = make_task_dict() ctx = create_task_context(task_dict) mock_embedding = create_mock_embedding_model(vector_size=128) mock_settings = create_mock_settings() @@ -930,11 +736,8 @@ class TestRecordingContextDataFlowAssertions: 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 + mock_thread_exec.side_effect = mock_thread_return_binary + mock_chunk_thread_exec.side_effect = mock_thread_return_binary handler = TaskHandler(ctx=ctx) await handler.handle() @@ -947,7 +750,7 @@ class TestRecordingContextDataFlowAssertions: @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() + task_dict = make_task_dict() ctx = create_task_context(task_dict) mock_embedding = create_mock_embedding_model(vector_size=128) mock_settings = create_mock_settings() @@ -976,11 +779,8 @@ class TestRecordingContextDataFlowAssertions: 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 + mock_thread_exec.side_effect = mock_thread_return_binary + mock_chunk_thread_exec.side_effect = mock_thread_return_binary handler = TaskHandler(ctx=ctx) await handler.handle() @@ -991,3 +791,66 @@ class TestRecordingContextDataFlowAssertions: 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}" + + +class TestGraphragPipelineIntegration: + """P2: Integration tests for GraphRAG pipeline default configuration.""" + + @pytest.mark.asyncio + async def test_graphrag_pipeline_configures_full_defaults(self): + """Verify that GraphRAG configures all default parameters when not already set.""" + task_dict = make_task_dict(doc_ids=["doc1", "doc2"], task_type="graphrag") + rec_ctx = RecordingContext() + ctx = create_task_context(task_dict, recording_context=rec_ctx) + 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 = {} + + 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_from_provider_instance") 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.run_graphrag_for_kb") as mock_run_graphrag, \ + patch("rag.svr.task_executor_refactor.task_handler.DocumentService"): + + 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_run_graphrag.return_value = {"status": "completed"} + + mock_thread_exec.side_effect = mock_thread_return_none + + handler = TaskHandler(ctx=ctx) + await handler.handle() + + # Verify update_by_id was called with full default config + mock_kb_service.update_by_id.assert_called_once() + call_args = mock_kb_service.update_by_id.call_args + config = call_args[0][1]["parser_config"]["graphrag"] + assert config["use_graphrag"] is True + assert "organization" in config["entity_types"] + assert "person" in config["entity_types"] + assert "geo" in config["entity_types"] + assert "event" in config["entity_types"] + assert "category" in config["entity_types"] + assert config["method"] == "light" + assert "batch_chunk_token_size" in config + assert "retry_attempts" in config + assert "retry_backoff_seconds" in config + assert "retry_backoff_max_seconds" in config + assert "build_subgraph_timeout_per_chunk_seconds" in config + assert "build_subgraph_min_timeout_seconds" in config + assert "merge_timeout_seconds" in config + assert "resolution_timeout_seconds" in config + assert "community_timeout_seconds" in config + assert "lock_acquire_timeout_seconds" in config, \ + "All GraphRAG default config parameters should be present"