feat(data-process): 接入三种文档切分引擎

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)