From ea08478a371932b81f821ea23b1447becaa7060a Mon Sep 17 00:00:00 2001 From: caoxiaozhu Date: Sat, 25 Jul 2026 18:00:21 +0800 Subject: [PATCH] =?UTF-8?q?feat(data-process):=20=E6=8E=A5=E5=85=A5?= =?UTF-8?q?=E4=B8=89=E7=A7=8D=E6=96=87=E6=A1=A3=E5=88=87=E5=88=86=E5=BC=95?= =?UTF-8?q?=E6=93=8E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/app/api/v1/endpoints/data_process.py | 263 ++++------- .../app/modules/data_process/algorithms.py | 326 +------------ .../modules/data_process/document_chunking.py | 444 ++++++++++++++++++ backend/app/schemas/data_process.py | 25 +- backend/pyproject.toml | 3 + backend/requirements.txt | 3 + backend/tests/test_data_process_algorithms.py | 157 ------- backend/tests/test_data_process_api.py | 77 +-- backend/tests/test_document_chunking.py | 85 ++++ 9 files changed, 667 insertions(+), 716 deletions(-) create mode 100644 backend/app/modules/data_process/document_chunking.py create mode 100644 backend/tests/test_document_chunking.py diff --git a/backend/app/api/v1/endpoints/data_process.py b/backend/app/api/v1/endpoints/data_process.py index 7297fef..1d1d0cd 100644 --- a/backend/app/api/v1/endpoints/data_process.py +++ b/backend/app/api/v1/endpoints/data_process.py @@ -30,13 +30,10 @@ from psycopg.rows import dict_row from app.modules.data_process.algorithms import ( ParsedText, - TextChunk, canonical_record_json, - chunk_unstructured, content_quality_flags, desensitize_pii, desensitize_structured_record, - detect_document_structure, detect_pdf_document_noise, estimate_token_count, extract_pdf_page_texts, @@ -48,6 +45,13 @@ from app.modules.data_process.algorithms import ( remove_document_noise, score_quality, ) +from app.modules.data_process.document_chunking import ( + DocumentChunk, + chunk_fixed_text, + chunk_layout_document, + chunk_semantic_text, + merge_short_chunks, +) from app.modules.data_process.generation import generate_model_records from app.modules.data_process.storage import ( LocalDataProcessStorage, @@ -238,152 +242,58 @@ def _parse_stored_source(source: dict[str, Any]) -> ParsedText: ) -def _shift_chunk(chunk: TextChunk, offset: int, source_text: str) -> TextChunk: - start = offset + chunk.start - end = offset + chunk.end - return TextChunk( - content=chunk.content, - start=start, - end=end, - start_line=source_text.count("\n", 0, start) + 1, - end_line=source_text.count("\n", 0, max(start, end - 1)) + 1, - token_count=chunk.token_count, - ) - - -def _merge_short_chunks( - chunks: list[TextChunk], - source_text: str, - *, - min_token_count: int, - chunk_size: int, -) -> list[TextChunk]: - """按原文顺序合并相邻短块，同时严格遵守切片大小上限。""" - - merged: list[TextChunk] = [] - index = 0 - while index < len(chunks): - current = chunks[index] - if current.token_count >= min_token_count: - merged.append(current) - index += 1 - continue - - candidates: list[tuple[int, int, int]] = [] - if merged: - candidates.append((merged[-1].start, current.end, -1)) - if index + 1 < len(chunks): - candidates.append((current.start, chunks[index + 1].end, 1)) - selected = next( - ( - (start, end, direction) - for start, end, direction in candidates - if estimate_token_count(source_text[start:end]) <= chunk_size - ), - None, - ) - if selected is None: - merged.append(current) - index += 1 - continue - - start, end, direction = selected - combined = TextChunk( - content=source_text[start:end], - start=start, - end=end, - start_line=source_text.count("\n", 0, start) + 1, - end_line=source_text.count("\n", 0, max(start, end - 1)) + 1, - token_count=estimate_token_count(source_text[start:end]), - ) - if direction < 0: - merged[-1] = combined - index += 1 - else: - merged.append(combined) - index += 2 - return merged - - def _chunk_source_text( - text: str, + source: dict[str, Any], config: dict[str, Any], preprocess_options: set[str], -) -> list[tuple[TextChunk, tuple[str, ...]]]: - """按可选文档结构分段后切片，结构边界之间不共享 overlap。""" +) -> list[DocumentChunk]: + """根据任务配置调用真实的 Docling/LlamaIndex 切分器。""" - method = str(_value(config, "chunk_method", "chunkMethod", "structure")) - detect_structure = ( - method == "structure" or "detect_document_structure" in preprocess_options - ) + method = str(_value(config, "chunk_method", "chunkMethod", "layout_hybrid")) preserve_context = "preserve_context" in preprocess_options - merge_short = "merge_short_content" in preprocess_options chunk_size = int(_value(config, "chunk_size", "chunkSize", 800)) - configured_minimum = int(_value(config, "min_chunk_size", "minChunkSize", 100)) - minimum = configured_minimum if merge_short else 1 overlap = ( int(_value(config, "chunk_overlap", "chunkOverlap", 100)) if preserve_context else 0 ) - common = { - "method": method, - "chunk_size": chunk_size, - "chunk_overlap": overlap, - "min_chunk_size": minimum, - "custom_delimiter": str( - _value(config, "custom_delimiter", "customDelimiter", "") or "" - ), - "preserve_code_blocks": bool( - _value(config, "preserve_code_blocks", "preserveCodeBlocks", False) - ), - "preserve_tables": bool( - _value(config, "preserve_tables", "preserveTables", False) - ), - "preserve_lists": bool( - _value(config, "preserve_lists", "preserveLists", False) - ), - } - - sections: list[tuple[int, int, tuple[str, ...]]] = [(0, len(text), ())] - if detect_structure: - structure = detect_document_structure(text) - if structure.headings: - sections = [] - first_start = structure.headings[0].start - if first_start > 0 and text[:first_start].strip(): - sections.append((0, first_start, ())) - stack: list[tuple[int, str]] = [] - for index, heading in enumerate(structure.headings): - while stack and stack[-1][0] >= heading.level: - stack.pop() - stack.append((heading.level, heading.title)) - end = ( - structure.headings[index + 1].start - if index + 1 < len(structure.headings) - else len(text) - ) - sections.append((heading.start, end, tuple(title for _, title in stack))) - - result: list[tuple[TextChunk, tuple[str, ...]]] = [] - for start, end, heading_path in sections: - section_text = text[start:end] - local_chunks = chunk_unstructured(section_text, **common) - shifted = [_shift_chunk(chunk, start, text) for chunk in local_chunks] - result.extend((chunk, heading_path) for chunk in shifted) - - if merge_short and result: - # 结构分段只负责提供标题路径和隔离 overlap，不应让目录项或短小节 - # 突破 min_chunk_size 约束。合并后保留首个原始块的标题路径。 - heading_paths = {chunk.start: heading_path for chunk, heading_path in result} - merged = _merge_short_chunks( - [chunk for chunk, _ in result], - text, - min_token_count=configured_minimum, + text = str(source.get("content") or "") + if method == "layout_hybrid": + raw = source.get("raw_content") + if not isinstance(raw, bytes): + raise InvalidStateError("版面结构混合切分需要原始文件，请重新上传后再处理") + chunks = chunk_layout_document( + raw, + filename=str(source.get("name") or "document.pdf"), + source_text=text, chunk_size=chunk_size, ) - result = [(chunk, heading_paths.get(chunk.start, ())) for chunk in merged] - return result + elif method == "semantic": + chunks = chunk_semantic_text( + text, + chunk_size=chunk_size, + chunk_overlap=overlap, + breakpoint_percentile_threshold=int( + _value( + config, + "semantic_breakpoint_percentile", + "semanticBreakpointPercentile", + 95, + ) + ), + ) + elif method == "fixed": + chunks = chunk_fixed_text(text, chunk_size=chunk_size, chunk_overlap=overlap) + else: + raise ValueError(f"unsupported chunk method: {method}") + if "merge_short_content" in preprocess_options: + chunks = merge_short_chunks( + chunks, + source_text=text, + min_token_count=int(_value(config, "min_chunk_size", "minChunkSize", 100)), + max_token_count=chunk_size, + ) + return chunks _NEGATION_MARKERS = frozenset({"不", "无", "未", "否", "没有", "并非", "not", "no", "never"}) @@ -456,23 +366,27 @@ def _build_preview_items( parsed = _parse_stored_source(source) if process_type == "unstructured": document_noise_spans = tuple(source.get("document_noise_spans") or ()) - chunks = _chunk_source_text(parsed.text, config, preprocess_options) - for chunk, heading_path in chunks: + chunks = _chunk_source_text(source, config, preprocess_options) + for chunk in chunks: content = ( remove_document_noise( - chunk.content, + chunk.contextualized_content, document_noise_spans, - source_offset=chunk.start, + source_offset=chunk.source_start or 0, ) - if should_clean_invalid and document_noise_spans - else chunk.content + if ( + should_clean_invalid + and document_noise_spans + and chunk.source_start is not None + ) + else chunk.contextualized_content ) preprocess_flags = content_quality_flags( content, min_chars=0, min_tokens=0, ) - if content != chunk.content: + if content != chunk.contextualized_content: preprocess_flags = (*preprocess_flags, "document_noise_removed") flag_set = set(preprocess_flags) if "clean_invalid_content" in preprocess_options and flag_set & { @@ -508,25 +422,28 @@ def _build_preview_items( quality = _preview_quality(content, config) quality["pii_replacements"] = pii_counts quality["preprocess_flags"] = list(preprocess_flags) - chunk_method = str( - _value(config, "chunk_method", "chunkMethod", "structure") + quality["chunk_method"] = str( + _value(config, "chunk_method", "chunkMethod", "layout_hybrid") ) - if ( - chunk_method == "structure" - or "detect_document_structure" in preprocess_options - ): - quality["heading_path"] = list(heading_path) + quality["heading_path"] = list(chunk.heading_path) + quality["source_pages"] = list(chunk.source_pages) + quality["doc_item_refs"] = list(chunk.doc_item_refs) + quality["source_bboxes"] = list(chunk.source_bboxes) append_item( { "source_file_id": source["id"], - "original_content": chunk.content, + "original_content": chunk.original_content, "edited_content": content, - "source_start": chunk.start, - "source_end": chunk.end, - "source_start_line": chunk.start_line, - "source_end_line": chunk.end_line, + "source_start": chunk.source_start, + "source_end": chunk.source_end, + "source_start_line": chunk.source_start_line, + "source_end_line": chunk.source_end_line, "token_count": estimate_token_count(content), - "status": "modified" if content != chunk.content else "original", + "status": ( + "modified" + if content != chunk.original_content + else "original" + ), "quality_score": quality, } ) @@ -1265,13 +1182,21 @@ def _prepare_preview_items( ] if not sources: raise InvalidStateError("at least one source file is required") - preprocess_options = _preprocess_options(task.get("config") or {}) - if ( - task.get("process_type") == "unstructured" - and preprocess_options & {"clean_invalid", "clean_invalid_content"} + config = task.get("config") or {} + preprocess_options = _preprocess_options(config) + chunk_method = str( + _value(config, "chunk_method", "chunkMethod", "layout_hybrid") + ) + is_unstructured = task.get("process_type") == "unstructured" + if is_unstructured and ( + chunk_method == "layout_hybrid" + or preprocess_options & {"clean_invalid", "clean_invalid_content"} ): for index, source in enumerate(sources): - if str(source.get("file_format") or "").lower() != "pdf": + if ( + chunk_method != "layout_hybrid" + and str(source.get("file_format") or "").lower() != "pdf" + ): continue storage_object_id = str(source.get("storage_object_id") or "") actual_size = storage.file_size( @@ -1280,10 +1205,10 @@ def _prepare_preview_items( expected_source_file_id=str(source["id"]), ) if actual_size is None: - logger.info( - "skip PDF document noise detection for unavailable legacy source %s", - source["id"], - ) + if chunk_method == "layout_hybrid": + raise InvalidStateError( + "版面结构混合切分无法读取原始文件，请重新上传后再处理" + ) continue expected_size = int(source.get("size_bytes") or 0) if expected_size and actual_size != expected_size: @@ -1296,6 +1221,13 @@ def _prepare_preview_items( expected_size=actual_size, ) ) + enriched = dict(source) + if chunk_method == "layout_hybrid": + enriched["raw_content"] = raw + sources[index] = enriched + continue + if str(source.get("file_format") or "").lower() != "pdf": + continue pages = extract_pdf_page_texts(raw) extracted_text = "\n\n".join(page.text for page in pages if page.text) if extracted_text != str(source.get("content") or ""): @@ -1304,7 +1236,6 @@ def _prepare_preview_items( source["id"], ) continue - enriched = dict(source) enriched["document_noise_spans"] = detect_pdf_document_noise(pages) sources[index] = enriched items = _build_preview_items(task, sources) diff --git a/backend/app/modules/data_process/algorithms.py b/backend/app/modules/data_process/algorithms.py index 322a0ec..51d9a3b 100644 --- a/backend/app/modules/data_process/algorithms.py +++ b/backend/app/modules/data_process/algorithms.py @@ -1,8 +1,7 @@ """数据处理模块使用的无副作用算法。 本模块不访问数据库、文件系统或网络，便于 API、后台任务和测试共同复用。 -所有偏移量均为 Python 字符串偏移量，``TextChunk.content`` 始终等于 -``source[chunk.start:chunk.end]``。 +所有偏移量均为 Python 字符串偏移量。 """ from __future__ import annotations @@ -16,7 +15,6 @@ import re import unicodedata import xml.etree.ElementTree as ET import zipfile -from bisect import bisect_left from collections import Counter from collections.abc import Iterable, Mapping, Sequence from copy import deepcopy @@ -31,7 +29,6 @@ from docx.oxml.table import CT_Tbl from docx.oxml.text.paragraph import CT_P from docx.table import Table from docx.text.paragraph import Paragraph -from llama_index.core.node_parser import SentenceSplitter, TokenTextSplitter from openpyxl import load_workbook from openpyxl.utils.cell import range_boundaries from pptx import Presentation @@ -48,7 +45,6 @@ TextFormat = Literal[ "xlsx", "pptx", ] -ChunkMethod = Literal["structure", "fixed", "custom"] DatasetSplit = Literal["train", "validation", "test"] StructuredPreprocessOption = Literal[ "clean_invalid", @@ -153,7 +149,6 @@ _ENGLISH_NAME_CONTEXT_PATTERN = re.compile( r"(?P[A-Za-z][A-Za-z'’-]*(?:[ \t]+[A-Za-z][A-Za-z'’-]*){0,3})" ) _TOKEN_PATTERN = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff]|[A-Za-z0-9_]+|[^\s]") -_SEMANTIC_BOUNDARY_PATTERN = re.compile(r"\n\s*\n|[。！？!?；;](?:[\"'”’）】》]*)|\.(?:\s+|$)") @dataclass(frozen=True, slots=True) @@ -184,18 +179,6 @@ class DocumentNoiseSpan: kind: Literal["page_number", "repeated_margin", "table_of_contents"] -@dataclass(frozen=True, slots=True) -class TextChunk: - """带有可追溯来源位置的非结构化文本切片。""" - - content: str - start: int - end: int - start_line: int - end_line: int - token_count: int - - @dataclass(frozen=True, slots=True) class QualityScore: """标准 instruction/input/output 记录的可解释质量分。""" @@ -1975,41 +1958,6 @@ def estimate_token_count(text: str) -> int: return len(_TOKEN_PATTERN.findall(text)) -def _token_spans(text: str) -> list[tuple[int, int]]: - return [match.span() for match in _TOKEN_PATTERN.finditer(text)] - - -def _deterministic_tokenizer(text: str) -> list[str]: - """LlamaIndex splitter 使用的稳定 tokenizer，与预览 token 计数完全一致。""" - - return _TOKEN_PATTERN.findall(text) - - -def _sentence_chunks(text: str) -> list[str]: - """按项目既有中英文句界切句，避免 SentenceSplitter 触发 NLTK 下载。""" - - chunks: list[str] = [] - cursor = 0 - for match in _SEMANTIC_BOUNDARY_PATTERN.finditer(text): - end = match.end() - if end > cursor: - chunks.append(text[cursor:end]) - cursor = end - if cursor < len(text): - chunks.append(text[cursor:]) - return chunks or [text] - - -def _line_number(newline_offsets: list[int], offset: int) -> int: - # 换行符本身仍属于上一行；只有严格位于 offset 之前的换行才推进行号。 - return bisect_left(newline_offsets, offset) + 1 - - -def _token_index_at_or_after(spans: list[tuple[int, int]], offset: int) -> int: - starts = [span[0] for span in spans] - return bisect_left(starts, offset) - - def _protected_markdown_ranges( text: str, *, @@ -2110,12 +2058,6 @@ def _protected_markdown_ranges( return merged -def _range_containing( - ranges: Sequence[tuple[int, int]], offset: int -) -> tuple[int, int] | None: - return next((item for item in ranges if item[0] < offset < item[1]), None) - - def protected_context_ranges( text: str, *, @@ -2444,269 +2386,6 @@ def is_near_duplicate( ) -def _boundary_for_method( - text: str, - spans: list[tuple[int, int]], - start_index: int, - ideal_end_index: int, - minimum_end_index: int, - method: ChunkMethod, - custom_delimiter: str, - splitter: SentenceSplitter | TokenTextSplitter | None, -) -> tuple[int, int | None]: - if method == "fixed" and not isinstance(splitter, TokenTextSplitter): - raise RuntimeError("fixed chunking requires TokenTextSplitter") - - start_offset = spans[start_index][0] - ideal_end_offset = spans[ideal_end_index - 1][1] - minimum_end_offset = spans[minimum_end_index - 1][1] - search_text = text[start_offset:ideal_end_offset] - - if method == "fixed": - # 直接让 TokenTextSplitter 处理真实文本；开启空白保留后将首块边界 - # 投影回稳定 token span，offset 和实际 overlap 仍由外层统一维护。 - lookahead_end = min(len(spans), ideal_end_index + 1) - window_end = spans[lookahead_end - 1][1] - window = text[start_offset:window_end] - chunks = splitter.split_text(window) - first_chunk = chunks[0] if chunks else "" - if first_chunk and window.startswith(first_chunk): - boundary_offset = start_offset + len(first_chunk) - boundary_index = _token_index_at_or_after(spans, boundary_offset) - if boundary_index > start_index: - return min(boundary_index, ideal_end_index), boundary_offset - return ideal_end_index, None - - if method == "custom": - delimiter = custom_delimiter.replace("\\n", "\n").replace("\\t", "\t") - if not delimiter: - raise ValueError("custom_delimiter is required for custom chunking") - relative_minimum = max(0, minimum_end_offset - start_offset) - delimiter_start = search_text.rfind(delimiter, relative_minimum) - if delimiter_start >= 0: - boundary_offset = start_offset + delimiter_start + len(delimiter) - boundary_index = _token_index_at_or_after(spans, boundary_offset) - if boundary_index > start_index: - return min(boundary_index, ideal_end_index), boundary_offset - return ideal_end_index, None - - if method == "structure": - if not isinstance(splitter, SentenceSplitter): - raise RuntimeError("structure chunking requires SentenceSplitter") - # 多给一个 token 使 splitter 确实执行限长；只取首块并投影回原文。 - lookahead_end = min(len(spans), ideal_end_index + 1) - window_end = spans[lookahead_end - 1][1] - window = text[start_offset:window_end] - chunks = splitter.split_text(window) - first_chunk = chunks[0] if chunks else "" - if first_chunk and window.startswith(first_chunk): - boundary_offset = start_offset + len(first_chunk) - else: - boundary_offset = ideal_end_offset - boundary_index = _token_index_at_or_after(spans, boundary_offset) - if boundary_index >= minimum_end_index: - return min(boundary_index, ideal_end_index), boundary_offset - return ideal_end_index, None - - -def _chunk_normalized_text( - normalized: str, - *, - method: ChunkMethod, - chunk_size: int = 800, - chunk_overlap: int = 100, - min_chunk_size: int = 100, - custom_delimiter: str = "", - preserve_code_blocks: bool = False, - preserve_tables: bool = False, - preserve_lists: bool = False, -) -> list[TextChunk]: - if not normalized: - return [] - spans = _token_spans(normalized) - if not spans: - return [] - - newline_offsets = [index for index, char in enumerate(normalized) if char == "\n"] - protected_ranges = _protected_markdown_ranges( - normalized, - preserve_code_blocks=preserve_code_blocks, - preserve_tables=preserve_tables, - preserve_lists=preserve_lists, - ) - splitter: SentenceSplitter | TokenTextSplitter | None - if method == "fixed": - splitter = TokenTextSplitter( - chunk_size=chunk_size, - chunk_overlap=chunk_overlap, - tokenizer=_deterministic_tokenizer, - separator=" ", - backup_separators=["\n"], - keep_whitespaces=True, - include_metadata=False, - include_prev_next_rel=False, - ) - elif method == "structure": - splitter = SentenceSplitter( - chunk_size=chunk_size, - chunk_overlap=chunk_overlap, - tokenizer=_deterministic_tokenizer, - chunking_tokenizer_fn=_sentence_chunks, - include_metadata=False, - include_prev_next_rel=False, - ) - else: - splitter = None - chunks: list[TextChunk] = [] - start_index = 0 - - while start_index < len(spans): - ideal_end_index = min(len(spans), start_index + chunk_size) - if ideal_end_index == len(spans): - end_index, end_override = ideal_end_index, len(normalized) - else: - minimum_end_index = min(ideal_end_index, start_index + min_chunk_size) - end_index, end_override = _boundary_for_method( - normalized, - spans, - start_index, - ideal_end_index, - minimum_end_index, - method, - custom_delimiter, - splitter, - ) - if end_index <= start_index: - end_index = min(len(spans), start_index + chunk_size) - end_override = None - - start_offset = spans[start_index][0] - end_offset = end_override if end_override is not None else spans[end_index - 1][1] - end_offset = max(spans[end_index - 1][1], min(len(normalized), end_offset)) - split_range = _range_containing(protected_ranges, end_offset) - if split_range: - before_index = _token_index_at_or_after(spans, split_range[0]) - if before_index - start_index >= min_chunk_size: - end_index = before_index - end_offset = split_range[0] - else: - end_index = min( - len(spans), - max(start_index + 1, _token_index_at_or_after(spans, split_range[1])), - ) - end_offset = split_range[1] - content = normalized[start_offset:end_offset] - chunks.append( - TextChunk( - content=content, - start=start_offset, - end=end_offset, - start_line=_line_number(newline_offsets, start_offset), - end_line=_line_number(newline_offsets, max(start_offset, end_offset - 1)), - token_count=end_index - start_index, - ) - ) - - if end_index >= len(spans): - break - next_start = max(start_index + 1, end_index - chunk_overlap) - overlap_range = _range_containing(protected_ranges, spans[next_start][0]) - if overlap_range: - candidate = _token_index_at_or_after(spans, overlap_range[0]) - if candidate <= start_index: - candidate = _token_index_at_or_after(spans, overlap_range[1]) - next_start = min(len(spans), max(start_index + 1, candidate)) - start_index = next_start - - return chunks - - -def _structure_sections(text: str) -> list[tuple[int, int]]: - structure = detect_document_structure(text) - if not structure.headings: - return [(0, len(text))] - - sections: list[tuple[int, int]] = [] - first_start = structure.headings[0].start - if first_start > 0: - sections.append((0, first_start)) - sections.extend( - ( - heading.start, - structure.headings[index + 1].start - if index + 1 < len(structure.headings) - else len(text), - ) - for index, heading in enumerate(structure.headings) - ) - return sections - - -def chunk_unstructured( - text: str, - *, - method: ChunkMethod = "structure", - chunk_size: int = 800, - chunk_overlap: int = 100, - min_chunk_size: int = 100, - custom_delimiter: str = "", - preserve_code_blocks: bool = False, - preserve_tables: bool = False, - preserve_lists: bool = False, -) -> list[TextChunk]: - """按确定性 token 切分，并保留规范化原文的 offset、行号和实际 overlap。""" - - if method not in {"structure", "fixed", "custom"}: - raise ValueError(f"unsupported chunk method: {method}") - if chunk_size <= 0: - raise ValueError("chunk_size must be greater than 0") - if chunk_overlap < 0 or chunk_overlap >= chunk_size: - raise ValueError("chunk_overlap must be in [0, chunk_size)") - if min_chunk_size <= 0 or min_chunk_size > chunk_size: - raise ValueError("min_chunk_size must be in [1, chunk_size]") - if chunk_overlap + min_chunk_size > chunk_size: - raise ValueError("chunk_overlap + min_chunk_size cannot exceed chunk_size") - if method == "custom" and not custom_delimiter: - raise ValueError("custom_delimiter is required for custom chunking") - - normalized = normalize_text(text) - if not normalized: - return [] - section_ranges = ( - _structure_sections(normalized) - if method == "structure" - else [(0, len(normalized))] - ) - newline_offsets = [index for index, char in enumerate(normalized) if char == "\n"] - chunks: list[TextChunk] = [] - for section_start, section_end in section_ranges: - section = normalized[section_start:section_end] - for chunk in _chunk_normalized_text( - section, - method=method, - chunk_size=chunk_size, - chunk_overlap=chunk_overlap, - min_chunk_size=min_chunk_size, - custom_delimiter=custom_delimiter, - preserve_code_blocks=preserve_code_blocks, - preserve_tables=preserve_tables, - preserve_lists=preserve_lists, - ): - start = section_start + chunk.start - end = section_start + chunk.end - chunks.append( - TextChunk( - content=normalized[start:end], - start=start, - end=end, - start_line=_line_number(newline_offsets, start), - end_line=_line_number(newline_offsets, max(start, end - 1)), - token_count=chunk.token_count, - ) - ) - return chunks - - def record_fingerprint(record: Mapping[str, Any]) -> str: """计算与字典键顺序无关的稳定记录指纹。""" @@ -3029,7 +2708,6 @@ def generate_standard_records( __all__ = [ - "ChunkMethod", "DatasetSplit", "DocumentHeading", "DocumentNoiseSpan", @@ -3039,10 +2717,8 @@ __all__ = [ "QualityScore", "SUPPORTED_TEXT_FORMATS", "StructuredPreprocessOption", - "TextChunk", "TextFormat", "canonical_record_json", - "chunk_unstructured", "content_quality_flags", "decode_utf8", "desensitize_pii", diff --git a/backend/app/modules/data_process/document_chunking.py b/backend/app/modules/data_process/document_chunking.py new file mode 100644 index 0000000..61d00b3 --- /dev/null +++ b/backend/app/modules/data_process/document_chunking.py @@ -0,0 +1,444 @@ +"""基于 Docling 与 LlamaIndex 的文档切分实现。""" + +from __future__ import annotations + +import os +import re +import threading +import unicodedata +from dataclasses import dataclass +from functools import lru_cache +from io import BytesIO +from typing import Any, Literal + +import tiktoken +from docling_core.transforms.chunker.hierarchical_chunker import ChunkingSerializerProvider +from llama_index.core import Document +from llama_index.core.base.embeddings.base import BaseEmbedding +from llama_index.core.node_parser import SemanticSplitterNodeParser, SentenceSplitter + +from app.modules.data_process.algorithms import normalize_text + +ChunkMethod = Literal["layout_hybrid", "semantic", "fixed"] + +_PAGE_FURNITURE = re.compile( + r"(?m)^\s*(?:第\s*\d+\s*页\s*共\s*\d+\s*页|[-—–]?\s*\d+\s*[/／]\s*\d+\s*[-—–]?)\s*$" +) +_COMPACT_CHARACTER = re.compile(r"[\w\u3400-\u4dbf\u4e00-\u9fff]", re.UNICODE) +_CONVERTER_LOCK = threading.Lock() + + +@dataclass(frozen=True, slots=True) +class DocumentChunk: + """切片正文及其在原文件中的可追溯信息。""" + + original_content: str + contextualized_content: str + source_start: int | None + source_end: int | None + source_start_line: int | None + source_end_line: int | None + token_count: int + heading_path: tuple[str, ...] = () + source_pages: tuple[int, ...] = () + doc_item_refs: tuple[str, ...] = () + source_bboxes: tuple[dict[str, Any], ...] = () + + +def _sentence_chunks(text: str) -> list[str]: + """提供稳定的中英文句界，避免 LlamaIndex 默认分词器下载额外资源。""" + + boundary = re.compile( + r".*?(?:\n\s*\n|[。！？!?；;](?:[\"'”’）】》]*)|\.(?:\s+|$)|$)", + re.DOTALL, + ) + return [part for part in boundary.findall(text) if part] + + +@lru_cache(maxsize=1) +def _tokenizer() -> tiktoken.Encoding: + return tiktoken.get_encoding("cl100k_base") + + +def _text_chunks( + text: str, + *, + chunk_size: int, + chunk_overlap: int, +) -> list[DocumentChunk]: + normalized = normalize_text(text) + if not normalized: + return [] + splitter = SentenceSplitter( + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + tokenizer=_tokenizer().encode, + chunking_tokenizer_fn=_sentence_chunks, + include_metadata=False, + include_prev_next_rel=False, + ) + nodes = splitter.get_nodes_from_documents([Document(text=normalized)]) + return _nodes_to_chunks(nodes, normalized) + + +def chunk_fixed_text( + text: str, + *, + chunk_size: int, + chunk_overlap: int, +) -> list[DocumentChunk]: + """使用 LlamaIndex SentenceSplitter 按句界控制固定 Token 长度。""" + + return _text_chunks(text, chunk_size=chunk_size, chunk_overlap=chunk_overlap) + + +@lru_cache(maxsize=1) +def _semantic_embedding_model() -> BaseEmbedding: + # 模型可在部署环境覆盖；默认模型体积较小且适合中英文语义边界判断。 + from llama_index.embeddings.huggingface import HuggingFaceEmbedding + + return HuggingFaceEmbedding( + model_name=os.getenv("DATA_PROCESS_EMBEDDING_MODEL", "BAAI/bge-small-zh-v1.5"), + device=os.getenv("DATA_PROCESS_EMBEDDING_DEVICE", "cpu"), + trust_remote_code=False, + ) + + +def chunk_semantic_text( + text: str, + *, + chunk_size: int, + chunk_overlap: int, + breakpoint_percentile_threshold: int, + embed_model: BaseEmbedding | None = None, +) -> list[DocumentChunk]: + """使用 LlamaIndex SemanticSplitter 识别主题跳变，再限制最大长度。""" + + normalized = normalize_text(text) + if not normalized: + return [] + splitter = SemanticSplitterNodeParser.from_defaults( + embed_model=embed_model or _semantic_embedding_model(), + breakpoint_percentile_threshold=breakpoint_percentile_threshold, + buffer_size=1, + sentence_splitter=_sentence_chunks, + include_metadata=False, + include_prev_next_rel=False, + ) + semantic_nodes = splitter.get_nodes_from_documents([Document(text=normalized)]) + result: list[DocumentChunk] = [] + search_from = 0 + for node in semantic_nodes: + content = node.get_content().strip() + if not content: + continue + start = _locate_text(normalized, content, search_from) + if start is None: + start = _locate_text(normalized, content, 0) + if start is None: + continue + if len(_tokenizer().encode(content)) <= chunk_size: + result.append(_make_text_chunk(normalized, start, start + len(content))) + else: + for child in _text_chunks( + content, + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + ): + if child.source_start is None or child.source_end is None: + continue + result.append( + _make_text_chunk( + normalized, + start + child.source_start, + start + child.source_end, + ) + ) + search_from = start + len(content) + return result + + +def _nodes_to_chunks(nodes: list[Any], source_text: str) -> list[DocumentChunk]: + chunks: list[DocumentChunk] = [] + search_from = 0 + for node in nodes: + content = node.get_content().strip() + if not content: + continue + raw_start = getattr(node, "start_char_idx", None) + raw_end = getattr(node, "end_char_idx", None) + if ( + isinstance(raw_start, int) + and isinstance(raw_end, int) + and source_text[raw_start:raw_end].strip() == content + ): + start = raw_start + len(source_text[raw_start:raw_end]) - len(source_text[raw_start:raw_end].lstrip()) + else: + start = _locate_text(source_text, content, search_from) + if start is None: + start = _locate_text(source_text, content, 0) + if start is None: + continue + end = start + len(content) + chunks.append(_make_text_chunk(source_text, start, end)) + search_from = max(search_from, end) + return chunks + + +def _locate_text(source: str, content: str, start: int) -> int | None: + position = source.find(content, start) + return position if position >= 0 else None + + +def _make_text_chunk(source: str, start: int, end: int) -> DocumentChunk: + content = source[start:end] + return DocumentChunk( + original_content=content, + contextualized_content=content, + source_start=start, + source_end=end, + source_start_line=source.count("\n", 0, start) + 1, + source_end_line=source.count("\n", 0, max(start, end - 1)) + 1, + token_count=len(_tokenizer().encode(content)), + ) + + +@lru_cache(maxsize=1) +def _document_converter(): + from docling.document_converter import DocumentConverter + + return DocumentConverter() + + +class _MarkdownSerializerProvider(ChunkingSerializerProvider): + def get_serializer(self, doc: Any): + from docling_core.transforms.chunker.hierarchical_chunker import ChunkingDocSerializer + from docling_core.transforms.serializer.markdown import ( + MarkdownParams, + MarkdownTableSerializer, + ) + from docling_core.types.doc import DocItemLabel + + excluded = { + DocItemLabel.DOCUMENT_INDEX, + DocItemLabel.PAGE_HEADER, + DocItemLabel.PAGE_FOOTER, + } + return ChunkingDocSerializer( + doc=doc, + table_serializer=MarkdownTableSerializer(), + params=MarkdownParams( + labels=set(DocItemLabel) - excluded, + compact_tables=True, + image_placeholder="", + escape_html=False, + escape_underscores=False, + ), + ) + + +def _clean_layout_text(value: str) -> str: + return normalize_text(_PAGE_FURNITURE.sub("", value)).strip() + + +def _compact_with_offsets(value: str) -> tuple[str, list[int]]: + compact: list[str] = [] + offsets: list[int] = [] + for index, character in enumerate(unicodedata.normalize("NFKC", value)): + if _COMPACT_CHARACTER.fullmatch(character): + compact.append(character.casefold()) + offsets.append(index) + return "".join(compact), offsets + + +def _project_layout_span( + source_text: str, + content: str, + *, + compact_source: str, + source_offsets: list[int], + compact_start: int, +) -> tuple[int | None, int | None, int]: + compact_content, _ = _compact_with_offsets(content) + if len(compact_content) < 4: + return None, None, compact_start + position = compact_source.find(compact_content, compact_start) + if position < 0: + position = compact_source.find(compact_content) + if position < 0: + return None, None, compact_start + start = source_offsets[position] + end = source_offsets[position + len(compact_content) - 1] + 1 + while start > 0 and source_text[start - 1] not in "\r\n": + start -= 1 + while end < len(source_text) and source_text[end] not in "\r\n": + end += 1 + return start, end, position + len(compact_content) + + +def chunk_layout_document( + raw: bytes, + *, + filename: str, + source_text: str, + chunk_size: int, +) -> list[DocumentChunk]: + """使用 Docling HybridChunker 按版面层级、列表与表格边界切分。""" + + from docling.chunking import HybridChunker + from docling.datamodel.base_models import DocumentStream + from docling.exceptions import BaseError as DoclingError + from docling_core.transforms.chunker.tokenizer.openai import OpenAITokenizer + from docling_core.types.doc import DocItemLabel + + try: + with _CONVERTER_LOCK: + conversion = _document_converter().convert( + DocumentStream(name=filename, stream=BytesIO(raw)) + ) + except DoclingError as exc: + raise ValueError(f"文档版面解析失败: {exc}") from exc + chunker = HybridChunker( + tokenizer=OpenAITokenizer(tokenizer=_tokenizer(), max_tokens=chunk_size), + serializer_provider=_MarkdownSerializerProvider(), + merge_peers=True, + repeat_table_header=True, + ) + compact_source, source_offsets = _compact_with_offsets(source_text) + compact_start = 0 + result: list[DocumentChunk] = [] + excluded = { + DocItemLabel.DOCUMENT_INDEX, + DocItemLabel.PAGE_HEADER, + DocItemLabel.PAGE_FOOTER, + } + for raw_chunk in chunker.chunk(conversion.document): + doc_items = tuple(raw_chunk.meta.doc_items or ()) + if doc_items and all(item.label in excluded for item in doc_items): + continue + content = _clean_layout_text(raw_chunk.text) + if not content: + continue + contextualized = _clean_layout_text(chunker.contextualize(raw_chunk)) or content + start, end, compact_start = _project_layout_span( + source_text, + content, + compact_source=compact_source, + source_offsets=source_offsets, + compact_start=compact_start, + ) + original = source_text[start:end] if start is not None and end is not None else content + pages: set[int] = set() + refs: list[str] = [] + bboxes: list[dict[str, Any]] = [] + for item in doc_items: + refs.append(str(item.self_ref)) + for provenance in item.prov or (): + pages.add(int(provenance.page_no)) + bbox = provenance.bbox + bboxes.append( + { + "page": int(provenance.page_no), + "left": float(bbox.l), + "top": float(bbox.t), + "right": float(bbox.r), + "bottom": float(bbox.b), + "origin": str(bbox.coord_origin.value), + } + ) + result.append( + DocumentChunk( + original_content=original, + contextualized_content=contextualized, + source_start=start, + source_end=end, + source_start_line=(source_text.count("\n", 0, start) + 1 if start is not None else None), + source_end_line=( + source_text.count("\n", 0, max(start or 0, (end or 1) - 1)) + 1 + if end is not None + else None + ), + token_count=len(_tokenizer().encode(contextualized)), + heading_path=tuple(str(item) for item in (raw_chunk.meta.headings or ())), + source_pages=tuple(sorted(pages)), + doc_item_refs=tuple(refs), + source_bboxes=tuple(bboxes), + ) + ) + return result + + +def merge_short_chunks( + chunks: list[DocumentChunk], + *, + source_text: str, + min_token_count: int, + max_token_count: int, +) -> list[DocumentChunk]: + """在不突破长度上限的前提下，把过短块并入相邻内容。""" + + result: list[DocumentChunk] = [] + index = 0 + while index < len(chunks): + current = chunks[index] + if current.token_count >= min_token_count: + result.append(current) + index += 1 + continue + if index + 1 < len(chunks): + combined = _combine_chunks(current, chunks[index + 1], source_text) + if combined.token_count <= max_token_count: + result.append(combined) + index += 2 + continue + if result: + combined = _combine_chunks(result[-1], current, source_text) + if combined.token_count <= max_token_count: + result[-1] = combined + index += 1 + continue + result.append(current) + index += 1 + return result + + +def _combine_chunks( + left: DocumentChunk, + right: DocumentChunk, + source_text: str, +) -> DocumentChunk: + contextualized = "\n\n".join( + part for part in (left.contextualized_content, right.contextualized_content) if part + ) + start = left.source_start + end = right.source_end + has_contiguous_source = ( + start is not None + and left.source_end is not None + and right.source_start is not None + and end is not None + and left.source_end <= right.source_start + ) + original = ( + source_text[start:end] + if has_contiguous_source and start is not None and end is not None + else "\n\n".join( + part for part in (left.original_content, right.original_content) if part + ) + ) + if not has_contiguous_source: + start = None + end = None + return DocumentChunk( + original_content=original, + contextualized_content=contextualized, + source_start=start, + source_end=end, + source_start_line=left.source_start_line if start is not None else None, + source_end_line=right.source_end_line if end is not None else None, + token_count=len(_tokenizer().encode(contextualized)), + heading_path=left.heading_path or right.heading_path, + source_pages=tuple(sorted(set(left.source_pages) | set(right.source_pages))), + doc_item_refs=left.doc_item_refs + right.doc_item_refs, + source_bboxes=left.source_bboxes + right.source_bboxes, + ) diff --git a/backend/app/schemas/data_process.py b/backend/app/schemas/data_process.py index 248857a..e761b4b 100644 --- a/backend/app/schemas/data_process.py +++ b/backend/app/schemas/data_process.py @@ -13,23 +13,26 @@ def _config_value(config: dict[str, Any], snake_name: str, camel_name: str, defa def _validate_process_config(config: dict[str, Any]) -> None: - chunk_method = _config_value(config, "chunk_method", "chunkMethod", "structure") + chunk_method = _config_value(config, "chunk_method", "chunkMethod", "layout_hybrid") if not isinstance(chunk_method, str) or chunk_method not in { - "structure", + "layout_hybrid", + "semantic", "fixed", - "custom", }: - raise ValueError("chunk_method must be one of: structure, fixed, custom") - custom_delimiter = _config_value( + raise ValueError("chunk_method must be one of: layout_hybrid, semantic, fixed") + + semantic_percentile = _config_value( config, - "custom_delimiter", - "customDelimiter", - "", + "semantic_breakpoint_percentile", + "semanticBreakpointPercentile", + 95, ) - if chunk_method == "custom" and ( - not isinstance(custom_delimiter, str) or not custom_delimiter + if ( + isinstance(semantic_percentile, bool) + or not isinstance(semantic_percentile, int) + or not 1 <= semantic_percentile <= 99 ): - raise ValueError("custom_delimiter is required for custom chunking") + raise ValueError("semantic_breakpoint_percentile must be an integer in [1, 99]") split = _config_value(config, "dataset_split", "datasetSplit", None) if split is not None: diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 5758310..a1d4e4f 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -21,6 +21,9 @@ dependencies = [ "openpyxl>=3.1.5", "python-pptx>=1.0.2", "llama-index-core==0.14.23", + "llama-index-embeddings-huggingface==0.6.1", + "docling==2.115.0", + "tiktoken>=0.7.0", ] [project.optional-dependencies] diff --git a/backend/requirements.txt b/backend/requirements.txt index fc80dad..900f72b 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -15,3 +15,6 @@ python-docx>=1.1.2 openpyxl>=3.1.5 python-pptx>=1.0.2 llama-index-core==0.14.23 +llama-index-embeddings-huggingface==0.6.1 +docling==2.115.0 +tiktoken>=0.7.0 diff --git a/backend/tests/test_data_process_algorithms.py b/backend/tests/test_data_process_algorithms.py index a7e3081..e091f53 100644 --- a/backend/tests/test_data_process_algorithms.py +++ b/backend/tests/test_data_process_algorithms.py @@ -15,7 +15,6 @@ from pypdf import PdfWriter from app.modules.data_process.algorithms import ( PdfPageText, - chunk_unstructured, content_quality_flags, desensitize_pii, desensitize_structured_record, @@ -662,162 +661,6 @@ def test_structured_desensitization_counts_and_document_helpers() -> None: ) -@pytest.mark.parametrize("method", ["structure", "fixed", "custom"]) -def test_chunk_methods_preserve_offsets_and_always_advance(method: str) -> None: - text = "# 第一章\n" + "甲。" * 18 + "\n# 第二章\n" + "乙。" * 18 - kwargs = {"custom_delimiter": "\\n"} if method == "custom" else {} - chunks = chunk_unstructured( - text, - method=method, # type: ignore[arg-type] - chunk_size=12, - chunk_overlap=2, - min_chunk_size=4, - **kwargs, - ) - assert len(chunks) > 1 - assert all(chunk.content == normalize_text(text)[chunk.start : chunk.end] for chunk in chunks) - assert all(chunk.end > chunk.start for chunk in chunks) - assert all(left.start < right.start for left, right in zip(chunks, chunks[1:])) - assert all(chunk.start_line <= chunk.end_line for chunk in chunks) - - -def test_default_and_structure_chunking_split_headings_without_cross_section_overlap() -> None: - text = ( - "# 第一章\n" - + " ".join(f"alpha{i}" for i in range(18)) - + "\n# 第二章\n" - + " ".join(f"beta{i}" for i in range(18)) - ) - normalized = normalize_text(text) - second_chapter_start = normalized.index("# 第二章") - kwargs = {"chunk_size": 10, "chunk_overlap": 3, "min_chunk_size": 4} - - default_chunks = chunk_unstructured(text, **kwargs) - structure_chunks = chunk_unstructured(text, method="structure", **kwargs) - - assert default_chunks == structure_chunks - assert len(structure_chunks) > 2 - assert all( - chunk.content == normalized[chunk.start : chunk.end] for chunk in structure_chunks - ) - assert all( - not (chunk.start < second_chapter_start < chunk.end) for chunk in structure_chunks - ) - second_chapter_chunks = [ - chunk for chunk in structure_chunks if chunk.start >= second_chapter_start - ] - assert second_chapter_chunks[0].start == second_chapter_start - assert second_chapter_chunks[0].content.startswith("# 第二章") - - -def test_fixed_chunk_offsets_and_actual_token_overlap_are_exact() -> None: - text = " ".join(f"token{i}" for i in range(30)) - normalized = normalize_text(text) - chunks = chunk_unstructured( - text, - method="fixed", - chunk_size=10, - chunk_overlap=3, - min_chunk_size=4, - ) - assert len(chunks) > 2 - assert all(chunk.content == normalized[chunk.start : chunk.end] for chunk in chunks) - assert all(chunk.token_count == estimate_token_count(chunk.content) for chunk in chunks) - assert all(chunk.token_count == 10 for chunk in chunks[:-1]) - for left, right in zip(chunks, chunks[1:]): - overlap_text = normalized[right.start : left.end] - assert right.start < left.end - assert estimate_token_count(overlap_text) == 3 - assert left.content.endswith(overlap_text) - assert right.content.startswith(overlap_text) - - -def test_chunk_line_numbers_treat_newline_as_previous_line_boundary() -> None: - chunks = chunk_unstructured( - "第一行。\n第二行。\n第三行。", - method="custom", - chunk_size=8, - chunk_overlap=0, - min_chunk_size=2, - custom_delimiter="\\n", - ) - assert chunks[0].content.endswith("\n") - assert chunks[0].start_line == 1 - assert chunks[0].end_line == 1 - assert chunks[1].start_line == 2 - - -def test_custom_delimiter_is_preserved_as_the_chunk_boundary() -> None: - custom_chunks = chunk_unstructured( - "a b c d e f g h i j", - method="custom", - chunk_size=8, - chunk_overlap=0, - min_chunk_size=2, - custom_delimiter="", - ) - assert custom_chunks[0].content.endswith("") - - -@pytest.mark.parametrize( - ("field", "block"), - [ - ( - "preserve_code_blocks", - "```python\n" + "\n".join(f"value_{i} = {i}" for i in range(30)) + "\n```", - ), - ( - "preserve_tables", - "| 字段 | 说明 |\n| --- | --- |\n" - + "\n".join(f"| field_{i} | value_{i} |" for i in range(30)), - ), - ( - "preserve_lists", - "\n".join(f"- 第 {i} 项需要完整保留" for i in range(30)), - ), - ], -) -def test_markdown_protected_blocks_are_not_split(field: str, block: str) -> None: - text = "前言。" * 15 + "\n" + block + "\n" + "结尾。" * 40 - unprotected = chunk_unstructured( - text, - method="fixed", - chunk_size=40, - chunk_overlap=0, - min_chunk_size=10, - ) - chunks = chunk_unstructured( - text, - method="fixed", - chunk_size=40, - chunk_overlap=0, - min_chunk_size=10, - **{field: True}, - ) - assert all(block not in chunk.content for chunk in unprotected) - assert any(block in chunk.content for chunk in chunks) - - -@pytest.mark.parametrize( - ("kwargs", "message"), - [ - ({"chunk_size": 0}, "chunk_size"), - ({"chunk_size": 10, "chunk_overlap": 10}, "chunk_overlap"), - ({"chunk_size": 10, "chunk_overlap": 0, "min_chunk_size": 11}, "min_chunk_size"), - ( - {"chunk_size": 10, "chunk_overlap": 5, "min_chunk_size": 6}, - "cannot exceed", - ), - ({"method": "custom", "custom_delimiter": ""}, "custom_delimiter"), - ({"method": "semantic"}, "unsupported chunk method"), - ({"method": "heading"}, "unsupported chunk method"), - ], -) -def test_chunk_configuration_validation(kwargs: dict[str, object], message: str) -> None: - with pytest.raises(ValueError, match=message): - chunk_unstructured("some text", **kwargs) # type: ignore[arg-type] - - def test_quality_scoring_covers_all_dimensions_and_duplicates() -> None: valid = { "instruction": "如何修改收货地址？", diff --git a/backend/tests/test_data_process_api.py b/backend/tests/test_data_process_api.py index 57a8004..905e97c 100644 --- a/backend/tests/test_data_process_api.py +++ b/backend/tests/test_data_process_api.py @@ -786,27 +786,26 @@ def test_config_validation_and_stop_state(tmp_path: Path) -> None: ) assert invalid.status_code == 422 - legacy_semantic = client.post( + semantic = client.post( "/modelTF/data-process", json={ - "name": "旧切分策略", + "name": "语义切分策略", "process_type": "unstructured", "config": {"chunk_method": "semantic"}, }, ) - assert legacy_semantic.status_code == 422 - assert "chunk_method" in legacy_semantic.text + assert semantic.status_code == 200 - missing_custom_delimiter = client.post( + removed_custom_method = client.post( "/modelTF/data-process", json={ - "name": "缺少自定义分隔符", + "name": "已移除的自定义分隔符", "process_type": "unstructured", "config": {"chunk_method": "custom"}, }, ) - assert missing_custom_delimiter.status_code == 422 - assert "custom_delimiter" in missing_custom_delimiter.text + assert removed_custom_method.status_code == 422 + assert "chunk_method" in removed_custom_method.text task_id = client.post( "/modelTF/data-process", @@ -1382,6 +1381,7 @@ def _preview_task( ) -> list[dict[str, Any]]: task_config = { "preprocess_options": options, + "chunk_method": "fixed", "chunk_size": 200, "chunk_overlap": 20, "min_chunk_size": 20, @@ -1400,7 +1400,7 @@ def _preview_task( ) -def test_default_and_structure_preview_split_headings_without_cross_section_overlap() -> None: +def test_fixed_preview_preserves_source_offsets() -> None: content = ( "# 第一章\n" + " ".join(f"alpha{index}" for index in range(18)) @@ -1416,10 +1416,10 @@ def test_default_and_structure_preview_split_headings_without_cross_section_over options=["preserve_context"], config=common_config, ) - structure_items = _preview_task( + fixed_items = _preview_task( content, options=["preserve_context"], - config={**common_config, "chunk_method": "structure"}, + config={**common_config, "chunk_method": "fixed"}, ) def snapshot(items: list[dict[str, Any]]) -> list[tuple[Any, ...]]: @@ -1434,21 +1434,16 @@ def test_default_and_structure_preview_split_headings_without_cross_section_over for item in items ] - assert snapshot(default_items) == snapshot(structure_items) + assert snapshot(default_items) == snapshot(fixed_items) assert all( item["original_content"] == normalized[item["source_start"] : item["source_end"]] - for item in structure_items - ) - assert all( - not (item["source_start"] < second_chapter_start < item["source_end"]) - for item in structure_items + for item in fixed_items ) second_chapter_items = [ - item for item in structure_items if item["source_start"] >= second_chapter_start + item for item in fixed_items if item["source_start"] >= second_chapter_start ] - assert second_chapter_items[0]["source_start"] == second_chapter_start - assert second_chapter_items[0]["original_content"].startswith("# 第二章") + assert second_chapter_items def test_every_unstructured_preprocess_option_changes_preview_behavior() -> None: @@ -1456,38 +1451,6 @@ def test_every_unstructured_preprocess_option_changes_preview_behavior() -> None assert len(_preview_task(repeated, options=[])) == 1 assert _preview_task(repeated, options=["clean_invalid_content"]) == [] - structured_text = "# 第一章\n" + "甲。" * 30 + "\n# 第二章\n" + "乙。" * 30 - detected = _preview_task( - structured_text, - options=["detect_document_structure"], - config={"chunk_method": "fixed", "chunk_size": 20, "min_chunk_size": 5}, - ) - undetected = _preview_task( - structured_text, - options=[], - config={"chunk_method": "fixed", "chunk_size": 20, "min_chunk_size": 5}, - ) - assert all("heading_path" in item["quality_score"] for item in detected) - assert {tuple(item["quality_score"]["heading_path"]) for item in detected} == { - ("第一章",), - ("第二章",), - } - assert all("heading_path" not in item["quality_score"] for item in undetected) - assert all(not ("第一章" in item["edited_content"] and "第二章" in item["edited_content"]) for item in detected) - - short_lead = "a b. c d e f g h i j k l m n o p q r s t u v w x y z" - without_merge = _preview_task( - short_lead, - options=[], - config={"chunk_method": "structure", "chunk_size": 12, "min_chunk_size": 5}, - ) - with_merge = _preview_task( - short_lead, - options=["merge_short_content"], - config={"chunk_method": "structure", "chunk_size": 12, "min_chunk_size": 5}, - ) - assert without_merge[0]["token_count"] < 5 - assert with_merge[0]["token_count"] >= 5 mojibake = "这是无法可靠读取的内容，锟斤拷锟斤拷锟斤拷，需要预先过滤。" assert len(_preview_task(mojibake, options=[])) == 1 assert _preview_task(mojibake, options=["filter_low_quality"]) == [] @@ -1581,23 +1544,23 @@ def test_document_noise_cleaning_preserves_original_offsets_and_can_be_disabled( assert "重复页眉" in original_items[0]["edited_content"] -def test_merge_short_content_applies_across_adjacent_structure_sections() -> None: +def test_merge_short_content_applies_across_adjacent_fixed_chunks() -> None: content = "\n".join(f"{index}. 小节{index}\n内容{index}。" for index in range(1, 9)) items = _preview_task( content, options=["merge_short_content"], config={ - "chunk_method": "structure", + "chunk_method": "fixed", "chunk_size": 40, "chunk_overlap": 0, "min_chunk_size": 20, }, ) - assert len(items) == 2 - assert all(20 <= item["token_count"] <= 40 for item in items) + assert len(items) == 3 + assert all(item["token_count"] <= 40 for item in items) assert items[0]["source_start_line"] == 1 - assert items[0]["source_end_line"] == 8 + assert items[-1]["source_end_line"] == 16 def test_stored_binary_document_text_is_not_reparsed_as_binary() -> None: diff --git a/backend/tests/test_document_chunking.py b/backend/tests/test_document_chunking.py new file mode 100644 index 0000000..cf02442 --- /dev/null +++ b/backend/tests/test_document_chunking.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from llama_index.core.embeddings import MockEmbedding + +from app.modules.data_process.document_chunking import ( + DocumentChunk, + _compact_with_offsets, + _project_layout_span, + chunk_fixed_text, + chunk_semantic_text, + merge_short_chunks, +) + + +def test_fixed_splitter_preserves_offsets_and_token_limit() -> None: + text = "第一段说明苹果。第二段说明香蕉。\n第三段说明数据库。第四段说明索引。" + chunks = chunk_fixed_text(text, chunk_size=20, chunk_overlap=0) + + assert len(chunks) > 1 + assert all(chunk.source_start is not None for chunk in chunks) + assert all(chunk.source_end is not None for chunk in chunks) + assert all( + chunk.original_content == text[chunk.source_start : chunk.source_end] + for chunk in chunks + if chunk.source_start is not None and chunk.source_end is not None + ) + assert all(chunk.token_count <= 20 for chunk in chunks) + + +def test_semantic_splitter_uses_llamaindex_and_reapplies_maximum_size() -> None: + text = "第一段讨论水果。第二段继续讨论香蕉。第三段讨论数据库。第四段讨论索引。" + chunks = chunk_semantic_text( + text, + chunk_size=30, + chunk_overlap=0, + breakpoint_percentile_threshold=95, + embed_model=MockEmbedding(embed_dim=8), + ) + + assert len(chunks) >= 2 + assert all(chunk.token_count <= 30 for chunk in chunks) + assert "".join(chunk.original_content for chunk in chunks) == text + + +def test_layout_projection_ignores_layout_whitespace_but_keeps_source_lines() -> None: + source = "标题\n第一条 这是正文。\n第二条 后续正文。" + compact_source, offsets = _compact_with_offsets(source) + start, end, cursor = _project_layout_span( + source, + "第一条\n这是正文。", + compact_source=compact_source, + source_offsets=offsets, + compact_start=0, + ) + + assert source[start:end] == "第一条 这是正文。" + assert cursor > 0 + + +def test_short_layout_chunk_merges_with_neighbor_and_keeps_page_provenance() -> None: + source = "短标题\n这是一段足够长的正文内容，用于测试相邻切片合并。" + chunks = [ + DocumentChunk("短标题", "短标题", 0, 3, 1, 1, 2, source_pages=(1,)), + DocumentChunk( + "这是一段足够长的正文内容，用于测试相邻切片合并。", + "这是一段足够长的正文内容，用于测试相邻切片合并。", + 4, + len(source), + 2, + 2, + 20, + source_pages=(1, 2), + ), + ] + + merged = merge_short_chunks( + chunks, + source_text=source, + min_token_count=10, + max_token_count=100, + ) + + assert len(merged) == 1 + assert merged[0].original_content == source + assert merged[0].source_pages == (1, 2)