first commit

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-16 11:12:17 +08:00
commit 4003624b8c
80 changed files with 9990 additions and 0 deletions
@@ -0,0 +1,120 @@
+"""End-to-end document chunking pipeline."""
+
+from __future__ import annotations
+
+import hashlib
+import shutil
+import uuid
+from pathlib import Path
+
+from rag_cut.layout_meta import enrich_layout_metadata
+from rag_cut.models import Chunk, ChunkResult, SplitConfig, SplitMode
+from rag_cut.parsers.pdf.noise_filter import filter_toc_blocks
+from rag_cut.parsers.pdf.table_merge import merge_cross_page_tables
+from rag_cut.parsers.registry import get_parser
+from rag_cut.renderer import render_blocks
+from rag_cut.split_policy import choose_split_config, split_config_summary
+from rag_cut.splitters import split_by_delimiter, split_by_row
+from rag_cut.splitters.default_splitter import split_default_with_meta
+from rag_cut.splitters.heading_splitter import assign_parent_chunk_ids, chunk_groups_to_block_groups
+from rag_cut.splitters.parent_child import CHUNK_STRATEGY as PARENT_CHILD_STRATEGY
+from rag_cut.splitters.parent_child import split_by_parent_child
+from rag_cut.splitters.pdf_semantic import CHUNK_STRATEGY, split_pdf_semantic
+from rag_cut.splitters.pdf_strategy import choose_pdf_chunk_strategy
+
+STORAGE_ROOT = Path(__file__).resolve().parent.parent.parent / "storage"
+PDF_SEMANTIC_EXTS = {".pdf", ".doc", ".docx"}
+
+
+def _doc_id(path: Path) -> str:
+    digest = hashlib.md5(f"{path.name}-{path.stat().st_mtime}".encode()).hexdigest()[:12]
+    return digest
+
+
+def chunk_document(
+    path: Path | str,
+    config: SplitConfig | None = None,
+    storage_root: Path | None = None,
+) -> ChunkResult:
+    path = Path(path)
+    if not path.exists():
+        raise FileNotFoundError(path)
+
+    root = storage_root or STORAGE_ROOT
+    doc_id = _doc_id(path)
+    assets_dir = root / "assets" / doc_id
+    uploads_dir = root / "uploads" / doc_id
+    uploads_dir.mkdir(parents=True, exist_ok=True)
+
+    stored = uploads_dir / path.name
+    if path.resolve() != stored.resolve():
+        shutil.copy2(path, stored)
+
+    parser = get_parser(path)
+    blocks = merge_cross_page_tables(parser.parse(stored, assets_dir))
+    blocks = enrich_layout_metadata(blocks)
+    # All formats: never parse/chunk document directories (目录 / Contents).
+    blocks = filter_toc_blocks(blocks)
+    config = config or choose_split_config(path, blocks)
+
+    for block in blocks:
+        if block.image_path:
+            rel = Path(block.image_path)
+            if "crops" in rel.parts:
+                block.image_path = f"assets/{doc_id}/crops/{rel.name}"
+            else:
+                block.image_path = f"assets/{doc_id}/{rel.name}"
+        crop = block.meta.get("crop_path")
+        if crop:
+            crop_name = Path(crop).name
+            url = f"assets/{doc_id}/crops/{crop_name}"
+            block.meta["crop_path"] = url
+            if block.type.value == "table" and not block.image_path:
+                block.image_path = url
+
+    ext = path.suffix.lower()
+    pdf_strategy = None
+    if ext in PDF_SEMANTIC_EXTS and config.mode == SplitMode.DEFAULT:
+        pdf_strategy = choose_pdf_chunk_strategy(path, blocks)
+
+    group_metas: list[dict] = []
+    if ext in PDF_SEMANTIC_EXTS and config.mode == SplitMode.DEFAULT:
+        chunks = split_pdf_semantic(blocks, config)
+    elif config.mode == SplitMode.BY_ROW:
+        groups = split_by_row(blocks, config)
+        chunks = []
+    elif config.mode == SplitMode.DELIMITER:
+        groups = split_by_delimiter(blocks, config)
+        chunks = []
+    elif config.mode == SplitMode.PARENT_CHILD:
+        parent_child_groups = split_by_parent_child(blocks, config)
+        groups, group_metas = chunk_groups_to_block_groups(parent_child_groups)
+        chunks = []
+    else:
+        groups, group_metas = split_default_with_meta(blocks, config)
+        chunks = []
+
+    if not chunks:
+        for i, group in enumerate(groups):
+            meta: dict = dict(group_metas[i]) if i < len(group_metas) else {}
+            chunks.append(render_blocks(group, index=i, meta=meta))
+
+    chunks = assign_parent_chunk_ids(chunks)
+
+    split_config = split_config_summary(config)
+    if pdf_strategy:
+        split_config["pdf_chunk_strategy"] = pdf_strategy
+        split_config["chunk_strategy"] = CHUNK_STRATEGY
+    elif config.mode == SplitMode.PARENT_CHILD:
+        split_config["chunk_strategy"] = PARENT_CHILD_STRATEGY
+
+    return ChunkResult(
+        filename=path.name,
+        doc_id=doc_id,
+        split_mode=config.mode,
+        split_config=split_config,
+        block_count=len(blocks),
+        chunk_count=len(chunks),
+        chunks=chunks,
+        assets_dir=str(assets_dir) if assets_dir.exists() else None,
+    )