zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-18 22:38:25 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: overhaul task executor with layered architecture and comprehensive test suite (#15471)

Browse Source 
## 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 <noreply@anthropic.com>
This commit is contained in:
Jack
2026-06-03 17:18:31 +08:00
committed by GitHub
parent d736f358ba
commit b363146997
21 changed files with 1317 additions and 1222 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
CLAUDE.md
View File

@@ -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/<page_name>`
- **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

1
rag/svr/task_executor.py
View File

@@ -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

10
rag/svr/task_executor_refactor/chunk_post_processor.py
View File

@@ -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):

3
rag/svr/task_executor_refactor/chunk_service.py
View File

@@ -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)

19
rag/svr/task_executor_refactor/dataflow_service.py
View File

@@ -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]:

35
rag/svr/task_executor_refactor/embedding_service.py
View File

@@ -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)

9
rag/svr/task_executor_refactor/raptor_service.py
View File

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

2
rag/svr/task_executor_refactor/task_context.py
View File

@@ -224,7 +224,7 @@ class TaskContext:
"parser_id": "",
"parser_config": {},
"kb_parser_config": {},
"language": "en",
"language": "Chinese",
"llm_id": "",
"embd_id": "",
"from_page": 0,

68
rag/svr/task_executor_refactor/task_handler.py
View File

@@ -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:

296
test/unit_test/rag/svr/task_executor_refactor/conftest.py
View File

@@ -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)

103
test/unit_test/rag/svr/task_executor_refactor/test_chunk_builder.py
View File

@@ -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)

524
test/unit_test/rag/svr/task_executor_refactor/test_chunk_post_processor.py
View File

@@ -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"}]

199
test/unit_test/rag/svr/task_executor_refactor/test_chunk_service.py
View File

@@ -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"}]

15
test/unit_test/rag/svr/task_executor_refactor/test_comparator.py
View File

@@ -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):

140
test/unit_test/rag/svr/task_executor_refactor/test_dataflow_service.py
View File

@@ -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
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)

179
test/unit_test/rag/svr/task_executor_refactor/test_embedding_service.py
View File

@@ -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

52
test/unit_test/rag/svr/task_executor_refactor/test_post_processor.py
View File

@@ -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."""

163
test/unit_test/rag/svr/task_executor_refactor/test_raptor_service.py
View File

@@ -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] == []

4
test/unit_test/rag/svr/task_executor_refactor/test_task_context.py
View File

@@ -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."""

282
test/unit_test/rag/svr/task_executor_refactor/test_task_handler.py
View File

@@ -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 <test.pdf> 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
assert handler._billing_hook is None

371
test/unit_test/rag/svr/task_executor_refactor/test_task_handler_integration.py
View File

@@ -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"