first commit

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
陈辅元
2026-07-16 11:12:17 +08:00
co-authored by Cursor
commit 4003624b8c
80 changed files with 9990 additions and 0 deletions
+1
View File
@@ -0,0 +1 @@
"""RAG-cut FastAPI app. Start with: python backend/run.py"""
+288
View File
@@ -0,0 +1,288 @@
"""FastAPI endpoints for document chunking + demo frontend."""
from __future__ import annotations
import asyncio
import json
import mimetypes
import re
import shutil
import tempfile
import time
from datetime import datetime, timezone
from pathlib import Path
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field
from rag_cut.models import ChunkResult, SplitConfig, SplitMode
from rag_cut.pipeline import chunk_document
from rag_cut.retrieval import recall_chunks
ROOT = Path(__file__).resolve().parent.parent.parent
FRONTEND_DIR = ROOT / "frontend"
STORAGE_DIR = ROOT / "storage"
RESULTS_DIR = STORAGE_DIR / "results"
UPLOADS_DIR = STORAGE_DIR / "uploads"
DOC_ID_RE = re.compile(r"^[a-f0-9]{12}$")
WORD_EXTS = {".doc", ".docx"}
app = FastAPI(title="RAG-cut Chunking API", version="0.1.0")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=False,
allow_methods=["*"],
allow_headers=["*"],
)
(STORAGE_DIR / "assets").mkdir(parents=True, exist_ok=True)
RESULTS_DIR.mkdir(parents=True, exist_ok=True)
UPLOADS_DIR.mkdir(parents=True, exist_ok=True)
app.mount("/assets", StaticFiles(directory=STORAGE_DIR / "assets"), name="assets")
class RecallRequest(BaseModel):
doc_id: str
query: str = Field(min_length=1, max_length=500)
top_k: int = Field(default=5, ge=1, le=20)
def _result_path(doc_id: str) -> Path:
if not DOC_ID_RE.fullmatch(doc_id):
raise HTTPException(status_code=400, detail="Invalid doc_id")
return RESULTS_DIR / f"{doc_id}.json"
def _save_result(result: ChunkResult) -> None:
_result_path(result.doc_id).write_text(result.model_dump_json(indent=2), encoding="utf-8")
def _find_original(doc_id: str) -> Path | None:
uploads_dir = UPLOADS_DIR / doc_id
if not uploads_dir.is_dir():
return None
files = sorted(p for p in uploads_dir.iterdir() if p.is_file())
return files[0] if files else None
def _find_preview(doc_id: str) -> Path | None:
"""Return a browser-previewable source, preferring converted Word PDFs."""
original = _find_original(doc_id)
if original is None:
return None
if original.suffix.lower() == ".pdf":
return original
if original.suffix.lower() not in WORD_EXTS:
return None
conversion_dir = STORAGE_DIR / "assets" / doc_id / "_conversion"
exact = conversion_dir / f"{original.stem}.pdf"
if exact.is_file():
return exact
if conversion_dir.is_dir():
return next(iter(sorted(conversion_dir.glob("*.pdf"))), None)
return None
def _load_result(doc_id: str) -> ChunkResult:
result_path = _result_path(doc_id)
if not result_path.exists():
raise HTTPException(status_code=404, detail="Chunk result not found")
try:
return ChunkResult.model_validate_json(result_path.read_text(encoding="utf-8"))
except Exception as exc:
raise HTTPException(status_code=500, detail=f"Invalid result file: {exc}") from exc
def _result_summary(path: Path) -> dict | None:
"""Build list summary without loading the full chunks array."""
try:
with path.open("r", encoding="utf-8") as fh:
head = fh.read(32768)
idx = head.find('\n "chunks"')
if idx == -1:
idx = head.find('"chunks"')
if idx != -1:
data = json.loads(head[:idx].rstrip().rstrip(",") + "\n}")
else:
data = json.loads(path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError):
return None
doc_id = data.get("doc_id") or path.stem
if not DOC_ID_RE.fullmatch(str(doc_id)):
return None
mtime = path.stat().st_mtime
return {
"doc_id": doc_id,
"filename": data.get("filename") or path.name,
"split_mode": data.get("split_mode"),
"chunk_count": data.get("chunk_count", 0),
"block_count": data.get("block_count", 0),
"has_original": _find_original(doc_id) is not None,
"saved_at": datetime.fromtimestamp(mtime, tz=timezone.utc).isoformat(),
"saved_ts": mtime,
}
@app.get("/health")
def health():
return {"status": "ok"}
@app.get("/api/results")
def list_results():
"""List persisted chunk results (newest first)."""
items: list[dict] = []
for path in RESULTS_DIR.glob("*.json"):
summary = _result_summary(path)
if summary:
items.append(summary)
items.sort(key=lambda item: item["saved_ts"], reverse=True)
for item in items:
item.pop("saved_ts", None)
return {"count": len(items), "results": items}
@app.get("/api/results/{doc_id}")
def get_result(doc_id: str):
"""Load a full persisted chunk result by doc_id."""
chunk_result = _load_result(doc_id)
payload = json.loads(chunk_result.model_dump_json())
payload["has_original"] = _find_original(doc_id) is not None
return JSONResponse(content=payload)
@app.get("/api/results/{doc_id}/original")
def get_original(doc_id: str):
"""Serve the original uploaded file for a historical result (if still on disk)."""
if not DOC_ID_RE.fullmatch(doc_id):
raise HTTPException(status_code=400, detail="Invalid doc_id")
original = _find_original(doc_id)
if original is None:
raise HTTPException(status_code=404, detail="Original file not found")
media_type, _ = mimetypes.guess_type(original.name)
return FileResponse(
path=original,
media_type=media_type or "application/octet-stream",
filename=original.name,
)
@app.get("/api/results/{doc_id}/preview")
def get_preview(doc_id: str):
"""Serve an inline PDF preview for PDF and converted Word documents."""
if not DOC_ID_RE.fullmatch(doc_id):
raise HTTPException(status_code=400, detail="Invalid doc_id")
preview = _find_preview(doc_id)
if preview is None:
raise HTTPException(status_code=404, detail="PDF preview not found")
return FileResponse(path=preview, media_type="application/pdf")
def _delete_tree(path: Path) -> bool:
if not path.exists():
return False
if path.is_dir():
shutil.rmtree(path)
else:
path.unlink()
return True
@app.delete("/api/results/{doc_id}")
def delete_result(doc_id: str):
"""Delete a persisted chunk result and its uploads/assets."""
result_path = _result_path(doc_id)
if not result_path.exists():
raise HTTPException(status_code=404, detail="Chunk result not found")
removed = {
"result": _delete_tree(result_path),
"uploads": _delete_tree(UPLOADS_DIR / doc_id),
"assets": _delete_tree(STORAGE_DIR / "assets" / doc_id),
}
return {"doc_id": doc_id, "deleted": True, "removed": removed}
@app.post("/api/chunk")
async def chunk_file(
file: UploadFile = File(...),
mode: str | None = Form(None),
delimiter: str | None = Form(None),
parent_delimiter: str | None = Form(None),
child_delimiter: str | None = Form(None),
max_chunk_size: int | None = Form(None),
child_max_size: int | None = Form(None),
overlap: int | None = Form(None),
header_row_start: int | None = Form(None),
header_row_end: int | None = Form(None),
start_row: int | None = Form(None),
rows_per_chunk: int | None = Form(None),
):
config = None
if mode is not None:
try:
split_mode = SplitMode(mode)
except ValueError as exc:
raise HTTPException(status_code=400, detail=f"Invalid mode: {mode}") from exc
config = SplitConfig(
mode=split_mode,
delimiter=delimiter,
parent_delimiter=parent_delimiter,
child_delimiter=child_delimiter,
max_chunk_size=1500 if max_chunk_size is None else max_chunk_size,
child_max_size=512 if child_max_size is None else child_max_size,
overlap=150 if overlap is None else overlap,
header_row_start=1 if header_row_start is None else header_row_start,
header_row_end=1 if header_row_end is None else header_row_end,
start_row=2 if start_row is None else start_row,
rows_per_chunk=1 if rows_per_chunk is None else rows_per_chunk,
)
filename = Path(file.filename or "upload").name
payload = await file.read()
with tempfile.TemporaryDirectory() as tmp:
tmp_path = Path(tmp) / filename
tmp_path.write_bytes(payload)
try:
# Heavy parsing (MinerU / LibreOffice / OCR) must not block the event loop,
# otherwise /health and other uploads hang while one document is processing.
result = await asyncio.to_thread(chunk_document, tmp_path, config)
_save_result(result)
return JSONResponse(content=json.loads(result.model_dump_json()))
except ValueError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
except FileNotFoundError as exc:
raise HTTPException(status_code=404, detail=str(exc)) from exc
except RuntimeError as exc:
raise HTTPException(status_code=500, detail=str(exc)) from exc
@app.post("/api/recall")
def recall(request: RecallRequest):
result_path = _result_path(request.doc_id)
if not result_path.exists():
raise HTTPException(status_code=404, detail="Chunk result not found; upload and split the document again")
started = time.perf_counter()
chunk_result = ChunkResult.model_validate_json(result_path.read_text(encoding="utf-8"))
results, candidate_count = recall_chunks(request.query, chunk_result.chunks, request.top_k)
return {
"doc_id": request.doc_id,
"query": request.query,
"top_k": request.top_k,
"candidate_count": candidate_count,
"returned_count": len(results),
"elapsed_ms": round((time.perf_counter() - started) * 1000, 3),
"results": results,
}
if FRONTEND_DIR.exists():
app.mount("/", StaticFiles(directory=FRONTEND_DIR, html=True), name="frontend")
+6
View File
@@ -0,0 +1,6 @@
"""RAG-cut document chunking pipeline."""
from rag_cut.models import Block, Chunk, SplitConfig, SplitMode
from rag_cut.pipeline import chunk_document
__all__ = ["Block", "Chunk", "SplitConfig", "SplitMode", "chunk_document"]
+244
View File
@@ -0,0 +1,244 @@
"""Layout metadata: reading order, heading binding, image context."""
from __future__ import annotations
from rag_cut.models import Block, BlockType
from rag_cut.parsers.pdf.tables import build_table_embedding_text, extract_table_keywords, guess_table_title
from rag_cut.splitters.heading_splitter import is_heading_block, normalize_heading_block
def assign_order_index(blocks: list[Block], start: int = 0) -> list[Block]:
"""Assign a global document-order index to every block."""
result: list[Block] = []
for idx, block in enumerate(blocks):
meta = dict(block.meta)
meta["order_index"] = start + idx
result.append(block.model_copy(update={"meta": meta}))
return result
def sort_blocks_reading_order(blocks: list[Block]) -> list[Block]:
"""Sort by page then top-to-bottom / left-to-right; stable for missing bboxes."""
def sort_key(item: tuple[int, Block]) -> tuple[int, float, float, int]:
idx, block = item
page = block.meta.get("page")
try:
page_key = int(page) if page is not None else 10**9
except (TypeError, ValueError):
page_key = 10**9
bbox = block.meta.get("bbox")
if isinstance(bbox, (list, tuple)) and len(bbox) >= 4:
try:
return (page_key, float(bbox[1]), float(bbox[0]), idx)
except (TypeError, ValueError):
pass
return (page_key, float(idx), 0.0, idx)
return [block for _, block in sorted(enumerate(blocks), key=sort_key)]
def bind_heading_context(blocks: list[Block]) -> list[Block]:
"""
Propagate heading hierarchy and bind images to nearest heading and adjacent text.
Chunk content should follow: heading → body → image → OCR/caption → subsequent body.
"""
heading_stack: list[tuple[int, str]] = []
result: list[Block] = []
for i, block in enumerate(blocks):
meta = dict(block.meta)
candidate = normalize_heading_block(block)
if is_heading_block(candidate):
level = candidate.level or 1
while heading_stack and heading_stack[-1][0] >= level:
heading_stack.pop()
heading_stack.append((level, candidate.text))
meta["section_boundary"] = True
if block.type != BlockType.HEADING or candidate is not block:
block = candidate.model_copy(update={"meta": {**dict(candidate.meta), **meta}})
meta = dict(block.meta)
else:
meta["section_boundary"] = False
# Demote overlong HEADING blobs (title+body merge) back to paragraph.
if block.type == BlockType.HEADING and candidate.type != BlockType.HEADING:
block = candidate.model_copy(update={"meta": {**dict(candidate.meta), **meta}})
meta = dict(block.meta)
if heading_stack:
meta["parent_heading"] = heading_stack[-1][1]
meta["heading_path"] = [text for _, text in heading_stack]
meta["nearest_heading"] = heading_stack[-1][1]
if (block.level or 1) <= 2 and block.type == BlockType.HEADING:
meta["chapter"] = block.text
elif "chapter" not in meta and len(heading_stack) >= 1:
# Keep chapter as nearest level-1/2 ancestor
for lvl, text in reversed(heading_stack):
if lvl <= 2:
meta["chapter"] = text
break
if block.type == BlockType.IMAGE:
_bind_image_context(blocks, i, meta)
elif block.type == BlockType.TABLE:
_bind_table_context(blocks, i, meta)
result.append(block.model_copy(update={"meta": meta}))
return result
def _heading_text(block: Block) -> str | None:
candidate = normalize_heading_block(block)
if is_heading_block(candidate):
return (candidate.text or "").strip() or None
return None
def _spatial_heading_above(
blocks: list[Block],
index: int,
page: object,
bbox: list[float] | None,
) -> str | None:
"""Pick same-page heading whose bottom edge is nearest above the image top."""
if page is None or not bbox or len(bbox) < 4:
return None
try:
page_key = int(page)
img_y0 = float(bbox[1])
except (TypeError, ValueError):
return None
best_text: str | None = None
best_dist = float("inf")
for j, block in enumerate(blocks):
if j == index:
continue
text = _heading_text(block)
if not text:
continue
try:
if int(block.meta.get("page")) != page_key:
continue
except (TypeError, ValueError):
continue
hb = block.meta.get("bbox")
if not isinstance(hb, (list, tuple)) or len(hb) < 4:
continue
try:
heading_y1 = float(hb[3])
except (TypeError, ValueError):
continue
if heading_y1 > img_y0 + 2:
continue
dist = img_y0 - heading_y1
if dist < best_dist:
best_dist = dist
best_text = text
return best_text
def _list_heading_above(blocks: list[Block], index: int) -> str | None:
for j in range(index - 1, -1, -1):
text = _heading_text(blocks[j])
if text:
return text
return None
def _bind_image_context(blocks: list[Block], index: int, meta: dict) -> None:
"""Bind image to nearest heading above and adjacent body text on the same page."""
page = meta.get("page")
bbox = meta.get("bbox") if isinstance(meta.get("bbox"), list) else None
bound = _spatial_heading_above(blocks, index, page, bbox) or _list_heading_above(blocks, index)
if bound:
meta["bound_heading"] = bound
for j in range(index - 1, -1, -1):
prev = blocks[j]
if _heading_text(prev):
break
if prev.type == BlockType.PARAGRAPH and prev.meta.get("page") == page:
meta["preceding_text"] = (prev.text or "")[:300]
break
if not meta.get("bound_heading") and meta.get("nearest_heading"):
meta["bound_heading"] = meta["nearest_heading"]
for j in range(index + 1, len(blocks)):
nxt = blocks[j]
if nxt.type == BlockType.IMAGE:
break
if nxt.type == BlockType.PARAGRAPH and nxt.meta.get("page") == page:
meta["following_text"] = (nxt.text or "")[:300]
break
if _heading_text(nxt):
break
def _bind_table_context(blocks: list[Block], index: int, meta: dict) -> None:
"""Bind table title, surrounding text and retrieval fields."""
page = meta.get("page")
block = blocks[index]
for j in range(index - 1, -1, -1):
prev = blocks[j]
if prev.type == BlockType.TABLE:
break
heading = _heading_text(prev)
if heading:
if not meta.get("table_title"):
meta["table_title"] = heading
meta.setdefault("bound_heading", heading)
break
if prev.type == BlockType.PARAGRAPH and prev.meta.get("page") == page:
title = guess_table_title(prev.text or "")
if title:
meta["table_title"] = title
meta["preceding_text"] = (prev.text or "")[:400]
break
if not meta.get("table_title") and meta.get("nearest_heading"):
meta.setdefault("table_title", meta["nearest_heading"])
for j in range(index + 1, len(blocks)):
nxt = blocks[j]
if nxt.type == BlockType.TABLE:
break
if nxt.type == BlockType.PARAGRAPH and nxt.meta.get("page") == page:
text = (nxt.text or "").strip()
if text and len(text) <= 300:
meta["following_text"] = text
if not meta.get("footnotes") and any(k in text for k in ("注", "备注", "说明", "Note")):
meta["table_description"] = text
break
if _heading_text(nxt):
break
if not meta.get("keywords"):
meta["keywords"] = extract_table_keywords(
meta.get("table_title") or "",
block.markdown or "",
meta.get("table_description") or meta.get("footnotes") or "",
block.ocr_text or "",
)
meta["embedding_text"] = build_table_embedding_text(
chapter=meta.get("chapter") or meta.get("nearest_heading") or "",
table_title=meta.get("table_title") or "",
markdown=block.markdown or "",
description=meta.get("table_description") or meta.get("preceding_text") or "",
footnotes=meta.get("footnotes") or meta.get("following_text") or "",
keywords=meta.get("keywords") or [],
ocr_text=block.ocr_text or "",
)
def enrich_layout_metadata(blocks: list[Block]) -> list[Block]:
"""Full post-parse enrichment: spatial order + order index + heading/image context."""
blocks = sort_blocks_reading_order(blocks)
blocks = assign_order_index(blocks)
return bind_heading_context(blocks)
+100
View File
@@ -0,0 +1,100 @@
"""Core data models for document chunking."""
from __future__ import annotations
from enum import Enum
from typing import Any
from pydantic import BaseModel, Field
class BlockType(str, Enum):
HEADING = "heading"
PARAGRAPH = "paragraph"
TABLE = "table"
IMAGE = "image"
LIST = "list"
CODE = "code"
class SplitMode(str, Enum):
DEFAULT = "default"
DELIMITER = "delimiter"
PARENT_CHILD = "parent_child"
BY_ROW = "by_row"
class Block(BaseModel):
"""Atomic document unit; images and tables must not be split across chunks."""
type: BlockType
text: str = ""
markdown: str = ""
level: int = 0
image_id: str | None = None
image_path: str | None = None
ocr_text: str = ""
meta: dict[str, Any] = Field(default_factory=dict)
def render(self) -> str:
if self.type == BlockType.TABLE:
parts: list[str] = []
title = self.meta.get("table_title")
if title:
parts.append(f"### {title}")
image_ref = self.image_path or self.meta.get("crop_path")
if image_ref:
alt = (title or self.image_id or "table").replace("\n", " ").strip()
parts.append(f"![{alt}]({image_ref})")
parts.append(self.markdown or self.text)
footnotes = self.meta.get("footnotes")
if footnotes and footnotes not in (self.markdown or ""):
parts.append(f"*{footnotes}*")
if self.ocr_text:
parts.append(f"*[表格 OCR]* {self.ocr_text}")
return "\n\n".join(p for p in parts if p.strip())
if self.type == BlockType.IMAGE:
alt = (self.image_id or self.text or "image").replace("\n", " ").strip()
path = self.image_path or self.image_id or ""
body = f"![{alt}]({path})"
if self.ocr_text:
body += f"\n\n*[OCR]* {self.ocr_text}"
return body
if self.type == BlockType.HEADING:
prefix = "#" * max(1, min(self.level, 6))
return f"{prefix} {self.text}".strip()
return self.text or self.markdown
class SplitConfig(BaseModel):
mode: SplitMode = SplitMode.DEFAULT
delimiter: str | None = None
parent_delimiter: str | None = None
child_delimiter: str | None = None
max_chunk_size: int = 1500
child_max_size: int = 512
overlap: int = 150
header_row_start: int = 1
header_row_end: int = 1
start_row: int = 2
rows_per_chunk: int = 1
table_format: str = "markdown"
class Chunk(BaseModel):
index: int
content: str
char_count: int
block_types: list[str]
meta: dict[str, Any] = Field(default_factory=dict)
class ChunkResult(BaseModel):
filename: str
doc_id: str
split_mode: SplitMode
split_config: dict[str, Any] = Field(default_factory=dict)
block_count: int
chunk_count: int
chunks: list[Chunk]
assets_dir: str | None = None
+5
View File
@@ -0,0 +1,5 @@
"""Document parsers."""
from rag_cut.parsers.registry import get_parser
__all__ = ["get_parser"]
+14
View File
@@ -0,0 +1,14 @@
"""Parser base class."""
from __future__ import annotations
from abc import ABC, abstractmethod
from pathlib import Path
from rag_cut.models import Block
class BaseParser(ABC):
@abstractmethod
def parse(self, path: Path, assets_dir: Path) -> list[Block]:
raise NotImplementedError
+24
View File
@@ -0,0 +1,24 @@
"""Standalone image parser."""
from __future__ import annotations
import shutil
from pathlib import Path
from rag_cut.models import Block, BlockType
from rag_cut.parsers.base import BaseParser
class ImageParser(BaseParser):
def parse(self, path: Path, assets_dir: Path) -> list[Block]:
assets_dir.mkdir(parents=True, exist_ok=True)
target = assets_dir / path.name
shutil.copy2(path, target)
return [
Block(
type=BlockType.IMAGE,
text=path.stem,
image_id=path.name,
image_path=str(target),
)
]
+421
View File
@@ -0,0 +1,421 @@
"""Optional MinerU parser adapter.
MinerU improves the parsing layer when installed, while RAG-cut keeps owning
chunking strategy and retrieval metadata.
"""
from __future__ import annotations
import json
import os
import signal
import shutil
import subprocess
from pathlib import Path
from typing import Any
from rag_cut.models import Block, BlockType
from rag_cut.parsers.pdf.ocr import describe_visual, ocr_image_file
def _mineru_command() -> str | None:
configured = os.getenv("RAG_CUT_MINERU_CMD")
if configured:
return configured
return shutil.which("mineru") or shutil.which("magic-pdf")
def _mineru_timeout_sec() -> float:
raw = os.getenv("RAG_CUT_MINERU_TIMEOUT", "540").strip()
try:
return max(30.0, float(raw))
except ValueError:
return 540.0
def _mineru_ocr_enabled() -> bool:
"""Fill empty OCR from image assets when pytesseract is available."""
raw = os.getenv("RAG_CUT_MINERU_OCR", "1").strip().lower()
return raw not in {"0", "false", "off", "no"}
def _terminate_process_tree(process: subprocess.Popen[str]) -> None:
"""Stop MinerU and any temporary API/model workers it started."""
try:
if os.name == "nt":
subprocess.run(
["taskkill", "/PID", str(process.pid), "/T", "/F"],
check=False,
capture_output=True,
timeout=10,
)
else:
os.killpg(process.pid, signal.SIGKILL)
except (OSError, subprocess.SubprocessError):
pass
if process.poll() is None:
process.kill()
try:
process.wait(timeout=5)
except (OSError, subprocess.SubprocessError):
pass
def _run_mineru(path: Path, output_dir: Path) -> bool:
cmd = _mineru_command()
if not cmd:
return False
output_dir.mkdir(parents=True, exist_ok=True)
if Path(cmd).name.lower() == "magic-pdf":
args = [cmd, "-p", str(path), "-o", str(output_dir)]
else:
args = [cmd, "-p", str(path), "-o", str(output_dir), "-b", "pipeline"]
api_url = os.getenv("RAG_CUT_MINERU_API_URL", "").strip()
if api_url:
args.extend(["--api-url", api_url])
try:
popen_kwargs: dict[str, Any] = {}
if os.name == "nt":
popen_kwargs["creationflags"] = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)
else:
popen_kwargs["start_new_session"] = True
process = subprocess.Popen(
args,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
encoding="utf-8",
errors="replace",
**popen_kwargs,
)
process.communicate(timeout=_mineru_timeout_sec())
except subprocess.TimeoutExpired:
_terminate_process_tree(process)
return False
except (OSError, subprocess.SubprocessError):
return False
return process.returncode == 0
def _find_content_list(output_dir: Path) -> Path | None:
"""Prefer stable content_list.json over content_list_v2.json."""
candidates = list(output_dir.rglob("*content_list*.json"))
if not candidates:
return None
def rank(path: Path) -> tuple[int, int, str]:
name = path.name.lower()
is_v2 = 1 if "v2" in name else 0
return (is_v2, len(path.parts), str(path).lower())
return sorted(candidates, key=rank)[0]
def _read_json(path: Path) -> Any:
return json.loads(path.read_text(encoding="utf-8"))
def _resolve_asset(path_text: str, base_dir: Path) -> Path | None:
if not path_text:
return None
raw = Path(path_text)
if raw.is_absolute() and raw.exists():
return raw
candidate = base_dir / raw
if candidate.exists():
return candidate
matches = list(base_dir.rglob(raw.name))
return matches[0] if matches else None
def _copy_asset(src: Path | None, assets_dir: Path) -> tuple[str | None, str | None]:
if not src or not src.exists():
return None, None
assets_dir.mkdir(parents=True, exist_ok=True)
target = assets_dir / src.name
if src.resolve() != target.resolve():
shutil.copy2(src, target)
return target.name, str(target)
def _page(item: dict[str, Any]) -> int | None:
for key in ("page", "page_no", "page_num"):
if item.get(key) is not None:
try:
return int(item[key])
except (TypeError, ValueError):
return None
if item.get("page_idx") is not None:
try:
return int(item["page_idx"]) + 1
except (TypeError, ValueError):
return None
return None
def _bbox(item: dict[str, Any]) -> list[float]:
raw = item.get("bbox") or item.get("poly") or []
if isinstance(raw, list) and len(raw) >= 4:
try:
if all(isinstance(v, (int, float)) for v in raw[:4]):
return [float(v) for v in raw[:4]]
if all(isinstance(p, list) and len(p) >= 2 for p in raw):
xs = [float(p[0]) for p in raw]
ys = [float(p[1]) for p in raw]
return [min(xs), min(ys), max(xs), max(ys)]
except (TypeError, ValueError):
return []
return []
def _meta(item: dict[str, Any]) -> dict[str, Any]:
meta: dict[str, Any] = {"parser": "mineru"}
page = _page(item)
bbox = _bbox(item)
kind = _item_type(item)
if page:
meta["page"] = page
if bbox:
meta["bbox"] = bbox
if kind:
meta["mineru_type"] = kind
return meta
def _text_value(item: dict[str, Any]) -> str:
for key in (
"text",
"content",
"table_caption",
"image_caption",
"code_body",
"code",
"equation",
"latex",
):
value = item.get(key)
if isinstance(value, str) and value.strip():
return value.strip()
if isinstance(value, list):
joined = " ".join(str(v).strip() for v in value if str(v).strip())
if joined:
return joined
list_items = item.get("list_items")
if isinstance(list_items, list) and list_items:
parts: list[str] = []
for entry in list_items:
if isinstance(entry, str) and entry.strip():
parts.append(entry.strip())
elif isinstance(entry, dict):
piece = str(entry.get("text") or entry.get("content") or "").strip()
if piece:
parts.append(piece)
if parts:
return "\n".join(parts)
return ""
def _table_markdown(item: dict[str, Any]) -> str:
for key in ("table_body", "html", "text", "content"):
value = item.get(key)
if isinstance(value, str) and value.strip():
return value.strip()
return ""
def _image_path_text(item: dict[str, Any]) -> str:
for key in ("img_path", "image_path", "path"):
value = item.get(key)
if isinstance(value, str) and value.strip():
return value.strip()
return ""
def _item_type(item: dict[str, Any]) -> str:
return str(item.get("type") or item.get("category") or "").lower()
def _ocr_from_item(item: dict[str, Any], image_path: str | None) -> str:
ocr_text = str(item.get("ocr_text") or item.get("image_ocr") or item.get("img_caption") or "").strip()
if not ocr_text:
caption = item.get("image_caption")
if isinstance(caption, list):
ocr_text = " ".join(str(v).strip() for v in caption if str(v).strip())
elif isinstance(caption, str):
ocr_text = caption.strip()
if ocr_text or not image_path or not _mineru_ocr_enabled():
return ocr_text
return ocr_image_file(Path(image_path))
_SKIP_MINERU_TYPES = frozenset(
{
"header",
"page_header",
"footer",
"page_footer",
"page_number",
"page_num",
"header_image",
"footer_image",
"aside_text",
"toc",
"contents",
"table_of_contents",
}
)
# Visual regions MinerU may label separately from plain "image".
_IMAGE_KINDS = frozenset(
{
"image",
"figure",
"chart",
"diagram",
"graphic",
"photo",
"screenshot",
"equation",
"formula",
}
)
def _build_image_block(
item: dict[str, Any],
base_dir: Path,
assets_dir: Path,
*,
visual_kind: str = "图片",
) -> Block | None:
src = _resolve_asset(_image_path_text(item), base_dir)
image_id, image_path = _copy_asset(src, assets_dir)
if not image_id and not image_path:
return None
caption = _text_value(item)
ocr_text = _ocr_from_item(item, image_path)
meta = _meta(item)
meta.update(
{
"caption": caption,
"source_image_path": str(src) if src else None,
}
)
return Block(
type=BlockType.IMAGE,
text=caption or describe_visual(ocr_text, visual_kind, image_id or "image"),
image_id=image_id,
image_path=image_path,
ocr_text=ocr_text,
meta=meta,
)
def _blocks_from_content_list(content_list: list[dict[str, Any]], json_path: Path, assets_dir: Path) -> list[Block]:
blocks: list[Block] = []
base_dir = json_path.parent
for item in content_list:
kind = _item_type(item)
if kind in _SKIP_MINERU_TYPES:
continue
meta = _meta(item)
img_path = _image_path_text(item)
if kind != "table" and (
kind in _IMAGE_KINDS or (img_path and kind not in {"text", "title", "heading", "list"})
):
visual = "图表" if kind in {"chart", "diagram"} else "图片"
image_block = _build_image_block(item, base_dir, assets_dir, visual_kind=visual)
if image_block:
blocks.append(image_block)
continue
# Chart/image without a resolvable asset: fall through if there is caption text.
if kind == "table":
markdown = _table_markdown(item)
src = _resolve_asset(img_path, base_dir) if img_path else None
image_id, image_path = _copy_asset(src, assets_dir)
ocr_text = _ocr_from_item(item, image_path) if image_path else ""
if not markdown and not image_path:
continue
caption = _text_value(item)
if caption:
meta["caption"] = caption
if src:
meta["source_image_path"] = str(src)
if image_path:
meta["crop_path"] = image_path
blocks.append(
Block(
type=BlockType.TABLE,
markdown=markdown,
text=caption,
image_id=image_id,
image_path=image_path,
ocr_text=ocr_text,
meta=meta,
)
)
continue
if kind == "code":
image_block = _build_image_block(item, base_dir, assets_dir, visual_kind="代码")
if image_block:
blocks.append(image_block)
continue
code_text = _text_value(item)
if code_text:
blocks.append(Block(type=BlockType.PARAGRAPH, text=code_text, meta=meta))
continue
if kind in {"page_footnote", "footnote", "ref_text"}:
note = _text_value(item)
if note:
meta["is_footnote"] = True
blocks.append(Block(type=BlockType.PARAGRAPH, text=note, meta=meta))
continue
text = _text_value(item)
if not text:
# Last resort: unknown typed asset with an image should not be dropped.
if img_path:
image_block = _build_image_block(item, base_dir, assets_dir)
if image_block:
blocks.append(image_block)
continue
level = item.get("text_level") or item.get("level")
try:
level_int = int(level)
except (TypeError, ValueError):
level_int = 0
if kind in {"title", "heading"} or level_int > 0:
meta["source_heading"] = True
meta["source_heading_level"] = max(level_int, 1)
blocks.append(Block(type=BlockType.HEADING, text=text, level=max(level_int, 1), meta=meta))
else:
blocks.append(Block(type=BlockType.PARAGRAPH, text=text, meta=meta))
return blocks
def parse_pdf_with_mineru(path: Path, assets_dir: Path) -> list[Block] | None:
"""Return MinerU blocks when available; otherwise None for fallback."""
output_dir = assets_dir / "_mineru"
if not _run_mineru(path, output_dir):
return None
content_json = _find_content_list(output_dir)
if not content_json:
return None
data = _read_json(content_json)
if not isinstance(data, list):
return None
blocks = _blocks_from_content_list(data, content_json, assets_dir)
return blocks or None
+164
View File
@@ -0,0 +1,164 @@
"""Convert Office documents to PDF for downstream parsing."""
from __future__ import annotations
import platform
import shutil
import subprocess
from pathlib import Path
_PPT_EXTENSIONS = {".ppt", ".pptx", ".ppsx"}
_WORD_EXTENSIONS = {".doc", ".docx"}
def _find_soffice() -> Path | None:
for candidate in (
"soffice",
r"C:\Program Files\LibreOffice\program\soffice.exe",
r"C:\Program Files (x86)\LibreOffice\program\soffice.exe",
"/usr/bin/libreoffice",
"/usr/bin/soffice",
):
if candidate in ("soffice", "libreoffice"):
found = shutil.which(candidate)
if found:
return Path(found)
elif Path(candidate).is_file():
return Path(candidate)
return None
def _convert_via_libreoffice(source: Path, out_dir: Path) -> Path:
soffice = _find_soffice()
if soffice is None:
raise RuntimeError("LibreOffice (soffice) not found")
out_dir.mkdir(parents=True, exist_ok=True)
subprocess.run(
[
str(soffice),
"--headless",
"--convert-to",
"pdf",
"--outdir",
str(out_dir),
str(source.resolve()),
],
check=True,
capture_output=True,
text=True,
timeout=180,
)
pdf_path = out_dir / f"{source.stem}.pdf"
if not pdf_path.is_file():
raise RuntimeError(f"LibreOffice did not produce PDF: {pdf_path}")
return pdf_path
def _convert_via_powerpoint(source: Path, out_dir: Path) -> Path:
if platform.system() != "Windows":
raise RuntimeError("PowerPoint COM conversion is only available on Windows")
out_dir.mkdir(parents=True, exist_ok=True)
pdf_path = out_dir / f"{source.stem}.pdf"
src = str(source.resolve())
dst = str(pdf_path.resolve())
script = f"""
$ErrorActionPreference = 'Stop'
$pp = New-Object -ComObject PowerPoint.Application
try {{
$pres = $pp.Presentations.Open('{src.replace("'", "''")}', $true, $true, $false)
try {{
$pres.SaveAs('{dst.replace("'", "''")}', 32)
}} finally {{
$pres.Close()
}}
}} finally {{
$pp.Quit()
}}
"""
subprocess.run(
["powershell", "-NoProfile", "-NonInteractive", "-Command", script],
check=True,
capture_output=True,
text=True,
timeout=180,
)
if not pdf_path.is_file():
raise RuntimeError(f"PowerPoint did not produce PDF: {pdf_path}")
return pdf_path
def _convert_via_word(source: Path, out_dir: Path) -> Path:
if platform.system() != "Windows":
raise RuntimeError("Word COM conversion is only available on Windows")
out_dir.mkdir(parents=True, exist_ok=True)
pdf_path = out_dir / f"{source.stem}.pdf"
src = str(source.resolve())
dst = str(pdf_path.resolve())
script = f"""
$ErrorActionPreference = 'Stop'
$word = New-Object -ComObject Word.Application
$word.Visible = $false
try {{
$doc = $word.Documents.Open('{src.replace("'", "''")}', $false, $true)
try {{
$doc.SaveAs2('{dst.replace("'", "''")}', 17)
}} finally {{
$doc.Close([ref]$false)
}}
}} finally {{
$word.Quit()
}}
"""
subprocess.run(
["powershell", "-NoProfile", "-NonInteractive", "-Command", script],
check=True,
capture_output=True,
text=True,
timeout=180,
)
if not pdf_path.is_file():
raise RuntimeError(f"Word did not produce PDF: {pdf_path}")
return pdf_path
def convert_to_pdf(source: Path, out_dir: Path) -> Path:
"""Convert an Office file to PDF. Tries LibreOffice, then format-specific COM on Windows."""
source = source.resolve()
if not source.is_file():
raise FileNotFoundError(source)
errors: list[str] = []
try:
return _convert_via_libreoffice(source, out_dir)
except Exception as exc: # noqa: BLE001 — collect and try next backend
errors.append(f"libreoffice: {exc}")
ext = source.suffix.lower()
fallbacks: list = []
if ext in _PPT_EXTENSIONS:
fallbacks.append(_convert_via_powerpoint)
elif ext in _WORD_EXTENSIONS:
fallbacks.append(_convert_via_word)
for converter in fallbacks:
try:
return converter(source, out_dir)
except Exception as exc: # noqa: BLE001
errors.append(f"{converter.__name__}: {exc}")
detail = "; ".join(errors) if errors else "no converter available"
raise RuntimeError(
f"Failed to convert {source.name} to PDF. "
f"Install LibreOffice or Microsoft Office. Details: {detail}"
)
def is_presentation(path: Path) -> bool:
return path.suffix.lower() in _PPT_EXTENSIONS
def is_word_document(path: Path) -> bool:
return path.suffix.lower() in _WORD_EXTENSIONS
+5
View File
@@ -0,0 +1,5 @@
"""PyMuPDF-based PDF parsing pipeline."""
from rag_cut.parsers.pdf.pipeline import parse_pdf
__all__ = ["parse_pdf"]
+349
View File
@@ -0,0 +1,349 @@
"""Page rendering and layout analysis."""
from __future__ import annotations
import re
from dataclasses import dataclass, field
from typing import Any
import fitz
from rag_cut.parsers.pdf.tables import looks_like_table
_NUMBERED_LINE_RE = re.compile(
r"^\s*(\d+(?:\.\d+)*)(?:\.|.)?\s*([A-Za-z0-9\u4e00-\u9fff].{2,})$"
)
@dataclass
class TextLine:
x0: float
y0: float
x1: float
y1: float
text: str
font_size: float
@dataclass
class LayoutRegion:
x0: float
y0: float
x1: float
y1: float
kind: str # text | table | image
data: dict[str, Any] = field(default_factory=dict)
@dataclass
class PageLayout:
page_index: int
page_width: float
page_height: float
body_font_size: float
regions: list[LayoutRegion] = field(default_factory=list)
def render_page(page: fitz.Page, zoom: float = 2.0) -> fitz.Pixmap:
"""Render page to pixmap for region cropping and OCR."""
return page.get_pixmap(matrix=fitz.Matrix(zoom, zoom), alpha=False)
def _median_body_size(text_dict: dict) -> float:
sizes: list[float] = []
for block in text_dict.get("blocks", []):
if block.get("type") != 0:
continue
for line in block.get("lines", []):
for span in line.get("spans", []):
sizes.append(span.get("size", 12))
return sorted(sizes)[len(sizes) // 2] if sizes else 12
def _collect_text_lines(page: fitz.Page, body_size: float) -> list[TextLine]:
text_dict = page.get_text("dict")
lines: list[TextLine] = []
for block in text_dict.get("blocks", []):
if block.get("type") != 0:
continue
for line in block.get("lines", []):
text = "".join(span.get("text", "") for span in line.get("spans", [])).strip()
if not text:
continue
x0, y0, x1, y1 = line["bbox"]
max_size = body_size
for span in line.get("spans", []):
max_size = max(max_size, span.get("size", body_size))
lines.append(TextLine(x0=x0, y0=y0, x1=x1, y1=y1, text=text, font_size=max_size))
return lines
def _detect_columns(
items: list[tuple[float, float, float, float]],
page_width: float,
) -> list[tuple[float, float, float, float]]:
if not items:
return items
mid = page_width / 2
left = [b for b in items if (b[0] + b[2]) / 2 < mid]
right = [b for b in items if (b[0] + b[2]) / 2 >= mid]
if len(left) >= 2 and len(right) >= 2:
left.sort(key=lambda b: (b[1], b[0]))
right.sort(key=lambda b: (b[1], b[0]))
return left + right
return sorted(items, key=lambda b: (b[1], b[0]))
def _sort_reading_order(
items: list[tuple[float, float, float, float]],
page_width: float,
) -> list[tuple[float, float, float, float]]:
"""Sort page boxes in a human reading order, preserving two-column flows."""
if not items:
return []
mid = page_width / 2
full_width: list[tuple[float, float, float, float]] = []
column_items: list[tuple[float, float, float, float]] = []
for box in items:
width = box[2] - box[0]
spans_mid = box[0] < mid < box[2]
if width >= page_width * 0.60 or (spans_mid and width >= page_width * 0.35):
full_width.append(box)
else:
column_items.append(box)
left = [b for b in column_items if (b[0] + b[2]) / 2 < mid]
right = [b for b in column_items if (b[0] + b[2]) / 2 >= mid]
if len(left) < 2 or len(right) < 2:
return sorted(items, key=lambda b: (b[1], b[0]))
full_width.sort(key=lambda b: (b[1], b[0]))
ordered: list[tuple[float, float, float, float]] = []
segment_top = float("-inf")
def add_columns_between(top: float, bottom: float) -> None:
segment = [b for b in column_items if b[1] >= top and b[1] < bottom]
segment_left = sorted([b for b in segment if (b[0] + b[2]) / 2 < mid], key=lambda b: (b[1], b[0]))
segment_right = sorted([b for b in segment if (b[0] + b[2]) / 2 >= mid], key=lambda b: (b[1], b[0]))
ordered.extend(segment_left)
ordered.extend(segment_right)
for box in full_width:
add_columns_between(segment_top, box[1])
ordered.append(box)
segment_top = box[3]
add_columns_between(segment_top, float("inf"))
seen: set[tuple[float, float, float, float]] = set(ordered)
ordered.extend(b for b in sorted(items, key=lambda b: (b[1], b[0])) if b not in seen)
return ordered
def _rect_overlap(a: tuple[float, float, float, float], b: tuple[float, float, float, float]) -> float:
x0 = max(a[0], b[0])
y0 = max(a[1], b[1])
x1 = min(a[2], b[2])
y1 = min(a[3], b[3])
if x1 <= x0 or y1 <= y0:
return 0.0
inter = (x1 - x0) * (y1 - y0)
area_a = max((a[2] - a[0]) * (a[3] - a[1]), 1e-6)
return inter / area_a
def _center_inside(inner: tuple[float, float, float, float], outer: tuple[float, float, float, float]) -> bool:
cx = (inner[0] + inner[2]) / 2
cy = (inner[1] + inner[3]) / 2
return outer[0] <= cx <= outer[2] and outer[1] <= cy <= outer[3]
def _horizontal_overlap_ratio(a: TextLine, b: TextLine) -> float:
overlap = min(a.x1, b.x1) - max(a.x0, b.x0)
if overlap <= 0:
return 0.0
return overlap / max(min(a.x1 - a.x0, b.x1 - b.x0), 1e-6)
def _line_on_image(
line: TextLine,
image_boxes: list[tuple[float, float, float, float]],
large_image_boxes: list[tuple[float, float, float, float]],
) -> bool:
# Never drop numbered section titles even if they sit near a screenshot.
if _NUMBERED_LINE_RE.match(line.text.strip()):
return False
bbox = (line.x0, line.y0, line.x1, line.y1)
# Large background/screenshot: require near-full coverage of the line, not
# merely center-inside (which wiped text sitting in margins of wide figures).
if large_image_boxes and any(
_center_inside(bbox, box) and _rect_overlap(bbox, box) >= 0.85 for box in large_image_boxes
):
return True
return any(_rect_overlap(bbox, box) >= 0.55 for box in image_boxes)
def _area(box: tuple[float, float, float, float]) -> float:
return max(box[2] - box[0], 0.0) * max(box[3] - box[1], 0.0)
def _filter_nested_image_regions(regions: list[LayoutRegion]) -> list[LayoutRegion]:
"""Drop image fragments that are already contained in a larger screenshot."""
result: list[LayoutRegion] = []
boxes = [(r.x0, r.y0, r.x1, r.y1) for r in regions]
for region, box in zip(regions, boxes):
box_area = _area(box)
nested = False
for other in boxes:
other_area = _area(other)
if other == box or other_area <= box_area * 1.5:
continue
if _center_inside(box, other) and _rect_overlap(box, other) >= 0.85:
nested = True
break
if not nested:
result.append(region)
return result
def _merge_lines_to_paragraphs(lines: list[TextLine], page_width: float) -> list[TextLine]:
if not lines:
return []
ordered_boxes = _sort_reading_order([(ln.x0, ln.y0, ln.x1, ln.y1) for ln in lines], page_width)
order = {(b[0], b[1], b[2], b[3]): i for i, b in enumerate(ordered_boxes)}
ordered = sorted(lines, key=lambda ln: order.get((ln.x0, ln.y0, ln.x1, ln.y1), (ln.y0, ln.x0)))
paragraphs: list[TextLine] = []
current = ordered[0]
for nxt in ordered[1:]:
vgap = nxt.y0 - current.y1
line_h = max(current.y1 - current.y0, nxt.y1 - nxt.y0, 8.0)
size_gap = abs(current.font_size - nxt.font_size)
size_ratio = max(current.font_size, nxt.font_size) / max(
min(current.font_size, nxt.font_size), 1e-6
)
# Keep title vs body separate: small absolute gap is enough when ratio is large.
same_style = size_gap <= 1.0 and size_ratio <= 1.15
either_numbered = bool(
_NUMBERED_LINE_RE.match(current.text.strip()) or _NUMBERED_LINE_RE.match(nxt.text.strip())
)
if (
same_style
and not either_numbered
and vgap <= line_h * 2.2
and _horizontal_overlap_ratio(current, nxt) >= 0.25
):
joiner = "" if current.text.endswith("-") or current.text.endswith(" ") else " "
current = TextLine(
x0=min(current.x0, nxt.x0),
y0=current.y0,
x1=max(current.x1, nxt.x1),
y1=nxt.y1,
text=f"{current.text}{joiner}{nxt.text}",
font_size=max(current.font_size, nxt.font_size),
)
else:
paragraphs.append(current)
current = nxt
paragraphs.append(current)
return paragraphs
def analyze_page_layout(page: fitz.Page, page_index: int) -> PageLayout:
"""Layout analysis: text paragraphs, table regions, image regions."""
text_dict = page.get_text("dict")
body_size = _median_body_size(text_dict)
page_width = page.rect.width
page_height = page.rect.height
page_area = page_width * page_height
table_regions: list[LayoutRegion] = []
try:
table_finder = page.find_tables()
for idx, table in enumerate(table_finder.tables):
bbox = table.bbox
rows = table.extract() or []
if not looks_like_table(rows):
continue
table_regions.append(
LayoutRegion(
x0=bbox[0],
y0=bbox[1],
x1=bbox[2],
y1=bbox[3],
kind="table",
data={"rows": rows, "table_index": idx, "source": "pymupdf"},
)
)
except Exception:
pass
image_regions: list[LayoutRegion] = []
for img_info in page.get_images(full=True):
xref = img_info[0]
try:
rects = page.get_image_rects(xref)
except Exception:
rects = []
if not rects:
continue
for rect_idx, rect in enumerate(rects):
image_regions.append(
LayoutRegion(
x0=rect.x0,
y0=rect.y0,
x1=rect.x1,
y1=rect.y1,
kind="image",
data={"xref": xref, "rect_index": rect_idx},
)
)
image_regions = _filter_nested_image_regions(image_regions)
image_boxes = [(r.x0, r.y0, r.x1, r.y1) for r in image_regions]
large_image_boxes = [
box
for box in image_boxes
if (box[2] - box[0]) * (box[3] - box[1]) >= page_area * 0.12
]
table_boxes = [(r.x0, r.y0, r.x1, r.y1) for r in table_regions]
raw_lines = _collect_text_lines(page, body_size)
filtered_lines = [
ln
for ln in raw_lines
if not _line_on_image(ln, image_boxes, large_image_boxes)
and not any(_rect_overlap((ln.x0, ln.y0, ln.x1, ln.y1), box) >= 0.55 for box in table_boxes)
]
paragraphs = _merge_lines_to_paragraphs(filtered_lines, page_width)
text_regions: list[LayoutRegion] = []
for para in paragraphs:
text_regions.append(
LayoutRegion(
x0=para.x0,
y0=para.y0,
x1=para.x1,
y1=para.y1,
kind="text",
data={"text": para.text, "font_size": para.font_size},
)
)
regions = text_regions + table_regions + image_regions
ordered_boxes = _sort_reading_order(
[(r.x0, r.y0, r.x1, r.y1) for r in regions],
page_width,
)
order = {box: i for i, box in enumerate(ordered_boxes)}
regions.sort(key=lambda r: order.get((r.x0, r.y0, r.x1, r.y1), len(order)))
return PageLayout(
page_index=page_index,
page_width=page_width,
page_height=page_height,
body_font_size=body_size,
regions=regions,
)
+446
View File
@@ -0,0 +1,446 @@
"""Filter page headers, footers, page numbers and other margin noise from PDF blocks."""
from __future__ import annotations
import re
from collections import defaultdict
from rag_cut.models import Block, BlockType
HEADER_ZONE_RATIO = 0.12
FOOTER_ZONE_RATIO = 0.10
RUNNING_HEADER_MIN_PAGES = 3
RUNNING_HEADER_PAGE_RATIO = 0.5
MAX_RUNNING_HEADER_LEN = 48
MIN_RUNNING_HEADER_LEN = 3
# Drop decorative fragments (icons/logos) from layout detectors like MinerU.
MIN_IMAGE_SIDE = 40.0
MIN_IMAGE_AREA = 1600.0
# Body paragraphs shorter than this can still exit a TOC zone when they
# clearly are not directory entries (keeps in-section catalogs like 形態指標).
TOC_BODY_EXIT_CHARS = 48
TOC_PAGE_ENTRY_RATIO = 0.55
TOC_PAGE_MIN_ENTRIES = 3
_PAGE_NUM_RE = re.compile(r"^\d{1,4}$")
_NUMBERED_SECTION_RE = re.compile(
r"^\s*(\d+(?:\.\d+)*)(?:\.|.)?\s*([A-Za-z0-9\u4e00-\u9fff][A-Za-z0-9\u4e00-\u9fff&/ \-_::]{2,})\s*$"
)
# Cover-page / front-matter directory headings only (exact-ish).
_TOC_TITLE_RE = re.compile(
r"^\s*(?:"
r"contents|table\s+of\s+contents|toc|"
r"\u76ee\u5f55|\u76ee\u9304|\u76ee\u6b21|" # 目录 / 目錄 / 目次
r"\u7ae0\u8282\u76ee\u5f55|\u7ae0\u7bc0\u76ee\u9304|" # 章节目录 / 章節目錄
r"list\s+of\s+(?:figures|tables|contents)"
r")[\s.::·•…-]*$",
re.I,
)
_DOT_LEADER_RE = re.compile(r"(?:\.{2,}|\u2026{2,}|\u00b7{2,}|\u2022{2,})")
# Classic TOC line: title …… 12 / 1. Login ..... 3
_TOC_ENTRY_LINE_RE = re.compile(
r"^\s*.{1,120}?"
r"(?:"
r"(?:\.{2,}|\u2026{2,}|\u00b7{2,}|\s{2,})"
r"\s*\d{1,4}"
r"|"
r"(?:\.{2,}|\u2026+)\s*\d{1,4}"
r")"
r"\s*$"
)
# Numbered entry with trailing page: "1.1 Account Status 12"
_TOC_NUMBERED_PAGE_RE = re.compile(
r"^\s*\d+(?:\.\d+)*(?:[\..]\s*|\s+)"
r".{1,100}?"
r"(?:\s{2,}|\s+)"
r"\d{1,4}\s*$"
)
_MD_TOC_LINK_RE = re.compile(r"^\s*[-*+]\s+\[[^\]]+\]\([^)]+\)\s*$")
def _looks_like_numbered_section(text: str) -> bool:
match = _NUMBERED_SECTION_RE.match(text.strip())
return bool(match and len(match.group(2).strip()) >= 3)
def _norm_text(text: str) -> str:
return " ".join((text or "").split())
def is_toc_title_text(text: str) -> bool:
"""True for standalone 目录 / Contents / TOC headings."""
return bool(_TOC_TITLE_RE.match(_norm_text(text)))
def is_toc_entry_line(text: str) -> bool:
"""True for a single TOC row (leaders / trailing page number / md link)."""
stripped = (text or "").strip()
if not stripped or len(stripped) > 200:
return False
if is_toc_title_text(stripped):
return False
if _MD_TOC_LINK_RE.match(stripped):
return True
if _DOT_LEADER_RE.search(stripped) and re.search(r"\d\s*$", stripped):
return True
if _TOC_ENTRY_LINE_RE.match(stripped):
return True
if _TOC_NUMBERED_PAGE_RE.match(stripped):
return True
return False
def is_toc_noise_text(text: str) -> bool:
"""True if the whole block text is a TOC title or TOC entries."""
stripped = (text or "").strip()
if not stripped:
return False
if is_toc_title_text(stripped):
return True
lines = [ln.strip() for ln in stripped.splitlines() if ln.strip()]
if not lines:
return False
if len(lines) == 1:
return is_toc_entry_line(lines[0])
hits = sum(1 for ln in lines if is_toc_entry_line(ln) or is_toc_title_text(ln))
return hits >= max(2, int(len(lines) * 0.6))
def _is_toc_zone_exit_block(block: Block) -> bool:
"""Substantial body / media ends a front-matter TOC stretch."""
if block.type in {BlockType.IMAGE, BlockType.TABLE}:
return True
text = (block.text or block.markdown or "").strip()
if not text or is_toc_noise_text(text):
return False
# Real section start right after TOC (no page-number trailer).
if _looks_like_numbered_section(text) and not is_toc_entry_line(text):
return True
if len(text) >= TOC_BODY_EXIT_CHARS and not is_toc_entry_line(text):
return True
return False
def _toc_heavy_pages(blocks: list[Block]) -> set[int]:
"""Pages dominated by directory lines are dropped wholesale (text only)."""
by_page: dict[int, list[Block]] = defaultdict(list)
for block in blocks:
if block.type in {BlockType.IMAGE, BlockType.TABLE}:
continue
page = block.meta.get("page")
if page is None:
continue
try:
page_key = int(page)
except (TypeError, ValueError):
continue
by_page[page_key].append(block)
heavy: set[int] = set()
for page_key, page_blocks in by_page.items():
texts = [(b.text or b.markdown or "").strip() for b in page_blocks]
texts = [t for t in texts if t]
if not texts:
continue
entry_hits = sum(1 for t in texts if is_toc_noise_text(t))
has_title = any(is_toc_title_text(t) for t in texts)
if entry_hits >= TOC_PAGE_MIN_ENTRIES and entry_hits / len(texts) >= TOC_PAGE_ENTRY_RATIO:
heavy.add(page_key)
elif has_title and entry_hits >= 1 and entry_hits / len(texts) >= 0.4:
heavy.add(page_key)
return heavy
def _toc_noise_bottoms(blocks: list[Block]) -> dict[int, float]:
"""Lowest TOC text position per page, so real content below it survives."""
bottoms: dict[int, float] = {}
for block in blocks:
if block.type in {BlockType.IMAGE, BlockType.TABLE}:
continue
text = (block.text or block.markdown or "").strip()
if not is_toc_noise_text(text):
continue
page = block.meta.get("page")
bbox = block.meta.get("bbox")
if page is None or not isinstance(bbox, (list, tuple)) or len(bbox) < 4:
continue
try:
page_key = int(page)
bottom = float(bbox[3])
except (TypeError, ValueError):
continue
bottoms[page_key] = max(bottoms.get(page_key, 0.0), bottom)
return bottoms
def filter_toc_blocks(blocks: list[Block]) -> list[Block]:
"""Drop document directories (目录 / Contents) — never chunk them.
Keeps in-section numbered catalogs without page leaders (e.g. 形態指標 list).
"""
if not blocks:
return blocks
toc_pages = _toc_heavy_pages(blocks)
toc_bottoms = _toc_noise_bottoms(blocks)
filtered: list[Block] = []
in_toc = False
for block in blocks:
page = block.meta.get("page")
try:
page_key = int(page) if page is not None else None
except (TypeError, ValueError):
page_key = None
text = (block.text or block.markdown or "").strip()
if page_key is not None and page_key in toc_pages and block.type not in {
BlockType.IMAGE,
BlockType.TABLE,
}:
bbox = block.meta.get("bbox")
toc_bottom = toc_bottoms.get(page_key)
if (
toc_bottom is None
or not isinstance(bbox, (list, tuple))
or len(bbox) < 4
or float(bbox[1]) <= toc_bottom
):
continue
if is_toc_title_text(text):
in_toc = True
continue
if in_toc:
if _is_toc_zone_exit_block(block):
in_toc = False
filtered.append(block)
continue
if is_toc_noise_text(text) or is_toc_entry_line(text):
continue
# Ambiguous short line right after TOC: treat as first real section.
in_toc = False
filtered.append(block)
continue
if is_toc_noise_text(text):
continue
filtered.append(block)
return filtered
def _bbox4(block: Block) -> list[float] | None:
bbox = block.meta.get("bbox")
if not isinstance(bbox, (list, tuple)) or len(bbox) < 4:
return None
try:
return [float(bbox[0]), float(bbox[1]), float(bbox[2]), float(bbox[3])]
except (TypeError, ValueError):
return None
def _bbox_area(bbox: list[float]) -> float:
return max(bbox[2] - bbox[0], 0.0) * max(bbox[3] - bbox[1], 0.0)
def _center_inside(inner: list[float], outer: list[float]) -> bool:
cx = (inner[0] + inner[2]) / 2
cy = (inner[1] + inner[3]) / 2
return outer[0] <= cx <= outer[2] and outer[1] <= cy <= outer[3]
def _overlap_ratio(inner: list[float], outer: list[float]) -> float:
x0 = max(inner[0], outer[0])
y0 = max(inner[1], outer[1])
x1 = min(inner[2], outer[2])
y1 = min(inner[3], outer[3])
if x1 <= x0 or y1 <= y0:
return 0.0
inter = (x1 - x0) * (y1 - y0)
return inter / max(_bbox_area(inner), 1e-6)
def is_tiny_image_block(block: Block) -> bool:
"""True for tiny image fragments that are usually logos/icons, not content figures."""
if block.type != BlockType.IMAGE:
return False
bbox = _bbox4(block)
if not bbox:
return False
width = bbox[2] - bbox[0]
height = bbox[3] - bbox[1]
if width <= 0 or height <= 0:
return True
if width < MIN_IMAGE_SIDE and height < MIN_IMAGE_SIDE:
return True
return width * height < MIN_IMAGE_AREA
def filter_nested_image_blocks(blocks: list[Block]) -> list[Block]:
"""Drop image fragments whose center lies inside a larger same-page image."""
image_boxes: list[tuple[int, int, list[float], float]] = []
for idx, block in enumerate(blocks):
if block.type != BlockType.IMAGE:
continue
bbox = _bbox4(block)
page = block.meta.get("page")
if bbox is None or page is None:
continue
try:
page_key = int(page)
except (TypeError, ValueError):
continue
image_boxes.append((idx, page_key, bbox, _bbox_area(bbox)))
drop: set[int] = set()
for idx, page_key, bbox, area in image_boxes:
for other_idx, other_page, other_bbox, other_area in image_boxes:
if idx == other_idx or page_key != other_page or other_area <= area * 1.5:
continue
if _center_inside(bbox, other_bbox) and _overlap_ratio(bbox, other_bbox) >= 0.85:
drop.add(idx)
break
if not drop:
return blocks
return [block for idx, block in enumerate(blocks) if idx not in drop]
def _infer_page_heights(blocks: list[Block]) -> dict[int, float]:
heights: dict[int, float] = {}
for block in blocks:
page = block.meta.get("page")
if page is None:
continue
try:
page_key = int(page)
except (TypeError, ValueError):
continue
explicit = block.meta.get("page_height")
if explicit:
heights[page_key] = max(heights.get(page_key, 0), float(explicit))
bbox = block.meta.get("bbox")
if bbox and len(bbox) >= 4:
heights[page_key] = max(heights.get(page_key, 0), float(bbox[3]))
return {page: max(height, 1.0) for page, height in heights.items()}
def _page_height(block: Block, heights: dict[int, float]) -> float | None:
page = block.meta.get("page")
if page is None:
explicit = block.meta.get("page_height")
return float(explicit) if explicit else None
try:
page_key = int(page)
except (TypeError, ValueError):
return block.meta.get("page_height")
return heights.get(page_key) or block.meta.get("page_height")
def _in_margin_zone(bbox: list[float], page_height: float, *, header: bool) -> bool:
y_mid = (float(bbox[1]) + float(bbox[3])) / 2
if header:
return y_mid < page_height * HEADER_ZONE_RATIO
return y_mid > page_height * (1 - FOOTER_ZONE_RATIO)
def is_margin_noise_block(block: Block, page_height: float | None) -> bool:
"""Heuristic margin noise filter for blocks missing explicit MinerU region types."""
if block.type == BlockType.IMAGE:
return is_tiny_image_block(block)
if block.type == BlockType.TABLE:
return False
text = (block.text or "").strip()
if not text:
return False
bbox = block.meta.get("bbox")
if not page_height or not bbox or len(bbox) < 4:
return block.type == BlockType.PARAGRAPH and bool(_PAGE_NUM_RE.match(text))
in_header = _in_margin_zone(bbox, page_height, header=True)
in_footer = _in_margin_zone(bbox, page_height, header=False)
if not in_header and not in_footer:
return False
if _PAGE_NUM_RE.match(text):
return True
if block.type != BlockType.PARAGRAPH:
return False
# Numbered section titles often sit at the top of a continued page — keep them.
if _looks_like_numbered_section(text):
return False
return len(text) <= MAX_RUNNING_HEADER_LEN
def detect_running_header_texts(blocks: list[Block]) -> set[str]:
"""Texts that repeat across many pages are likely running headers/footers.
Only paragraphs in the header/footer margin bands are considered, so
real section headings that happen to repeat are not wiped document-wide.
"""
pages_by_text: dict[str, set[int]] = defaultdict(set)
all_pages: set[int] = set()
heights = _infer_page_heights(blocks)
for block in blocks:
if block.type != BlockType.PARAGRAPH:
continue
text = (block.text or "").strip()
if not (MIN_RUNNING_HEADER_LEN <= len(text) <= MAX_RUNNING_HEADER_LEN):
continue
if _looks_like_numbered_section(text):
continue
page = block.meta.get("page")
if page is None:
continue
try:
page_key = int(page)
except (TypeError, ValueError):
continue
page_height = _page_height(block, heights)
bbox = _bbox4(block)
if page_height and bbox:
in_margin = _in_margin_zone(bbox, float(page_height), header=True) or _in_margin_zone(
bbox, float(page_height), header=False
)
if not in_margin:
continue
all_pages.add(page_key)
pages_by_text[text].add(page_key)
if len(all_pages) < RUNNING_HEADER_MIN_PAGES:
return set()
threshold = max(RUNNING_HEADER_MIN_PAGES, int(len(all_pages) * RUNNING_HEADER_PAGE_RATIO))
return {text for text, pages in pages_by_text.items() if len(pages) >= threshold}
def filter_noise_blocks(blocks: list[Block]) -> list[Block]:
"""Remove headers, footers, page numbers, TOC and other repeated margin noise."""
if not blocks:
return blocks
heights = _infer_page_heights(blocks)
running_headers = detect_running_header_texts(blocks)
filtered: list[Block] = []
for block in blocks:
text = (block.text or "").strip()
# Never wipe real section headings via running-header equality.
if text and text in running_headers and block.type == BlockType.PARAGRAPH:
continue
page_height = _page_height(block, heights)
if is_margin_noise_block(block, float(page_height) if page_height else None):
continue
filtered.append(block)
return filter_nested_image_blocks(filter_toc_blocks(filtered))
+57
View File
@@ -0,0 +1,57 @@
"""Optional OCR for cropped image/table regions."""
from __future__ import annotations
import io
from pathlib import Path
import fitz
def ocr_pixmap(pix: fitz.Pixmap) -> str:
"""Run OCR on a pixmap; returns empty string when OCR is unavailable."""
try:
import pytesseract
from PIL import Image
except ImportError:
return ""
try:
image = Image.open(io.BytesIO(pix.tobytes("png")))
return _ocr_pil(image)
except Exception:
return ""
def ocr_image_file(path: Path) -> str:
try:
import pytesseract # noqa: F401
from PIL import Image
except ImportError:
return ""
try:
return _ocr_pil(Image.open(path))
except Exception:
return ""
def _ocr_pil(image) -> str:
import pytesseract
for lang in ("chi_tra+eng", "chi_sim+eng", "eng"):
try:
text = pytesseract.image_to_string(image, lang=lang)
cleaned = " ".join(text.split())
if cleaned:
return cleaned
except Exception:
continue
return ""
def describe_visual(ocr_text: str, kind: str, label: str) -> str:
"""Lightweight image/table description without an external vision model."""
if ocr_text:
return f"{kind}:{ocr_text[:300]}"
return f"{kind}:{label}"
+213
View File
@@ -0,0 +1,213 @@
"""End-to-end PyMuPDF PDF pipeline matching the architecture diagram."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import pdfplumber
from rag_cut.models import Block
from rag_cut.parsers.pdf.layout import (
LayoutRegion,
PageLayout,
_rect_overlap,
_sort_reading_order,
analyze_page_layout,
render_page,
)
from rag_cut.parsers.pdf.standardize import open_document
from rag_cut.parsers.pdf.tables import rows_to_markdown
from rag_cut.parsers.pdf.text_extract import extract_text_blocks
from rag_cut.parsers.pdf.visual_extract import extract_visual_blocks
@dataclass
class _PageItem:
y0: float
x0: float
kind: str
block: Block
def _pdfplumber_tables(path: Path) -> dict[int, list[dict]]:
"""Extract tables per page with bbox from pdfplumber."""
tables_by_page: dict[int, list[dict]] = {}
try:
with pdfplumber.open(path) as pdf:
for i, page in enumerate(pdf.pages):
found = []
try:
for idx, table in enumerate(page.find_tables()):
rows = table.extract() or []
if not rows:
continue
bbox = table.bbox
found.append(
{
"rows": rows,
"table_index": idx,
"source": "pdfplumber",
"bbox": list(bbox) if bbox else None,
}
)
except Exception:
rows_list = page.extract_tables() or []
for idx, rows in enumerate(rows_list):
found.append({"rows": rows, "table_index": idx, "source": "pdfplumber", "bbox": None})
if found:
tables_by_page[i] = found
except Exception:
pass
return tables_by_page
def _inject_pdfplumber_tables(layout: PageLayout, plumber_tables: list[dict]) -> PageLayout:
if not plumber_tables:
return layout
if any(r.kind == "table" for r in layout.regions):
return layout
for item in plumber_tables:
rows = item.get("rows") or []
md = rows_to_markdown(rows)
if not md:
continue
bbox = item.get("bbox")
if bbox and len(bbox) == 4:
x0, y0, x1, y1 = bbox
else:
idx = int(item.get("table_index", 0))
x0, y0 = 0, layout.page_height * (idx + 1) / (len(plumber_tables) + 1)
x1, y1 = layout.page_width, layout.page_height * (idx + 2) / (len(plumber_tables) + 1)
layout.regions.append(
LayoutRegion(
x0=x0,
y0=y0,
x1=x1,
y1=y1,
kind="table",
data={
"rows": rows,
"table_index": item.get("table_index", 0),
"source": item.get("source", "pdfplumber"),
},
)
)
ordered_boxes = _sort_reading_order(
[(r.x0, r.y0, r.x1, r.y1) for r in layout.regions],
layout.page_width,
)
order = {box: i for i, box in enumerate(ordered_boxes)}
layout.regions.sort(key=lambda r: order.get((r.x0, r.y0, r.x1, r.y1), len(order)))
return layout
def _bbox_key(bbox: list[float] | tuple[float, ...]) -> tuple[float, ...]:
return tuple(round(v, 1) for v in bbox)
def _region_box(region: LayoutRegion) -> tuple[float, float, float, float]:
return (region.x0, region.y0, region.x1, region.y1)
def _match_block_to_region(
region: LayoutRegion,
candidates: list[Block],
used: set[int],
) -> Block | None:
"""Match a layout region to a parsed block via exact bbox key or overlap."""
region_box = _region_box(region)
key = _bbox_key(region_box)
for block in candidates:
if id(block) in used:
continue
bbox = block.meta.get("bbox")
if bbox and _bbox_key(bbox) == key:
return block
best: Block | None = None
best_overlap = 0.35
for block in candidates:
if id(block) in used:
continue
bbox = block.meta.get("bbox")
if not bbox:
continue
overlap = _rect_overlap(region_box, tuple(bbox))
if overlap > best_overlap:
best_overlap = overlap
best = block
return best
def _merge_page_blocks(
text_blocks: list[Block],
visual_blocks: list[Block],
layout: PageLayout,
) -> list[Block]:
"""Merge text and visual branches in page reading order."""
used: set[int] = set()
merged: list[_PageItem] = []
for region in layout.regions:
if region.kind == "text":
candidates = text_blocks
elif region.kind in {"table", "image"}:
candidates = visual_blocks
else:
continue
block = _match_block_to_region(region, candidates, used)
if block is not None:
merged.append(_PageItem(y0=region.y0, x0=region.x0, kind=region.kind, block=block))
used.add(id(block))
leftovers: list[_PageItem] = []
for block in text_blocks + visual_blocks:
if id(block) not in used:
bbox = block.meta.get("bbox", [0, 0, 0, 0])
leftovers.append(_PageItem(y0=bbox[1], x0=bbox[0], kind=block.type.value, block=block))
merged.extend(sorted(leftovers, key=lambda it: (it.y0, it.x0)))
return [it.block for it in merged]
def parse_pdf(path: Path, assets_dir: Path) -> list[Block]:
"""
PDF pipeline:
1. 文档标准化
2. 页面渲染 + 版面分析
3. 文本块提取 ∥ 图片/表格区域裁剪 + OCR
4. 章节级语义切片(合并两路结果,标注章节上下文)
"""
assets_dir.mkdir(parents=True, exist_ok=True)
renders_dir = assets_dir / "pages"
renders_dir.mkdir(parents=True, exist_ok=True)
doc = open_document(path)
plumber_tables = _pdfplumber_tables(path)
all_blocks: list[Block] = []
chapter_title: str | None = None
img_counter = 0
try:
for page_index, page in enumerate(doc):
layout = analyze_page_layout(page, page_index)
layout = _inject_pdfplumber_tables(layout, plumber_tables.get(page_index, []))
page_render = render_page(page)
render_path = renders_dir / f"page{page_index + 1}.png"
page_render.save(str(render_path))
text_blocks, chapter_title = extract_text_blocks(layout, chapter_title)
visual_blocks, img_counter = extract_visual_blocks(
page, layout, doc, assets_dir, chapter_title, img_counter
)
page_blocks = _merge_page_blocks(text_blocks, visual_blocks, layout)
all_blocks.extend(page_blocks)
finally:
doc.close()
return all_blocks
@@ -0,0 +1,16 @@
"""Document standardization for PDF input."""
from __future__ import annotations
from pathlib import Path
import fitz
def open_document(path: Path) -> fitz.Document:
"""Open and normalize a PDF for downstream layout processing."""
doc = fitz.open(path)
if doc.is_encrypted and not doc.authenticate(""):
doc.close()
raise ValueError(f"Encrypted PDF cannot be opened: {path.name}")
return doc
@@ -0,0 +1,104 @@
"""Build rich table blocks from PDF layout regions."""
from __future__ import annotations
from pathlib import Path
import fitz
from rag_cut.models import Block, BlockType
from rag_cut.parsers.pdf.layout import LayoutRegion, PageLayout
from rag_cut.parsers.pdf.ocr import describe_visual, ocr_pixmap
from rag_cut.parsers.pdf.tables import (
detect_header_row_count,
extract_table_keywords,
header_signature,
normalize_rows,
rows_to_markdown,
split_body_and_footnotes,
)
def _crop_region(page: fitz.Page, region: LayoutRegion, zoom: float = 2.0) -> fitz.Pixmap:
clip = fitz.Rect(region.x0, region.y0, region.x1, region.y1)
return page.get_pixmap(matrix=fitz.Matrix(zoom, zoom), clip=clip, alpha=False)
def _save_pixmap(pix: fitz.Pixmap, path: Path) -> None:
if pix.n - pix.alpha > 3:
pix = fitz.Pixmap(fitz.csRGB, pix)
pix.save(str(path))
def build_table_block(
page: fitz.Page,
region: LayoutRegion,
layout: PageLayout,
assets_dir: Path,
chapter_title: str | None,
) -> Block | None:
"""Extract a table block with rows, screenshot, OCR and structural metadata."""
raw_rows = region.data.get("rows") or []
normalized = normalize_rows(raw_rows)
if not normalized:
return None
header_rows = detect_header_row_count(normalized)
body_rows, footnote_rows, footnotes = split_body_and_footnotes(normalized, header_rows)
data_rows = normalized[:header_rows] + body_rows
if not data_rows:
return None
page_no = layout.page_index + 1
table_index = int(region.data.get("table_index", 0)) + 1
crop_id = f"page{page_no}_table{table_index}.png"
crops_dir = assets_dir / "crops"
crops_dir.mkdir(parents=True, exist_ok=True)
crop_path = crops_dir / crop_id
ocr_text = ""
try:
crop_pix = _crop_region(page, region)
_save_pixmap(crop_pix, crop_path)
ocr_text = ocr_pixmap(crop_pix)
except Exception:
crop_path_str = ""
else:
crop_path_str = str(crop_path)
sig = header_signature(normalized, header_rows)
md = rows_to_markdown(normalized, header_rows=header_rows, include_footnotes=footnotes)
if not md:
return None
header_text = " ".join(" ".join(r for r in row if r) for row in normalized[:header_rows])
keywords = extract_table_keywords(header_text, md, footnotes, ocr_text)
meta: dict = {
"page": page_no,
"pages": [page_no],
"bbox": [region.x0, region.y0, region.x1, region.y1],
"bboxes": [{"page": page_no, "bbox": [region.x0, region.y0, region.x1, region.y1]}],
"table_source": region.data.get("source", "pymupdf"),
"crop_path": crop_path_str,
"rows": normalized,
"row_count": len(normalized),
"col_count": max((len(r) for r in normalized), default=0),
"header_rows": header_rows,
"header_signature": [list(r) for r in sig],
"footnotes": footnotes,
"keywords": keywords,
"cross_page": False,
}
if chapter_title:
meta["chapter"] = chapter_title
return Block(
type=BlockType.TABLE,
markdown=md,
ocr_text=ocr_text,
text=describe_visual(ocr_text, "表格", crop_id) if ocr_text else "",
image_id=crop_id,
image_path=crop_path_str or None,
meta=meta,
)
+177
View File
@@ -0,0 +1,177 @@
"""Cross-page table detection and merging."""
from __future__ import annotations
from rag_cut.models import Block, BlockType
from rag_cut.parsers.pdf.tables import (
detect_header_row_count,
header_signature,
normalize_rows,
rows_to_markdown,
split_body_and_footnotes,
)
def _page(block: Block) -> int | None:
page = block.meta.get("page")
return int(page) if page is not None else None
def _header_rows(block: Block) -> int:
rows = block.meta.get("rows") or []
return int(block.meta.get("header_rows") or detect_header_row_count(normalize_rows(rows)) or 1)
def _header_sig(block: Block) -> tuple[tuple[str, ...], ...]:
if block.meta.get("header_signature"):
raw = block.meta["header_signature"]
return tuple(tuple(r) for r in raw)
rows = normalize_rows(block.meta.get("rows") or [])
return header_signature(rows, _header_rows(block))
def _column_count(block: Block) -> int:
rows = normalize_rows(block.meta.get("rows") or [])
return max((len(r) for r in rows), default=0)
def _similar_columns(a: Block, b: Block) -> bool:
ca, cb = _column_count(a), _column_count(b)
if ca == 0 or cb == 0:
return False
return ca == cb or abs(ca - cb) <= 1
def _repeated_header(rows: list[list[str]], sig: tuple[tuple[str, ...], ...]) -> int:
"""Return number of leading rows in `rows` that repeat the header signature."""
if not sig:
return 0
n = len(sig)
if len(rows) < n:
return 0
if header_signature(rows, n) == sig:
return n
if n == 1 and rows and tuple(rows[0]) == sig[0]:
return 1
return 0
def _can_merge_continuation(prev: Block, nxt: Block) -> bool:
if prev.type != BlockType.TABLE or nxt.type != BlockType.TABLE:
return False
prev_page = _page(prev)
nxt_page = _page(nxt)
if prev_page is None or nxt_page is None or nxt_page != prev_page + 1:
return False
if not _similar_columns(prev, nxt):
return False
sig = _header_sig(prev)
if not sig:
return False
rows_b = normalize_rows(nxt.meta.get("rows") or [])
if not rows_b:
return False
if _repeated_header(rows_b, sig) > 0:
return True
# Continuation without repeated header: similar width and no title on next table
if nxt.meta.get("table_title"):
return False
prev_bbox = prev.meta.get("bbox") or []
nxt_bbox = nxt.meta.get("bbox") or []
if len(prev_bbox) == 4 and len(nxt_bbox) == 4:
prev_width = prev_bbox[2] - prev_bbox[0]
nxt_width = nxt_bbox[2] - nxt_bbox[0]
if prev_width > 0 and abs(prev_width - nxt_width) / prev_width <= 0.15:
return True
return _column_count(prev) == _column_count(nxt)
def _merge_two_tables(prev: Block, nxt: Block) -> Block:
rows_a = normalize_rows(prev.meta.get("rows") or [])
rows_b = normalize_rows(nxt.meta.get("rows") or [])
header_rows = _header_rows(prev)
sig = _header_sig(prev)
skip = _repeated_header(rows_b, sig)
merged_rows = rows_a + rows_b[skip:]
body_rows, _, foot_a = split_body_and_footnotes(rows_a, header_rows)
_, _, foot_b = split_body_and_footnotes(rows_b, skip or header_rows)
footnotes = " ".join(x for x in (prev.meta.get("footnotes") or foot_a, foot_b) if x).strip()
pages = sorted(set((prev.meta.get("pages") or [_page(prev)]) + [_page(nxt)]))
pages = [p for p in pages if p is not None]
bboxes = list(prev.meta.get("bboxes") or [])
if prev.meta.get("bbox"):
bboxes.append({"page": _page(prev), "bbox": prev.meta["bbox"]})
if nxt.meta.get("bbox"):
bboxes.append({"page": _page(nxt), "bbox": nxt.meta["bbox"]})
prev_bbox = prev.meta.get("bbox") or [0, 0, 0, 0]
nxt_bbox = nxt.meta.get("bbox") or prev_bbox
merged_bbox = [
min(prev_bbox[0], nxt_bbox[0]),
min(prev_bbox[1], nxt_bbox[1]),
max(prev_bbox[2], nxt_bbox[2]),
max(prev_bbox[3], nxt_bbox[3]),
]
md = rows_to_markdown(merged_rows, header_rows=header_rows, include_footnotes=footnotes)
meta = dict(prev.meta)
meta.update(
{
"rows": merged_rows,
"row_count": len(merged_rows),
"col_count": max((len(r) for r in merged_rows), default=0),
"header_rows": header_rows,
"header_signature": [list(r) for r in sig],
"footnotes": footnotes,
"pages": pages,
"page": pages[0] if pages else prev.meta.get("page"),
"bbox": merged_bbox,
"bboxes": bboxes,
"cross_page": len(pages) > 1,
"merged_table_count": int(prev.meta.get("merged_table_count") or 1) + 1,
}
)
if nxt.meta.get("following_text") and not meta.get("following_text"):
meta["following_text"] = nxt.meta.get("following_text")
return prev.model_copy(update={"markdown": md, "meta": meta})
def merge_cross_page_tables(blocks: list[Block]) -> list[Block]:
"""Merge consecutive cross-page table blocks that share headers/structure."""
if not blocks:
return []
result: list[Block] = []
i = 0
while i < len(blocks):
current = blocks[i]
if current.type != BlockType.TABLE:
result.append(current)
i += 1
continue
merged = current
j = i + 1
while j < len(blocks):
nxt = blocks[j]
if nxt.type == BlockType.TABLE and _can_merge_continuation(merged, nxt):
merged = _merge_two_tables(merged, nxt)
j += 1
continue
break
result.append(merged)
i = j
return result
+313
View File
@@ -0,0 +1,313 @@
"""Table helpers: normalization, header detection, Markdown export, cross-page merge."""
from __future__ import annotations
import re
from typing import Any
FOOTNOTE_ROW_RE = re.compile(r"^[\s*※①②③④⑤]*(?:注[::]?|备注[::]?|说明[::]?|Note[::]?)", re.I)
DATA_FIRST_CELL_RE = re.compile(r"^[a-z_][a-z0-9_.-]*$", re.I)
TABLE_TITLE_RE = re.compile(
r"^(?:表\s*\d+[::.]?|Table\s*\d+[::.]?|图\s*\d+[::.]?)?\s*.{2,80}$",
re.I,
)
KEYWORD_TERMS = (
"字段", "参数", "必填", "选填", "状态", "类型", "说明", "含义", "取值",
"field", "parameter", "required", "optional", "status", "description",
)
def normalize_cell(value: str | None) -> str:
"""Merge in-cell line breaks; escape pipe chars for Markdown tables."""
if not value:
return ""
text = str(value).replace("\r\n", "\n").replace("\r", "\n")
parts = [p.strip() for p in text.split("\n") if p.strip()]
merged = " ".join(parts) if parts else ""
return merged.replace("|", "\\|")
def normalize_rows(rows: list[list[str | None]]) -> list[list[str]]:
"""Pad columns and normalize every cell without losing row alignment."""
if not rows:
return []
cleaned = [[normalize_cell(c) for c in row] for row in rows]
col_count = max((len(r) for r in cleaned), default=0)
return [row + [""] * (col_count - len(row)) for row in cleaned]
def looks_like_table(rows: list[list[str | None]]) -> bool:
"""Return True only for rows that have a real table-like grid."""
if len(rows) < 2:
return False
cleaned = normalize_rows(rows)
col_count = max((len(r) for r in cleaned), default=0)
if col_count < 2:
return False
non_empty_cells = sum(1 for row in cleaned for cell in row if cell)
rows_with_two_cells = sum(1 for row in cleaned if sum(1 for cell in row if cell) >= 2)
return non_empty_cells >= 4 and rows_with_two_cells >= 2
DESCRIPTION_HINTS = ("必填", "选填", "格式要求", "required", "optional", "格式", "用于标识")
QA_HEADER_TERMS = ("query", "question", "用户输入", "reference_output", "answer", "标准答案", "session")
def _row_fill(row: list[str]) -> int:
return sum(1 for c in row if c)
def looks_like_column_header_row(row: list[str]) -> bool:
"""True when a row looks like short spreadsheet column names."""
filled = [c.strip() for c in row if c and c.strip()]
if len(filled) < 2:
return False
if any(len(c) > 40 for c in filled):
return False
identifier_like = sum(
1
for c in filled
if DATA_FIRST_CELL_RE.match(c) or re.match(r"^[a-z][a-z0-9_]*$", c, re.I)
)
return identifier_like >= max(2, (len(filled) + 1) // 2)
def looks_like_description_row(row: list[str]) -> bool:
"""True when a row is a template field-description line (not data/header)."""
filled = [c.strip() for c in row if c and c.strip()]
if not filled:
return False
if max(len(c) for c in filled) >= 48:
return True
return sum(1 for c in filled if any(h in c for h in DESCRIPTION_HINTS)) >= 2
def detect_header_row_count(rows: list[list[str]]) -> int:
"""Detect 1-2 header rows from content patterns."""
if len(rows) < 2:
return 1 if rows else 0
first_fill = _row_fill(rows[0])
second_fill = _row_fill(rows[1]) if len(rows) > 1 else 0
if first_fill < 2:
return 0
if len(rows) > 2 and second_fill >= 2:
first_short = all(len(c) <= 24 for c in rows[0] if c)
second_short = all(len(c) <= 24 for c in rows[1] if c)
second_is_data = bool(rows[1][0]) and DATA_FIRST_CELL_RE.match(rows[1][0])
third_data_like = _row_fill(rows[2]) >= max(1, first_fill - 1)
if first_short and second_short and third_data_like and not second_is_data:
return 2
return 1
def detect_spreadsheet_layout(rows: list[list[str]]) -> dict[str, int]:
"""
Detect spreadsheet preamble/header/data boundaries (1-based row numbers).
Common template: row 1 = field descriptions, row 2 = column names, row 3+ = data.
"""
normalized = normalize_rows(rows)
if not normalized:
return {
"preamble_rows": 0,
"header_rows": 1,
"header_row_start": 1,
"header_row_end": 1,
"data_start_row": 2,
}
preamble = 0
header_index = 0
if (
len(normalized) >= 3
and looks_like_description_row(normalized[0])
and looks_like_column_header_row(normalized[1])
):
preamble = 1
header_index = 1
header_rows = 1
else:
header_rows = detect_header_row_count(normalized)
header_index = preamble
header_end_index = header_index + header_rows - 1
data_start_index = header_end_index + 1
return {
"preamble_rows": preamble,
"header_rows": header_rows,
"header_row_start": header_index + 1,
"header_row_end": header_end_index + 1,
"data_start_row": data_start_index + 1,
}
def is_qa_style_table(rows: list[list[str]], layout: dict[str, int] | None = None) -> bool:
"""True for evaluation/Q&A sheets where each row should become one chunk."""
if not rows:
return False
layout = layout or detect_spreadsheet_layout(rows)
h_start = layout["header_row_start"] - 1
h_end = layout["header_row_end"]
header_text = " ".join(
(cell or "").lower() for row in rows[h_start:h_end] for cell in row if cell
)
return sum(1 for term in QA_HEADER_TERMS if term in header_text) >= 2
def split_body_and_footnotes(rows: list[list[str]], header_rows: int) -> tuple[list[list[str]], list[list[str]], str]:
"""Separate data rows from trailing footnote rows."""
if header_rows >= len(rows):
return [], [], ""
body = rows[header_rows:]
footnote_rows: list[list[str]] = []
while body:
first_cell = (body[-1][0] if body[-1] else "") or ""
joined = " ".join(c for c in body[-1] if c)
if FOOTNOTE_ROW_RE.match(first_cell) or FOOTNOTE_ROW_RE.match(joined):
footnote_rows.insert(0, body.pop())
elif len(joined) <= 80 and any(k in joined for k in ("注", "备注", "说明", "Note")):
footnote_rows.insert(0, body.pop())
else:
break
footnotes = " ".join(" ".join(c for c in row if c) for row in footnote_rows).strip()
return body, footnote_rows, footnotes
def header_signature(rows: list[list[str]], header_rows: int) -> tuple[tuple[str, ...], ...]:
if header_rows <= 0:
return ()
return tuple(tuple(row) for row in rows[:header_rows])
def rows_to_markdown(
rows: list[list[str | None]],
header_rows: int | None = None,
include_footnotes: str = "",
) -> str:
"""Render rows as a standard Markdown table."""
if not looks_like_table(rows):
return ""
normalized = normalize_rows(rows)
if header_rows is None:
header_rows = detect_header_row_count(normalized)
body_rows, _, inline_footnotes = split_body_and_footnotes(normalized, header_rows)
data_rows = normalized[:header_rows] + body_rows
if not data_rows:
return ""
col_count = max(len(r) for r in data_rows)
lines: list[str] = []
for i, row in enumerate(data_rows):
padded = row + [""] * (col_count - len(row))
lines.append("| " + " | ".join(padded) + " |")
if i == header_rows - 1:
lines.append("| " + " | ".join(["---"] * col_count) + " |")
md = "\n".join(lines)
footnotes = include_footnotes or inline_footnotes
if footnotes:
md += f"\n\n*{footnotes}*"
return md
def guess_table_title(text: str) -> str | None:
"""Guess table title from a short preceding line."""
cleaned = normalize_cell(text)
if not cleaned or len(cleaned) > 120:
return None
if TABLE_TITLE_RE.match(cleaned):
return cleaned
if cleaned.endswith("表") or cleaned.endswith("列表") or cleaned.endswith("说明"):
return cleaned
if re.match(r"^表\s*\d+", cleaned):
return cleaned
return None
def extract_table_keywords(*texts: str, limit: int = 20) -> list[str]:
"""Extract retrieval keywords from table title, headers and body."""
source = " ".join(t for t in texts if t)
words = re.findall(r"[A-Za-z][A-Za-z0-9_-]{2,}|[\u4e00-\u9fff]{2,}", source)
seen: set[str] = set()
result: list[str] = []
for word in words:
if word in seen:
continue
seen.add(word)
result.append(word)
if len(result) >= limit:
break
for term in KEYWORD_TERMS:
if term.lower() in source.lower() and term not in seen:
result.append(term)
seen.add(term)
if len(result) >= limit:
break
return result[:limit]
def build_table_embedding_text(
*,
chapter: str = "",
table_title: str = "",
markdown: str = "",
description: str = "",
footnotes: str = "",
keywords: list[str] | None = None,
ocr_text: str = "",
) -> str:
"""Compose embedding text: chapter + title + markdown + description + keywords."""
parts: list[str] = []
if chapter:
parts.append(f"章节标题:{chapter}")
if table_title:
parts.append(f"表格标题:{table_title}")
if markdown:
parts.append(markdown)
if description:
parts.append(f"表格说明:{description}")
if footnotes:
parts.append(f"脚注说明:{footnotes}")
if ocr_text:
parts.append(f"表格 OCR:{ocr_text}")
if keywords:
parts.append(f"关键词:{','.join(keywords)}")
return "\n".join(parts)
def table_meta_summary(meta: dict[str, Any]) -> dict[str, Any]:
"""Pick table-specific fields for chunk metadata."""
keys = (
"table_title",
"table_description",
"header_rows",
"header_signature",
"footnotes",
"keywords",
"chapter",
"pages",
"page",
"bbox",
"bboxes",
"crop_path",
"image_path",
"row_count",
"col_count",
"cross_page",
"table_source",
"preceding_text",
"following_text",
"nearest_heading",
"parent_heading",
)
return {k: meta[k] for k in keys if k in meta and meta[k] not in (None, "", [], {})}
+104
View File
@@ -0,0 +1,104 @@
"""Text block extraction from analyzed layout."""
from __future__ import annotations
import re
from rag_cut.models import Block, BlockType
from rag_cut.parsers.pdf.layout import LayoutRegion, PageLayout
_UI_LABEL_RE = re.compile(r"^[\d\s\W]{0,6}[\u4e00-\u9fff]{1,6}$")
# "10.上升三角形態" / "1.2 ACCOUNT STATUS" (space after number optional)
_NUMBERED_HEADING_RE = re.compile(
r"^\s*(\d+(?:\.\d+)*)(?:\.|.)?\s*([A-Za-z0-9\u4e00-\u9fff][A-Za-z0-9\u4e00-\u9fff&/ \-_::]{2,})\s*$"
)
# Title+body merges must not become headings; real section titles stay shorter.
MAX_HEADING_CHARS = 100
def _numbered_heading_level(text: str) -> int | None:
match = _NUMBERED_HEADING_RE.match(text.strip())
if not match or len(match.group(2).strip()) < 3:
return None
return match.group(1).count(".") + 1
def _looks_like_ui_label(text: str) -> bool:
"""True for short UI chips; exclude numbered / CJK section titles."""
stripped = text.strip()
if _numbered_heading_level(stripped) is not None:
return False
if re.match(r"^\d+(?:\.\d+)*(?:\.|.)", stripped):
return False
# Short Chinese section banners like 「形態指標」are not toolbar labels.
if re.fullmatch(r"[\u4e00-\u9fff]{2,12}", stripped):
return False
return bool(_UI_LABEL_RE.match(stripped))
def _font_heading_level(size: float, body_size: float, text: str) -> int | None:
stripped = text.strip()
numbered = _numbered_heading_level(stripped)
if numbered and len(stripped) <= MAX_HEADING_CHARS:
return numbered
if len(stripped) < 4 or len(stripped) > MAX_HEADING_CHARS:
return None
if _looks_like_ui_label(stripped):
return None
# Compact CJK section titles such as 「形態指標」.
if (
size >= body_size + 3
and re.fullmatch(r"[\u4e00-\u9fff]{2,12}", stripped)
and not _numbered_heading_level(stripped)
):
return 1
if size >= body_size + 6:
return 1 if len(stripped) >= 10 else 2
if size >= body_size + 3:
return 2 if len(stripped) >= 8 else 3
if size >= body_size + 1.5:
return 3
return None
def extract_text_blocks(layout: PageLayout, chapter_title: str | None = None) -> tuple[list[Block], str | None]:
"""Extract heading/paragraph blocks; update chapter title when headings appear."""
blocks: list[Block] = []
current_chapter = chapter_title
for region in layout.regions:
if region.kind != "text":
continue
text = region.data.get("text", "").strip()
if not text:
continue
font_size = float(region.data.get("font_size", layout.body_font_size))
level = _font_heading_level(font_size, layout.body_font_size, text)
page_no = layout.page_index + 1
meta = {
"page": page_no,
"bbox": [region.x0, region.y0, region.x1, region.y1],
"font_size": font_size,
"body_font_size": layout.body_font_size,
"page_height": layout.page_height,
}
if current_chapter:
meta["chapter"] = current_chapter
if level:
if level <= 2:
current_chapter = text
if current_chapter:
meta["chapter"] = current_chapter
blk = Block(type=BlockType.HEADING, text=text, level=level, meta=meta)
else:
blk = Block(type=BlockType.PARAGRAPH, text=text, meta=meta)
blocks.append(blk)
return blocks, current_chapter
@@ -0,0 +1,91 @@
"""Image/table region cropping with OCR and descriptions."""
from __future__ import annotations
from pathlib import Path
import fitz
from rag_cut.models import Block, BlockType
from rag_cut.parsers.pdf.layout import LayoutRegion, PageLayout
from rag_cut.parsers.pdf.ocr import describe_visual, ocr_image_file, ocr_pixmap
from rag_cut.parsers.pdf.table_extract import build_table_block
def _crop_region(page: fitz.Page, region: LayoutRegion, zoom: float = 2.0) -> fitz.Pixmap:
clip = fitz.Rect(region.x0, region.y0, region.x1, region.y1)
return page.get_pixmap(matrix=fitz.Matrix(zoom, zoom), clip=clip, alpha=False)
def _save_pixmap(pix: fitz.Pixmap, path: Path) -> None:
if pix.n - pix.alpha > 3:
pix = fitz.Pixmap(fitz.csRGB, pix)
pix.save(str(path))
def extract_visual_blocks(
page: fitz.Page,
layout: PageLayout,
doc: fitz.Document,
assets_dir: Path,
chapter_title: str | None,
img_counter: int,
) -> tuple[list[Block], int]:
blocks: list[Block] = []
page_no = layout.page_index + 1
for region in layout.regions:
if region.kind == "table":
table_block = build_table_block(page, region, layout, assets_dir, chapter_title)
if table_block:
blocks.append(table_block)
continue
if region.kind != "image":
continue
img_counter += 1
xref = region.data.get("xref")
img_id = f"page{page_no}_img{img_counter}.png"
img_path = assets_dir / img_id
try:
# Prefer on-page crop so scaled/clipped placements OCR the visible region.
crop_pix = _crop_region(page, region)
_save_pixmap(crop_pix, img_path)
except Exception:
try:
pix = fitz.Pixmap(doc, xref)
if pix.n - pix.alpha > 3:
pix = fitz.Pixmap(fitz.csRGB, pix)
pix.save(str(img_path))
except Exception:
continue
ocr_text = ocr_image_file(img_path)
if not ocr_text:
try:
ocr_text = ocr_pixmap(_crop_region(page, region))
except Exception:
ocr_text = ""
meta = {
"page": page_no,
"bbox": [region.x0, region.y0, region.x1, region.y1],
"xref": xref,
}
if chapter_title:
meta["chapter"] = chapter_title
blocks.append(
Block(
type=BlockType.IMAGE,
image_id=img_id,
image_path=str(img_path),
ocr_text=ocr_text,
text=describe_visual(ocr_text, "图片", img_id),
meta=meta,
)
)
return blocks, img_counter
+31
View File
@@ -0,0 +1,31 @@
"""PDF parser entry point — delegates to the PyMuPDF pipeline."""
from __future__ import annotations
import os
from pathlib import Path
from rag_cut.models import Block
from rag_cut.parsers.base import BaseParser
from rag_cut.parsers.mineru_adapter import parse_pdf_with_mineru
from rag_cut.parsers.pdf.noise_filter import filter_noise_blocks
from rag_cut.parsers.pdf.pipeline import parse_pdf
class PdfParser(BaseParser):
def parse(self, path: Path, assets_dir: Path) -> list[Block]:
engine = os.getenv("RAG_CUT_PDF_ENGINE", "auto").strip().lower()
if engine not in {"auto", "mineru", "pymupdf"}:
engine = "auto"
blocks: list[Block] | None = None
if engine in {"auto", "mineru"}:
blocks = parse_pdf_with_mineru(path, assets_dir)
if blocks is None and engine == "mineru":
raise RuntimeError("MinerU parsing failed or MinerU is not installed")
if blocks is None:
blocks = parse_pdf(path, assets_dir)
# filter_noise_blocks already includes TOC removal.
return filter_noise_blocks(blocks)
+30
View File
@@ -0,0 +1,30 @@
"""PPT/PPTX parser: convert to PDF, then reuse PdfParser."""
from __future__ import annotations
from pathlib import Path
from rag_cut.models import Block
from rag_cut.parsers.base import BaseParser
from rag_cut.parsers.office_to_pdf import convert_to_pdf
from rag_cut.parsers.pdf_parser import PdfParser
class PptParser(BaseParser):
"""Parse presentations by converting to PDF and delegating to PdfParser."""
def __init__(self) -> None:
self._pdf_parser = PdfParser()
def parse(self, path: Path, assets_dir: Path) -> list[Block]:
assets_dir.mkdir(parents=True, exist_ok=True)
pdf_dir = assets_dir / "_conversion"
pdf_path = convert_to_pdf(path, pdf_dir)
blocks = self._pdf_parser.parse(pdf_path, assets_dir)
source_fmt = path.suffix.lower().lstrip(".")
for block in blocks:
block.meta.setdefault("source_format", source_fmt)
if page := block.meta.get("page"):
block.meta["slide"] = page
return blocks
+44
View File
@@ -0,0 +1,44 @@
"""Route file extensions to parsers."""
from __future__ import annotations
from pathlib import Path
from rag_cut.parsers.base import BaseParser
from rag_cut.parsers.word_parser import WordParser
from rag_cut.parsers.image_parser import ImageParser
from rag_cut.parsers.pdf_parser import PdfParser
from rag_cut.parsers.ppt_parser import PptParser
from rag_cut.parsers.text_parser import TextParser
from rag_cut.parsers.xlsx_parser import SpreadsheetParser
PARSERS: dict[str, BaseParser] = {
".docx": WordParser(),
".doc": WordParser(),
".pdf": PdfParser(),
".ppt": PptParser(),
".pptx": PptParser(),
".ppsx": PptParser(),
".xlsx": SpreadsheetParser(),
".xls": SpreadsheetParser(),
".csv": SpreadsheetParser(),
".md": TextParser(),
".txt": TextParser(),
".html": TextParser(),
".htm": TextParser(),
".json": TextParser(),
".xml": TextParser(),
".log": TextParser(),
".jpg": ImageParser(),
".jpeg": ImageParser(),
".png": ImageParser(),
".bmp": ImageParser(),
".gif": ImageParser(),
}
def get_parser(path: Path) -> BaseParser:
ext = path.suffix.lower()
if ext not in PARSERS:
raise ValueError(f"Unsupported file type: {ext}")
return PARSERS[ext]
+123
View File
@@ -0,0 +1,123 @@
"""Plain-text and markup parsers."""
from __future__ import annotations
import json
import re
from html.parser import HTMLParser
from pathlib import Path
from rag_cut.models import Block, BlockType
from rag_cut.parsers.base import BaseParser
class _HTMLTextExtractor(HTMLParser):
def __init__(self) -> None:
super().__init__()
self._parts: list[str] = []
self._heading: tuple[int, str] | None = None
self._blocks: list[Block] = []
self._current: list[str] = []
self._in_heading = False
def handle_starttag(self, tag: str, attrs) -> None:
if tag in ("h1", "h2", "h3", "h4", "h5", "h6"):
self._flush_paragraph()
self._in_heading = True
self._heading_level = int(tag[1])
def handle_endtag(self, tag: str) -> None:
if tag in ("h1", "h2", "h3", "h4", "h5", "h6"):
text = "".join(self._current).strip()
self._current = []
self._in_heading = False
if text:
self._blocks.append(
Block(type=BlockType.HEADING, text=text, level=self._heading_level)
)
elif tag in ("p", "div", "br", "li"):
self._flush_paragraph()
def handle_data(self, data: str) -> None:
self._current.append(data)
def _flush_paragraph(self) -> None:
text = "".join(self._current).strip()
self._current = []
if text:
self._blocks.append(Block(type=BlockType.PARAGRAPH, text=text))
def get_blocks(self) -> list[Block]:
self._flush_paragraph()
return self._blocks
def _parse_markdown(text: str) -> list[Block]:
blocks: list[Block] = []
for line in text.splitlines():
stripped = line.strip()
if not stripped:
continue
m = re.match(r"^(#{1,6})\s+(.+)$", stripped)
if m:
blocks.append(
Block(type=BlockType.HEADING, text=m.group(2).strip(), level=len(m.group(1)))
)
else:
blocks.append(Block(type=BlockType.PARAGRAPH, text=stripped))
return blocks
def _parse_json(text: str) -> list[Block]:
data = json.loads(text)
blocks: list[Block] = []
if isinstance(data, list):
for i, item in enumerate(data):
blocks.append(
Block(
type=BlockType.CODE,
text=json.dumps(item, ensure_ascii=False, indent=2),
meta={"json_index": i},
)
)
elif isinstance(data, dict):
for key, value in data.items():
blocks.append(
Block(
type=BlockType.CODE,
text=json.dumps({key: value}, ensure_ascii=False, indent=2),
meta={"json_key": key},
)
)
else:
blocks.append(Block(type=BlockType.PARAGRAPH, text=str(data)))
return blocks
class TextParser(BaseParser):
def parse(self, path: Path, assets_dir: Path) -> list[Block]:
for encoding in ("utf-8-sig", "utf-8", "gbk", "latin-1"):
try:
text = path.read_text(encoding=encoding)
break
except UnicodeDecodeError:
continue
else:
raise ValueError(f"Cannot decode text file: {path}")
ext = path.suffix.lower()
if ext == ".md":
return _parse_markdown(text)
if ext in (".html", ".htm"):
parser = _HTMLTextExtractor()
parser.feed(text)
return parser.get_blocks()
if ext == ".json":
return _parse_json(text)
# txt, xml, log — paragraph split on blank lines
blocks: list[Block] = []
for para in re.split(r"\n\s*\n", text):
para = para.strip()
if para:
blocks.append(Block(type=BlockType.PARAGRAPH, text=para))
return blocks
+29
View File
@@ -0,0 +1,29 @@
"""DOC/DOCX parser: convert to PDF, then reuse PdfParser."""
from __future__ import annotations
from pathlib import Path
from rag_cut.models import Block
from rag_cut.parsers.base import BaseParser
from rag_cut.parsers.office_to_pdf import convert_to_pdf
from rag_cut.parsers.pdf_parser import PdfParser
class WordParser(BaseParser):
"""Parse Word documents by converting to PDF and delegating to PdfParser."""
def __init__(self) -> None:
self._pdf_parser = PdfParser()
def parse(self, path: Path, assets_dir: Path) -> list[Block]:
assets_dir.mkdir(parents=True, exist_ok=True)
pdf_dir = assets_dir / "_conversion"
pdf_path = convert_to_pdf(path, pdf_dir)
blocks = self._pdf_parser.parse(pdf_path, assets_dir)
source_fmt = path.suffix.lower().lstrip(".")
for block in blocks:
block.meta.setdefault("source_format", source_fmt)
block.meta.setdefault("converted_from", source_fmt)
return blocks
+89
View File
@@ -0,0 +1,89 @@
"""Spreadsheet parser (xlsx/xls/csv) producing row-oriented table blocks."""
from __future__ import annotations
import csv
from pathlib import Path
from rag_cut.models import Block, BlockType
from rag_cut.parsers.base import BaseParser
from rag_cut.parsers.pdf.tables import (
build_table_embedding_text,
detect_spreadsheet_layout,
extract_table_keywords,
normalize_rows,
rows_to_markdown,
)
def _read_csv(path: Path) -> list[list[str]]:
for encoding in ("utf-8-sig", "utf-8", "gbk", "latin-1"):
try:
with open(path, newline="", encoding=encoding) as f:
return [list(row) for row in csv.reader(f)]
except UnicodeDecodeError:
continue
raise ValueError(f"Cannot decode CSV: {path}")
def _read_xlsx(path: Path, sheet_name: str | None = None) -> tuple[str, list[list[str]]]:
import openpyxl
wb = openpyxl.load_workbook(path, read_only=True, data_only=True)
name = sheet_name or wb.sheetnames[0]
ws = wb[name]
rows: list[list[str]] = []
for row in ws.iter_rows(values_only=True):
rows.append(["" if v is None else str(v) for v in row])
wb.close()
# trim trailing empty rows/cols
while rows and all(not c for c in rows[-1]):
rows.pop()
return name, rows
class SpreadsheetParser(BaseParser):
def parse(self, path: Path, assets_dir: Path) -> list[Block]:
ext = path.suffix.lower()
if ext == ".csv":
rows = _read_csv(path)
sheet_name = path.stem
else:
sheet_name, rows = _read_xlsx(path)
if not rows:
return []
normalized = normalize_rows(rows)
layout = detect_spreadsheet_layout(normalized)
preamble = layout["preamble_rows"]
header_rows = layout["header_rows"]
display_rows = normalized[preamble:]
md = rows_to_markdown(display_rows, header_rows=header_rows)
description = ""
if preamble:
description = " ".join(cell for cell in normalized[0] if cell)[:500]
keywords = extract_table_keywords(sheet_name, md, description)
embedding_text = build_table_embedding_text(
table_title=sheet_name,
markdown=md,
description=description,
keywords=keywords,
)
return [
Block(
type=BlockType.TABLE,
markdown=md,
meta={
"sheet": sheet_name,
"table_title": sheet_name,
"row_count": len(normalized),
"col_count": max(len(r) for r in normalized),
"rows": normalized,
"keywords": keywords,
"embedding_text": embedding_text,
"table_description": description or None,
**layout,
},
)
]
+120
View File
@@ -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,
)
+164
View File
@@ -0,0 +1,164 @@
"""Render block groups into chunk markdown strings with layout metadata."""
from __future__ import annotations
from rag_cut.models import Block, BlockType, Chunk
from rag_cut.parsers.pdf.tables import table_meta_summary
def block_to_layout_dict(block: Block) -> dict:
"""Serialize a block for chunk metadata and frontend positional rendering."""
entry: dict = {
"type": block.type.value,
"order_index": block.meta.get("order_index"),
"page": block.meta.get("page"),
"pages": block.meta.get("pages"),
"bbox": block.meta.get("bbox"),
"bboxes": block.meta.get("bboxes"),
"parent_heading": block.meta.get("parent_heading"),
"nearest_heading": block.meta.get("nearest_heading"),
"bound_heading": block.meta.get("bound_heading"),
"preceding_text": block.meta.get("preceding_text"),
"following_text": block.meta.get("following_text"),
}
if block.type == BlockType.HEADING:
entry["text"] = block.text
entry["level"] = block.level
elif block.type == BlockType.IMAGE:
entry["text"] = block.text
entry["image_path"] = block.image_path
entry["image_id"] = block.image_id
entry["ocr_text"] = block.ocr_text
elif block.type == BlockType.TABLE:
entry["text"] = block.markdown or block.text
entry["markdown"] = block.markdown or block.text
entry["ocr_text"] = block.ocr_text
entry["image_path"] = block.image_path
entry["image_id"] = block.image_id
entry["crop_path"] = block.meta.get("crop_path")
entry.update(table_meta_summary(block.meta))
entry["embedding_text"] = block.meta.get("embedding_text")
else:
entry["text"] = block.text or block.markdown
return entry
def collect_chunk_layout_meta(blocks: list[Block]) -> dict:
"""Aggregate page/bbox/image/table metadata for a chunk group."""
pages = sorted({b.meta.get("page") for b in blocks if b.meta.get("page") is not None})
for b in blocks:
for p in b.meta.get("pages") or []:
if p is not None:
pages.append(p)
pages = sorted(set(pages))
bboxes = [
{
"order_index": b.meta.get("order_index"),
"page": b.meta.get("page"),
"bbox": b.meta.get("bbox"),
"type": b.type.value,
}
for b in blocks
if b.meta.get("bbox")
]
images = [
{
"order_index": b.meta.get("order_index"),
"image_id": b.image_id,
"image_path": b.image_path,
"page": b.meta.get("page"),
"bbox": b.meta.get("bbox"),
"ocr_text": b.ocr_text,
"bound_heading": b.meta.get("bound_heading"),
"preceding_text": b.meta.get("preceding_text"),
"following_text": b.meta.get("following_text"),
}
for b in blocks
if b.type == BlockType.IMAGE
]
tables = [
{
"order_index": b.meta.get("order_index"),
"table_title": b.meta.get("table_title"),
"image_path": b.image_path or b.meta.get("crop_path"),
"crop_path": b.meta.get("crop_path"),
"page": b.meta.get("page"),
"pages": b.meta.get("pages"),
"bbox": b.meta.get("bbox"),
"bboxes": b.meta.get("bboxes"),
"markdown": b.markdown,
"ocr_text": b.ocr_text,
"footnotes": b.meta.get("footnotes"),
"keywords": b.meta.get("keywords"),
"nearest_heading": b.meta.get("nearest_heading"),
"chapter": b.meta.get("chapter"),
"row_count": b.meta.get("row_count"),
"col_count": b.meta.get("col_count"),
"header_rows": b.meta.get("header_rows"),
"table_source": b.meta.get("table_source"),
"preceding_text": b.meta.get("preceding_text"),
"following_text": b.meta.get("following_text"),
"cross_page": b.meta.get("cross_page"),
"embedding_text": b.meta.get("embedding_text"),
}
for b in blocks
if b.type == BlockType.TABLE
]
meta: dict = {
"blocks": [block_to_layout_dict(b) for b in blocks],
"bboxes": bboxes,
"images": images,
"tables": tables,
}
if pages:
meta["pages"] = pages
if len(pages) == 1:
meta["page"] = pages[0]
headings = [b.text for b in blocks if b.type == BlockType.HEADING]
if headings:
meta["heading"] = headings[0]
meta["nearest_heading"] = headings[0]
elif blocks and blocks[0].meta.get("nearest_heading"):
meta["nearest_heading"] = blocks[0].meta.get("nearest_heading")
table_blocks = [b for b in blocks if b.type == BlockType.TABLE]
if table_blocks:
primary = table_blocks[0]
meta["table_title"] = primary.meta.get("table_title") or meta.get("nearest_heading")
if not meta.get("chunk_strategy"):
meta["chunk_strategy"] = "table_with_context"
meta["retrieval"] = meta.get("retrieval", True)
if primary.meta.get("embedding_text"):
meta["embedding_text"] = primary.meta["embedding_text"]
if primary.meta.get("keywords"):
meta["keywords"] = primary.meta["keywords"]
order_indices = [b.meta.get("order_index") for b in blocks if b.meta.get("order_index") is not None]
if order_indices:
meta["order_range"] = meta.get("order_range") or [min(order_indices), max(order_indices)]
return meta
def render_blocks(blocks: list[Block], index: int, meta: dict | None = None) -> Chunk:
"""Render blocks in reading order: heading → body → table/image → notes."""
parts: list[str] = []
for block in blocks:
rendered = block.render().strip()
if rendered:
parts.append(rendered)
content = "\n\n".join(parts)
layout_meta = collect_chunk_layout_meta(blocks)
chunk_meta = {**(meta or {}), **layout_meta}
return Chunk(
index=index,
content=content,
char_count=len(content),
block_types=[b.type.value for b in blocks],
meta=chunk_meta,
)
+78
View File
@@ -0,0 +1,78 @@
"""Small dependency-free lexical retriever for the recall demo."""
from __future__ import annotations
import math
import re
from collections import Counter
from rag_cut.models import Chunk
LATIN_TOKEN_RE = re.compile(r"[a-z0-9_]+", re.I)
CJK_RE = re.compile(r"[\u3400-\u9fff]")
def _tokens(text: str) -> list[str]:
normalized = text.lower()
tokens = LATIN_TOKEN_RE.findall(normalized)
cjk = CJK_RE.findall(normalized)
tokens.extend(cjk)
tokens.extend("".join(cjk[i : i + 2]) for i in range(len(cjk) - 1))
return tokens
def _search_text(chunk: Chunk) -> str:
meta = chunk.meta
fields = [
meta.get("heading") or "",
meta.get("nearest_heading") or "",
meta.get("embedding_text") or "",
" ".join(meta.get("keywords") or []),
chunk.content,
]
return "\n".join(str(value) for value in fields if value)
def recall_chunks(query: str, chunks: list[Chunk], top_k: int = 5) -> tuple[list[dict], int]:
"""Rank retrievable chunks with a compact BM25-style lexical score."""
candidates = [chunk for chunk in chunks if chunk.meta.get("retrieval", True)]
query_tokens = _tokens(query)
if not candidates or not query_tokens:
return [], len(candidates)
documents = [_tokens(_search_text(chunk)) for chunk in candidates]
document_frequency = Counter(token for tokens in documents for token in set(tokens))
average_length = sum(len(tokens) for tokens in documents) / len(documents)
query_frequency = Counter(query_tokens)
scored: list[tuple[float, Chunk]] = []
for chunk, tokens in zip(candidates, documents):
frequency = Counter(tokens)
length_normalizer = 1.2 * (0.25 + 0.75 * len(tokens) / max(average_length, 1))
score = 0.0
for token, query_count in query_frequency.items():
term_count = frequency[token]
if not term_count:
continue
inverse_frequency = math.log(1 + (len(documents) - document_frequency[token] + 0.5) / (document_frequency[token] + 0.5))
score += inverse_frequency * ((term_count * 2.2) / (term_count + length_normalizer)) * min(query_count, 2)
if score > 0:
scored.append((score, chunk))
scored.sort(key=lambda item: (-item[0], item[1].index))
results = []
for rank, (score, chunk) in enumerate(scored[:top_k], start=1):
results.append(
{
"rank": rank,
"chunk_index": chunk.index,
"score": round(score, 6),
"content": chunk.content,
"heading": chunk.meta.get("heading") or chunk.meta.get("nearest_heading"),
"pages": chunk.meta.get("pages") or ([chunk.meta["page"]] if chunk.meta.get("page") is not None else []),
"block_types": chunk.block_types,
"parent_chunk_id": chunk.meta.get("parent_chunk_id"),
"is_sub_chunk": bool(chunk.meta.get("is_sub_chunk")),
}
)
return results, len(candidates)
+82
View File
@@ -0,0 +1,82 @@
"""Automatic split policy selection based on parsed document shape."""
from __future__ import annotations
from pathlib import Path
from rag_cut.models import Block, BlockType, SplitConfig, SplitMode
from rag_cut.parsers.pdf.tables import detect_spreadsheet_layout, is_qa_style_table
SPREADSHEET_EXTS = {".xlsx", ".xls", ".csv"}
PRESENTATION_EXTS = {".ppt", ".pptx", ".ppsx"}
TEXT_EXTS = {".md", ".txt", ".html", ".htm", ".json", ".xml", ".log"}
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".gif"}
def _table_rows(blocks: list[Block]) -> int:
if len(blocks) != 1 or blocks[0].type != BlockType.TABLE:
return 0
rows = blocks[0].meta.get("rows") or []
return len(rows)
def _rows_per_chunk(data_row_count: int) -> int:
if data_row_count <= 20:
return max(1, data_row_count)
if data_row_count <= 80:
return 10
if data_row_count <= 300:
return 20
return 40
def choose_split_config(path: Path, blocks: list[Block]) -> SplitConfig:
"""Choose conservative defaults that preserve document structure first."""
ext = path.suffix.lower()
row_count = _table_rows(blocks)
if ext in SPREADSHEET_EXTS or row_count:
rows = blocks[0].meta.get("rows") or [] if blocks else []
layout = detect_spreadsheet_layout(rows) if rows else {
"header_row_start": 1,
"header_row_end": 1,
"data_start_row": 2,
}
data_row_count = max(0, row_count - layout["data_start_row"] + 1)
rows_per = 1 if is_qa_style_table(rows, layout) else _rows_per_chunk(data_row_count)
return SplitConfig(
mode=SplitMode.BY_ROW,
max_chunk_size=2400,
overlap=0,
header_row_start=layout["header_row_start"],
header_row_end=layout["header_row_end"],
start_row=layout["data_start_row"],
rows_per_chunk=rows_per,
)
if ext == ".pdf" or ext in PRESENTATION_EXTS:
return SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2600, overlap=120)
if ext in IMAGE_EXTS:
return SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2200, overlap=0)
if ext in TEXT_EXTS:
return SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=1800, overlap=120)
return SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2200, overlap=120)
def split_config_summary(config: SplitConfig) -> dict[str, int | str | None]:
return {
"mode": config.mode.value,
"delimiter": config.delimiter,
"parent_delimiter": config.parent_delimiter,
"child_delimiter": config.child_delimiter,
"max_chunk_size": config.max_chunk_size,
"child_max_size": config.child_max_size,
"overlap": config.overlap,
"header_row_start": config.header_row_start,
"header_row_end": config.header_row_end,
"start_row": config.start_row,
"rows_per_chunk": config.rows_per_chunk,
}
+8
View File
@@ -0,0 +1,8 @@
"""Chunk splitting strategies."""
from rag_cut.splitters.by_row import split_by_row
from rag_cut.splitters.default_splitter import split_default
from rag_cut.splitters.delimiter import split_by_delimiter
from rag_cut.splitters.parent_child import split_by_parent_child
__all__ = ["split_default", "split_by_delimiter", "split_by_parent_child", "split_by_row"]
+87
View File
@@ -0,0 +1,87 @@
"""Row-based splitting for spreadsheet documents."""
from __future__ import annotations
from rag_cut.models import Block, BlockType, SplitConfig
from rag_cut.parsers.pdf.tables import (
build_table_embedding_text,
detect_spreadsheet_layout,
extract_table_keywords,
rows_to_markdown,
)
def _chunk_meta(table_block: Block, chunk_rows: list[list[str]], md: str, row_range: list[int]) -> dict:
"""Build per-chunk metadata without leaking the full source table."""
header_rows = table_block.meta.get("header_rows") or 1
keywords = extract_table_keywords(
table_block.meta.get("table_title") or "",
md,
table_block.meta.get("table_description") or "",
)
embedding_text = build_table_embedding_text(
table_title=table_block.meta.get("table_title") or "",
markdown=md,
description=table_block.meta.get("table_description") or "",
keywords=keywords,
)
return {
"sheet": table_block.meta.get("sheet"),
"table_title": table_block.meta.get("table_title"),
"table_description": table_block.meta.get("table_description"),
"header_rows": header_rows,
"header_row_start": table_block.meta.get("header_row_start"),
"header_row_end": table_block.meta.get("header_row_end"),
"data_start_row": table_block.meta.get("data_start_row"),
"preamble_rows": table_block.meta.get("preamble_rows", 0),
"row_count": len(chunk_rows),
"col_count": max((len(r) for r in chunk_rows), default=0),
"chunk_rows": chunk_rows,
"row_range": row_range,
"keywords": keywords,
"embedding_text": embedding_text,
}
def split_by_row(blocks: list[Block], config: SplitConfig) -> list[list[Block]]:
table_block = next((b for b in blocks if b.type == BlockType.TABLE and b.meta.get("rows")), None)
if not table_block:
return [blocks] if blocks else []
rows: list[list[str]] = table_block.meta["rows"]
layout = {
"header_row_start": table_block.meta.get("header_row_start"),
"header_row_end": table_block.meta.get("header_row_end"),
"data_start_row": table_block.meta.get("data_start_row"),
"header_rows": table_block.meta.get("header_rows"),
}
if not layout["header_row_start"]:
layout = detect_spreadsheet_layout(rows)
h_start = max(1, config.header_row_start) - 1
h_end = max(h_start + 1, config.header_row_end)
header = rows[h_start:h_end]
header_rows = layout.get("header_rows") or len(header)
data_start = max(config.start_row - 1, h_end)
data_rows = rows[data_start:]
rows_per = max(1, config.rows_per_chunk)
groups: list[list[Block]] = []
for i in range(0, len(data_rows), rows_per):
slice_rows = data_rows[i : i + rows_per]
chunk_rows = header + slice_rows
md = rows_to_markdown(chunk_rows, header_rows=header_rows)
row_range = [
data_start + i + 1,
data_start + i + len(slice_rows),
]
groups.append(
[
Block(
type=BlockType.TABLE,
markdown=md,
meta=_chunk_meta(table_block, chunk_rows, md, row_range),
)
]
)
return groups
@@ -0,0 +1,210 @@
"""Default structure-aware splitting."""
from __future__ import annotations
from rag_cut.models import Block, BlockType, SplitConfig
from rag_cut.splitters.heading_splitter import (
chunk_groups_to_block_groups,
has_meaningful_headings,
split_by_heading_hierarchy,
)
ATOMIC_TYPES = {BlockType.TABLE, BlockType.IMAGE}
def _same_layout_context(prev: Block, curr: Block) -> bool:
"""True when blocks should stay together to preserve image/text position."""
if prev.meta.get("page") != curr.meta.get("page"):
return False
if prev.type == BlockType.PARAGRAPH and curr.type in {BlockType.IMAGE, BlockType.TABLE}:
return True
if prev.type == BlockType.HEADING and curr.type in {BlockType.PARAGRAPH, BlockType.IMAGE, BlockType.TABLE}:
return prev.meta.get("parent_heading") == curr.meta.get("parent_heading") or not prev.meta.get("parent_heading")
if prev.type == BlockType.IMAGE and curr.type == BlockType.PARAGRAPH:
return True
if prev.type == BlockType.TABLE and curr.type == BlockType.PARAGRAPH:
return True
if prev.type == BlockType.IMAGE and curr.type == BlockType.IMAGE:
return True
return False
def _is_heading(block: Block) -> bool:
return block.type == BlockType.HEADING
def _meaningful_headings(blocks: list[Block]) -> bool:
headings = [b for b in blocks if _is_heading(b)]
if len(headings) < 2:
return False
substantial = [h for h in headings if len((h.text or "").strip()) >= 8]
return len(substantial) >= 2
def _split_by_headings(blocks: list[Block]) -> list[list[Block]]:
"""Outline-aware: split when a heading of same-or-higher level appears."""
if not any(_is_heading(b) for b in blocks):
return []
sections: list[list[Block]] = []
current: list[Block] = []
stack: list[int] = []
for block in blocks:
if _is_heading(block):
level = block.level or 1
while stack and stack[-1] >= level:
stack.pop()
if current:
sections.append(current)
current = []
stack.append(level)
current.append(block)
else:
current.append(block)
if current:
sections.append(current)
return sections
def _split_by_page(blocks: list[Block], config: SplitConfig) -> list[list[Block]]:
"""Prefer page boundaries for PDF/manual style documents."""
if not any(b.meta.get("page") for b in blocks):
return []
groups: list[list[Block]] = []
current: list[Block] = []
current_page: int | None = None
for block in blocks:
page = block.meta.get("page")
if current and page != current_page:
groups.append(current)
current = []
current_page = page
current.append(block)
if current:
groups.append(current)
return _merge_oversized_sections(groups, config)
def _split_by_size(blocks: list[Block], config: SplitConfig) -> list[list[Block]]:
"""Fallback: pack blocks up to max_chunk_size without splitting atomic blocks."""
groups: list[list[Block]] = []
current: list[Block] = []
current_len = 0
def flush() -> None:
nonlocal current, current_len
if current:
groups.append(current)
current = []
current_len = 0
for block in blocks:
rendered = block.render()
block_len = len(rendered) + 2
page = block.meta.get("page")
if block.type in ATOMIC_TYPES and current_len + block_len > config.max_chunk_size and current:
if not _same_layout_context(current[-1], block):
flush()
if block.type not in ATOMIC_TYPES and block_len > config.max_chunk_size:
if current:
flush()
text = block.text or block.markdown
start = 0
while start < len(text):
end = min(start + config.max_chunk_size, len(text))
piece = Block(type=block.type, text=text[start:end], level=block.level, meta=block.meta)
groups.append([piece])
if end >= len(text):
break
start = max(end - config.overlap, start + 1)
continue
if current_len + block_len > config.max_chunk_size and current:
# Keep image with preceding heading/body on the same page
if _same_layout_context(current[-1], block):
pass
else:
flush()
elif (
current
and page is not None
and current[-1].meta.get("page") != page
and current_len >= min(400, config.max_chunk_size // 3)
):
flush()
current.append(block)
current_len += block_len
flush()
return groups
def _merge_oversized_sections(sections: list[list[Block]], config: SplitConfig) -> list[list[Block]]:
result: list[list[Block]] = []
for section in sections:
rendered_len = sum(len(b.render()) + 2 for b in section)
if rendered_len <= config.max_chunk_size:
result.append(section)
else:
result.extend(_split_by_size(section, config))
return result
def split_default(blocks: list[Block], config: SplitConfig) -> list[list[Block]]:
if not blocks:
return []
# Spreadsheet: single table block — default = chunk by groups of rows with header
if len(blocks) == 1 and blocks[0].type == BlockType.TABLE and blocks[0].meta.get("rows"):
from rag_cut.splitters.by_row import split_by_row
row_config = SplitConfig(
mode=config.mode,
header_row_start=1,
header_row_end=1,
start_row=2,
rows_per_chunk=max(1, min(10, len(blocks[0].meta["rows"]) // 5 or 1)),
max_chunk_size=config.max_chunk_size,
overlap=config.overlap,
)
return split_by_row(blocks, row_config)
# Heading hierarchy first: same section keeps body/images/tables/captions together
if has_meaningful_headings(blocks):
heading_groups = split_by_heading_hierarchy(blocks, config)
if heading_groups:
block_groups, _ = chunk_groups_to_block_groups(heading_groups)
return block_groups
page_sections = _split_by_page(blocks, config)
if len(page_sections) > 1:
return page_sections
return _split_by_size(blocks, config)
def split_default_with_meta(blocks: list[Block], config: SplitConfig) -> tuple[list[list[Block]], list[dict]]:
"""Like split_default but also returns per-group metadata (heading sections)."""
if not blocks:
return [], []
if len(blocks) == 1 and blocks[0].type == BlockType.TABLE and blocks[0].meta.get("rows"):
groups = split_default(blocks, config)
return groups, [{} for _ in groups]
if has_meaningful_headings(blocks):
heading_groups = split_by_heading_hierarchy(blocks, config)
if heading_groups:
return chunk_groups_to_block_groups(heading_groups)
groups = split_default(blocks, config)
return groups, [{} for _ in groups]
+63
View File
@@ -0,0 +1,63 @@
"""Delimiter-based splitting for non-tabular documents."""
from __future__ import annotations
from rag_cut.models import Block, BlockType, SplitConfig
from rag_cut.splitters.default_splitter import _split_by_size
ATOMIC_TYPES = {BlockType.TABLE, BlockType.IMAGE}
def partition_blocks_by_delimiter(blocks: list[Block], delimiter: str) -> list[list[Block]]:
"""Split block stream on delimiter; delimiter text is discarded from chunks."""
if not delimiter:
return [blocks] if blocks else []
groups: list[list[Block]] = []
current: list[Block] = []
def flush() -> None:
nonlocal current
if current:
groups.append(current)
current = []
for block in blocks:
if block.type in ATOMIC_TYPES:
current.append(block)
continue
text = block.text or block.markdown
if delimiter not in text:
current.append(block)
continue
parts = text.split(delimiter)
for i, part in enumerate(parts):
part = part.strip()
if part:
piece = Block(type=block.type, text=part, level=block.level, meta=dict(block.meta))
current.append(piece)
if i < len(parts) - 1:
flush()
flush()
return groups
def split_by_delimiter(blocks: list[Block], config: SplitConfig) -> list[list[Block]]:
if not config.delimiter:
raise ValueError("delimiter is required for delimiter split mode")
groups = partition_blocks_by_delimiter(blocks, config.delimiter)
if not groups:
return _split_by_size(blocks, config)
sized: list[list[Block]] = []
for group in groups:
rendered_len = sum(len(b.render()) + 2 for b in group)
if rendered_len <= config.max_chunk_size:
sized.append(group)
else:
sized.extend(_split_by_size(group, config))
return sized
@@ -0,0 +1,344 @@
"""Heading-hierarchy-first document splitting."""
from __future__ import annotations
import re
from dataclasses import dataclass, field
from rag_cut.models import Block, BlockType, SplitConfig
# e.g. "1.2 ACCOUNT STATUS CODE MASTER", "10.上升三角形態" (space after '.' optional)
NUMBERED_HEADING_RE = re.compile(
r"^\s*(\d+(?:\.\d+)*)(?:\.|.)?\s*([A-Za-z0-9\u4e00-\u9fff][A-Za-z0-9\u4e00-\u9fff&/ \-_::]{2,})\s*$"
)
FIGURE_TABLE_RE = re.compile(
r"^\s*(?:Figure|Fig\.|图|表|Table)\s*[\d.]+",
re.I,
)
STEP_RE = re.compile(
r"^\s*(?:Step\s*\d+|步骤\s*\d+|\d{1,2}[..、]\s*(?:点击|點擊|选择|選擇|输入|輸入))",
re.I,
)
ATOMIC_TYPES = {BlockType.TABLE, BlockType.IMAGE}
MAX_HEADING_CHARS = 100
@dataclass
class ChunkGroup:
blocks: list[Block]
meta: dict = field(default_factory=dict)
def _text(block: Block) -> str:
return (block.text or block.markdown or "").strip()
def _rendered_len(blocks: list[Block]) -> int:
return sum(len(b.render()) + 2 for b in blocks)
def numbered_heading_level(text: str) -> int | None:
match = NUMBERED_HEADING_RE.match(text.strip())
if not match:
return None
return match.group(1).count(".") + 1
def infer_heading_level(block: Block) -> int:
if block.type == BlockType.HEADING and block.level:
numbered = numbered_heading_level(block.text)
if numbered:
return numbered
return block.level
numbered = numbered_heading_level(_text(block))
if numbered:
return numbered
return block.level or 1
def is_heading_block(block: Block) -> bool:
text = _text(block)
if not text or len(text) > MAX_HEADING_CHARS:
return False
if block.type == BlockType.HEADING:
return True
if NUMBERED_HEADING_RE.match(text):
return True
if block.meta.get("font_size") and block.meta.get("body_font_size"):
return block.meta["font_size"] >= block.meta["body_font_size"] + 1.5
return False
def normalize_heading_block(block: Block) -> Block:
text = _text(block)
if block.type == BlockType.HEADING and len(text) > MAX_HEADING_CHARS:
# Parser sometimes merges title+body then marks the blob as heading.
return Block(type=BlockType.PARAGRAPH, text=text, meta=dict(block.meta))
numbered = numbered_heading_level(text)
if block.type == BlockType.HEADING:
level = numbered or block.level or 1
return block.model_copy(update={"level": level})
if numbered and NUMBERED_HEADING_RE.match(text):
return Block(
type=BlockType.HEADING,
text=text,
level=numbered,
meta=dict(block.meta),
)
return block
def has_meaningful_headings(blocks: list[Block]) -> bool:
headings = [normalize_heading_block(b) for b in blocks]
count = sum(1 for b in headings if is_heading_block(b) or b.type == BlockType.HEADING)
return count >= 2
def _section_key(block: Block | None, fallback: int) -> str:
if block is None:
return f"section-{fallback}"
oi = block.meta.get("order_index", fallback)
title = re.sub(r"\W+", "-", (_text(block) or "heading"))[:48]
return f"h-{oi}-{title}"
@dataclass
class _OpenSection:
level: int
start_index: int
blocks: list[Block]
def _split_primary_sections(blocks: list[Block]) -> list[list[Block]]:
"""
Split so each heading owns its content until the next same-or-higher-level heading.
Example: 1.2 section runs until 1.3 (same level) or 2.0 (higher level).
Nested sub-headings (1.2.1) stay inside the 1.2 section.
"""
if not blocks:
return []
open_sections: list[_OpenSection] = []
finished: list[tuple[int, list[Block]]] = []
for raw in blocks:
block = normalize_heading_block(raw)
if is_heading_block(block) or block.type == BlockType.HEADING:
level = infer_heading_level(block)
while open_sections and open_sections[-1].level >= level:
sec = open_sections.pop()
finished.append((sec.start_index, sec.blocks))
start = int(block.meta.get("order_index", len(finished)))
open_sections.append(_OpenSection(level=level, start_index=start, blocks=[block]))
elif open_sections:
open_sections[-1].blocks.append(block)
while open_sections:
sec = open_sections.pop()
finished.append((sec.start_index, sec.blocks))
finished.sort(key=lambda item: item[0])
return [sec_blocks for _, sec_blocks in finished]
def _section_heading(section: list[Block]) -> Block | None:
for block in section:
nb = normalize_heading_block(block)
if nb.type == BlockType.HEADING or is_heading_block(nb):
return nb
return None
def _split_by_child_headings(section: list[Block], parent_level: int) -> list[list[Block]]:
"""Split an oversized section by deeper sub-headings."""
child_sections: list[list[Block]] = []
current: list[Block] = []
parent_heading = _section_heading(section)
for block in section:
nb = normalize_heading_block(block)
if (
block is not parent_heading
and (nb.type == BlockType.HEADING or is_heading_block(nb))
and infer_heading_level(nb) > parent_level
):
if current:
child_sections.append(current)
current = [block]
else:
current.append(block)
if current:
child_sections.append(current)
return child_sections if len(child_sections) > 1 else [section]
def _is_split_marker(block: Block) -> bool:
if block.type in ATOMIC_TYPES:
return False
text = _text(block)
if not text:
return False
return bool(FIGURE_TABLE_RE.match(text) or STEP_RE.match(text))
def _split_by_content_markers(section: list[Block], config: SplitConfig) -> list[list[Block]]:
"""Fallback: split at figure/table/step markers while keeping atomic blocks intact."""
if _rendered_len(section) <= config.max_chunk_size:
return [section]
groups: list[list[Block]] = []
current: list[Block] = []
current_len = 0
def flush() -> None:
nonlocal current, current_len
if current:
groups.append(current)
current = []
current_len = 0
for block in section:
blen = len(block.render()) + 2
if (
current
and _is_split_marker(block)
and current_len >= min(500, config.max_chunk_size // 4)
and current_len + blen > config.max_chunk_size
):
flush()
elif current_len + blen > config.max_chunk_size and current:
if block.type in ATOMIC_TYPES:
flush()
elif not _is_split_marker(block):
flush()
current.append(block)
current_len += blen
flush()
return groups if groups else [section]
def _prepend_parent_heading(section: list[Block], parent: Block | None) -> list[Block]:
if not parent:
return section
parent_text = _text(parent)
if section and _text(normalize_heading_block(section[0])) == parent_text:
return section
return [parent] + section
def _split_oversized_section(
section: list[Block],
config: SplitConfig,
section_id: str,
) -> list[ChunkGroup]:
"""Split an oversized heading section by sub-headings/markers/size — no parent chunk."""
parent_heading = _section_heading(section)
parent_level = infer_heading_level(parent_heading) if parent_heading else 1
parent_title = _text(parent_heading) if parent_heading else ""
child_sections = _split_by_child_headings(section, parent_level)
if len(child_sections) == 1:
child_sections = _split_by_content_markers(section, config)
if len(child_sections) == 1:
from rag_cut.splitters.default_splitter import _split_by_size
child_sections = _split_by_size(section, config)
parts: list[ChunkGroup] = []
for idx, child in enumerate(child_sections):
child_heading = _section_heading(child)
blocks = _prepend_parent_heading(child, parent_heading)
parts.append(
ChunkGroup(
blocks=blocks,
meta={
"section_id": f"{section_id}-part-{idx}",
"heading": _text(child_heading) if child_heading else parent_title,
"heading_level": infer_heading_level(child_heading) if child_heading else parent_level,
"parent_heading": parent_title,
"chunk_strategy": "heading_hierarchy_part",
"part_index": idx,
"is_sub_chunk": False,
"retrieval": True,
},
)
)
return parts
def split_by_heading_hierarchy(blocks: list[Block], config: SplitConfig) -> list[ChunkGroup]:
"""
Heading-first splitting:
- Same heading section stays together (body, images, tables, captions).
- Boundaries at same/higher-level headings.
- Oversized sections are split by sub-headings / markers / length (no parent+child pair).
"""
if not blocks:
return []
if not has_meaningful_headings(blocks):
return []
groups: list[ChunkGroup] = []
sections = _split_primary_sections(blocks)
for i, section in enumerate(sections):
heading = _section_heading(section)
sid = _section_key(heading, i)
title = _text(heading) if heading else ""
level = infer_heading_level(heading) if heading else 1
if _rendered_len(section) <= config.max_chunk_size:
groups.append(
ChunkGroup(
blocks=section,
meta={
"section_id": sid,
"heading": title,
"heading_level": level,
"is_sub_chunk": False,
"chunk_strategy": "heading_hierarchy",
"retrieval": True,
"order_range": [
section[0].meta.get("order_index"),
section[-1].meta.get("order_index"),
],
},
)
)
continue
groups.extend(_split_oversized_section(section, config, sid))
return groups
def chunk_groups_to_block_groups(groups: list[ChunkGroup]) -> tuple[list[list[Block]], list[dict]]:
"""Convert ChunkGroups to block groups + per-chunk meta for renderer."""
block_groups: list[list[Block]] = []
metas: list[dict] = []
for group in groups:
block_groups.append(group.blocks)
metas.append(group.meta)
return block_groups, metas
def assign_parent_chunk_ids(chunks: list) -> list:
"""Resolve parent_section_id -> parent_chunk_id (chunk index)."""
section_index: dict[str, int] = {}
for i, chunk in enumerate(chunks):
sid = chunk.meta.get("section_id")
if sid and chunk.meta.get("is_section_parent"):
section_index[sid] = i
updated = []
for chunk in chunks:
meta = dict(chunk.meta)
parent_sid = meta.get("parent_section_id")
if parent_sid and parent_sid in section_index:
meta["parent_chunk_id"] = section_index[parent_sid]
updated.append(chunk.model_copy(update={"meta": meta}))
return updated
+103
View File
@@ -0,0 +1,103 @@
"""Parent/child delimiter splitting for fine retrieval + coarse recall."""
from __future__ import annotations
from rag_cut.models import Block, SplitConfig
from rag_cut.splitters.default_splitter import _split_by_size
from rag_cut.splitters.delimiter import partition_blocks_by_delimiter
from rag_cut.splitters.heading_splitter import ChunkGroup
CHUNK_STRATEGY = "parent_child_delimiter"
CHILD_MAX_HARD_LIMIT = 1500
def _rendered_len(blocks: list[Block]) -> int:
return sum(len(b.render()) + 2 for b in blocks)
def _validate(config: SplitConfig) -> tuple[str, str | None, int, int]:
parent_delimiter = (config.parent_delimiter or config.delimiter or "").strip()
if not parent_delimiter:
raise ValueError("parent_delimiter is required for parent_child split mode")
child_delimiter = (config.child_delimiter or "").strip() or None
parent_max = max(200, int(config.max_chunk_size or 1500))
child_max = int(config.child_max_size or 512)
child_max = max(50, min(child_max, CHILD_MAX_HARD_LIMIT, parent_max))
return parent_delimiter, child_delimiter, parent_max, child_max
def _size_cap(groups: list[list[Block]], max_size: int, overlap: int) -> list[list[Block]]:
sized: list[list[Block]] = []
size_config = SplitConfig(max_chunk_size=max_size, overlap=overlap)
for group in groups:
if _rendered_len(group) <= max_size:
sized.append(group)
else:
sized.extend(_split_by_size(group, size_config))
return sized or groups
def split_by_parent_child(blocks: list[Block], config: SplitConfig) -> list[ChunkGroup]:
"""
Split into parent chunks (context) and child chunks (retrieval).
1. Partition by parent_delimiter (then cap by parent max length).
2. For each parent: keep a full parent chunk (retrieval=false).
3. Partition parent by child_delimiter (or by length) into children
capped by child_max_size (retrieval=true, linked via parent_section_id).
"""
if not blocks:
return []
parent_delimiter, child_delimiter, parent_max, child_max = _validate(config)
overlap = max(0, int(config.overlap or 0))
parents = partition_blocks_by_delimiter(blocks, parent_delimiter)
if not parents:
parents = [blocks]
parents = _size_cap(parents, parent_max, overlap)
groups: list[ChunkGroup] = []
for parent_idx, parent_blocks in enumerate(parents):
section_id = f"pc-{parent_idx}"
groups.append(
ChunkGroup(
blocks=list(parent_blocks),
meta={
"section_id": section_id,
"is_section_parent": True,
"is_sub_chunk": False,
"retrieval": False,
"chunk_strategy": f"{CHUNK_STRATEGY}_parent",
"parent_index": parent_idx,
},
)
)
if child_delimiter:
children = partition_blocks_by_delimiter(parent_blocks, child_delimiter)
else:
children = [parent_blocks]
if not children:
children = [parent_blocks]
children = _size_cap(children, child_max, overlap)
for child_idx, child_blocks in enumerate(children):
groups.append(
ChunkGroup(
blocks=list(child_blocks),
meta={
"section_id": f"{section_id}-sub-{child_idx}",
"parent_section_id": section_id,
"is_section_parent": False,
"is_sub_chunk": True,
"retrieval": True,
"chunk_strategy": f"{CHUNK_STRATEGY}_child",
"parent_index": parent_idx,
"sub_chunk_index": child_idx,
},
)
)
return groups
+427
View File
@@ -0,0 +1,427 @@
"""Generic heading/layout multimodal splitter for PDF-derived documents."""
from __future__ import annotations
import re
from dataclasses import dataclass, field
from rag_cut.models import Block, BlockType, Chunk, SplitConfig
from rag_cut.parsers.pdf.noise_filter import is_toc_noise_text, is_toc_title_text
from rag_cut.renderer import collect_chunk_layout_meta, render_blocks
CHUNK_STRATEGY = "heading_layout_multimodal"
PAGE_NUMBER_RE = re.compile(r"^\s*(?:[-\u2013\u2014]?\s*)?\d{1,4}(?:\s*/\s*\d{1,4})?\s*$")
NUMBERED_HEADING_RE = re.compile(
r"^\s*(?P<num>\d+(?:\.\d+)*)(?:\.|.)?\s*(?P<title>[A-Za-z0-9\u4e00-\u9fff][^\n]{1,120})\s*$"
)
LETTER_HEADING_RE = re.compile(
r"^\s*(?P<letter>[A-Z])[\.)]\s+(?P<title>[A-Za-z0-9\u4e00-\u9fff][^\n]{1,100})\s*$"
)
# Procedural "Step N …" lines are body content, not section boundaries.
# Promoting them to headings caused heading-only groups to be dropped on flush
# (e.g. Steps 1–3 under 5.7.1 vanished while only Step 4 with following body survived).
STEP_INSTRUCTION_RE = re.compile(
r"^\s*(?:Step\s*\d+|STEP\s*\d+|步骤\s*\d+)(?:\s*[:.:)\-]?\s+\S|\s*$)",
re.I,
)
CHAPTER_HEADING_RE = re.compile(r"^\s*Chapter\s+\d+(?:\s*[:.-]?\s+[^\n]{1,100})?\s*$", re.I)
CN_HEADING_RE = re.compile(
r"^\s*(?:[\u4e00-\u9fff]{1,3}[、..]|[((][\u4e00-\u9fff]{1,3}[))])\s*[^\n]{1,100}$"
)
@dataclass
class _Group:
heading_path: list[str] = field(default_factory=list)
blocks: list[Block] = field(default_factory=list)
@dataclass
class _HeadingSignal:
text: str
level: int
def _text(block: Block) -> str:
return (block.text or block.markdown or "").strip()
def _page(block: Block) -> int | None:
page = block.meta.get("page")
try:
return int(page) if page is not None else None
except (TypeError, ValueError):
return None
def _bbox(block: Block) -> list[float]:
bbox = block.meta.get("bbox")
return list(bbox) if isinstance(bbox, (list, tuple)) else []
def _copy_with_meta(block: Block, **meta_updates) -> Block:
meta = dict(block.meta)
meta.update({k: v for k, v in meta_updates.items() if v not in (None, [], {})})
return block.model_copy(update={"meta": meta})
def _is_empty_text_block(block: Block) -> bool:
return block.type not in {BlockType.IMAGE, BlockType.TABLE} and not _text(block)
def _is_toc_text(text: str) -> bool:
return is_toc_noise_text(text) or is_toc_title_text(text)
def _is_decorative_image(block: Block) -> bool:
bbox = _bbox(block)
if len(bbox) != 4:
return False
width = bbox[2] - bbox[0]
height = bbox[3] - bbox[1]
if width <= 0 or height <= 0:
return True
if width < 28 or height < 16:
return True
page_height = block.meta.get("page_height") or 0
if page_height and bbox[3] <= page_height * 0.10 and width * height < 12000:
return True
return False
def _is_margin_noise(block: Block) -> bool:
if block.type == BlockType.HEADING:
return False
text = _text(block)
bbox = _bbox(block)
page_height = block.meta.get("page_height") or 0
if not bbox or not page_height:
return False
y0, y1 = bbox[1], bbox[3]
in_top = y1 <= page_height * 0.08
in_bottom = y0 >= page_height * 0.92
if PAGE_NUMBER_RE.match(text) and (in_top or in_bottom):
return True
return len(text) <= 80 and (in_top or in_bottom) and block.meta.get("running_header")
def _running_header_texts(blocks: list[Block]) -> set[str]:
page_count = len({_page(b) for b in blocks if _page(b) is not None})
if page_count < 3:
return set()
locations: dict[str, set[int]] = {}
for block in blocks:
if block.type in {BlockType.IMAGE, BlockType.TABLE, BlockType.HEADING}:
continue
text = " ".join(_text(block).split())
page = _page(block)
bbox = _bbox(block)
page_height = block.meta.get("page_height") or 0
if not text or page is None or not bbox or not page_height or len(text) > 100:
continue
if bbox[3] <= page_height * 0.10 or bbox[1] >= page_height * 0.90:
locations.setdefault(text, set()).add(page)
threshold = max(3, int(page_count * 0.5))
return {text for text, pages in locations.items() if len(pages) >= threshold}
def _is_noise(block: Block, running_headers: set[str]) -> bool:
if _is_empty_text_block(block):
return True
text = " ".join(_text(block).split())
if block.type == BlockType.IMAGE:
return _is_decorative_image(block)
if PAGE_NUMBER_RE.match(text):
return True
if _is_toc_text(text):
return True
if text in running_headers:
return True
return _is_margin_noise(block)
def _heading_signal(block: Block, current_top_level: bool = False) -> _HeadingSignal | None:
text = _text(block)
if not text or len(text) > 180 or _is_toc_text(text):
return None
# Keep procedural steps inside the parent section; do not open a new group.
if STEP_INSTRUCTION_RE.match(text):
return None
if block.type == BlockType.HEADING:
level = block.level or 1
numbered = NUMBERED_HEADING_RE.match(text)
if numbered:
level = numbered.group("num").count(".") + 1
elif LETTER_HEADING_RE.match(text):
level = 2 if current_top_level else max(2, level)
return _HeadingSignal(text=" ".join(text.split()), level=max(1, min(level, 6)))
numbered = NUMBERED_HEADING_RE.match(text)
if numbered:
return _HeadingSignal(text=" ".join(text.split()), level=numbered.group("num").count(".") + 1)
if LETTER_HEADING_RE.match(text):
return _HeadingSignal(text=" ".join(text.split()), level=2)
if CHAPTER_HEADING_RE.match(text):
return _HeadingSignal(text=" ".join(text.split()), level=1)
if CN_HEADING_RE.match(text):
return _HeadingSignal(text=" ".join(text.split()), level=2 if current_top_level else 1)
font_size = block.meta.get("font_size")
body_size = block.meta.get("body_font_size")
if font_size and body_size and font_size >= body_size + 1.5 and len(text) <= 100:
level = block.level or 2
# Short CJK section banners (形態指標 / 策略指標) are chapter peers, not
# subsections of the preceding numbered person/indicator entry.
if (
font_size >= body_size + 3
and re.fullmatch(r"[\u4e00-\u9fffA-Za-z0-9//\s]{2,24}", text)
and not NUMBERED_HEADING_RE.match(text)
):
level = 1
return _HeadingSignal(text=" ".join(text.split()), level=max(1, min(level, 6)))
return None
def _heading_block(block: Block, signal: _HeadingSignal) -> Block:
meta = dict(block.meta)
meta["section_boundary"] = True
meta["heading_level"] = signal.level
return Block(type=BlockType.HEADING, text=signal.text, level=signal.level, meta=meta)
def _enrich_block(block: Block, heading_path: list[str], group_blocks: list[Block]) -> Block:
meta = dict(block.meta)
meta["heading_path"] = list(heading_path)
if heading_path:
meta["nearest_heading"] = heading_path[-1]
meta["parent_heading"] = heading_path[-1]
meta["chapter"] = heading_path[0]
meta.setdefault("section_boundary", block.type == BlockType.HEADING)
if block.type == BlockType.HEADING:
meta["heading_level"] = block.level or len(heading_path) or 1
if block.type in {BlockType.IMAGE, BlockType.TABLE}:
meta.setdefault("bound_heading", heading_path[-1] if heading_path else "")
_bind_adjacent_text(meta, block, group_blocks)
return block.model_copy(update={"meta": meta})
def _bind_adjacent_text(meta: dict, block: Block, group_blocks: list[Block]) -> None:
page = _page(block)
for prev in reversed(group_blocks):
if prev.type == BlockType.PARAGRAPH and _page(prev) == page and _text(prev):
meta.setdefault("preceding_text", _text(prev)[:400])
break
if prev.type == BlockType.HEADING:
break
if block.type == BlockType.TABLE and not meta.get("table_title"):
for prev in reversed(group_blocks):
if prev.type == BlockType.HEADING:
meta["table_title"] = _text(prev)
break
if prev.type == BlockType.PARAGRAPH and _page(prev) == page and 0 < len(_text(prev)) <= 120:
meta["table_title"] = _text(prev)
break
def _fill_following_text(group: _Group) -> _Group:
blocks = list(group.blocks)
for i, block in enumerate(blocks):
if block.type not in {BlockType.IMAGE, BlockType.TABLE}:
continue
meta = dict(block.meta)
page = _page(block)
for nxt in blocks[i + 1 :]:
if nxt.type == BlockType.HEADING:
break
if nxt.type == BlockType.PARAGRAPH and _page(nxt) == page and _text(nxt):
meta.setdefault("following_text", _text(nxt)[:400])
break
blocks[i] = block.model_copy(update={"meta": meta})
return _Group(heading_path=group.heading_path, blocks=blocks)
def _group_has_body(group: _Group) -> bool:
for block in group.blocks:
if block.type in {BlockType.IMAGE, BlockType.TABLE}:
return True
if block.type != BlockType.HEADING and _text(block):
return True
# Safety net: long instructional "headings" are themselves searchable content.
return any(block.type == BlockType.HEADING and len(_text(block)) >= 48 for block in group.blocks)
def _attach_heading_only_to_previous(groups: list[_Group], orphan: _Group) -> None:
"""Keep catalog-style heading-only groups instead of deleting them.
Example: after \"共包含以下11個形態指標說明:\", lines like \"1.頭肩頂形態\"
briefly open a group with no body before the next number arrives. Flush used
to drop them; fold those leaf titles into the previous section as paragraphs.
"""
if not groups or not orphan.blocks:
return
prev = groups[-1]
prev_keys = {(b.type, _text(b), tuple(b.meta.get("bbox") or [])) for b in prev.blocks}
path_set = set(prev.heading_path)
for block in orphan.blocks:
text = _text(block)
if not text:
continue
# Skip copies of ancestor headings already injected into the orphan group.
if text in path_set:
continue
key = (block.type, text, tuple(block.meta.get("bbox") or []))
if key in prev_keys:
continue
if block.type == BlockType.HEADING:
meta = dict(block.meta)
meta["catalog_item"] = True
meta["section_boundary"] = False
demoted = Block(type=BlockType.PARAGRAPH, text=text, meta=meta)
prev.blocks.append(_enrich_block(demoted, prev.heading_path, prev.blocks))
else:
prev.blocks.append(_enrich_block(block, prev.heading_path, prev.blocks))
prev_keys.add(key)
def _group_is_retrievable(group: _Group) -> bool:
"""Return whether a rendered group contains searchable semantic context."""
if group.heading_path:
return True
for block in group.blocks:
if block.type == BlockType.TABLE:
return True
if block.type == BlockType.IMAGE:
if block.ocr_text.strip():
return True
continue
if block.type != BlockType.HEADING and _text(block):
return True
return False
def _render_group(group: _Group, index: int, config: SplitConfig, meta: dict | None = None) -> Chunk:
group = _fill_following_text(group)
layout_meta = collect_chunk_layout_meta(group.blocks)
heading_path = group.heading_path or layout_meta.get("heading_path") or []
heading = heading_path[-1] if heading_path else layout_meta.get("heading") or "Untitled section"
extra_meta = {
**layout_meta,
"chunk_strategy": CHUNK_STRATEGY,
"heading": heading,
"heading_path": heading_path,
"retrieval": _group_is_retrievable(group),
**(meta or {}),
}
if "order_range" not in extra_meta:
order_indices = [b.meta.get("order_index") for b in group.blocks if b.meta.get("order_index") is not None]
if order_indices:
extra_meta["order_range"] = [min(order_indices), max(order_indices)]
chunk = render_blocks(group.blocks, index=index, meta=extra_meta)
chunk_meta = dict(chunk.meta)
chunk_meta.update(extra_meta)
return chunk.model_copy(update={"meta": chunk_meta})
def _split_oversized_group(group: _Group, start_index: int, config: SplitConfig) -> list[Chunk]:
"""Split an oversized section by length only (no parent+child pair)."""
rendered_len = sum(len(b.render()) + 2 for b in group.blocks)
if rendered_len <= config.max_chunk_size or len(group.blocks) <= 3:
return [_render_group(group, start_index, config)]
heading_blocks = [b for b in group.blocks if b.type == BlockType.HEADING]
prefix = heading_blocks[:1]
chunks: list[Chunk] = []
current: list[Block] = list(prefix)
current_len = sum(len(b.render()) + 2 for b in current)
part_index = 0
def flush() -> None:
nonlocal current, current_len, part_index
body = [b for b in current if b not in prefix]
if not body:
return
part_group = _Group(group.heading_path, list(current))
chunks.append(
_render_group(
part_group,
start_index + len(chunks),
config,
meta={
"chunk_strategy": f"{CHUNK_STRATEGY}_part",
"part_index": part_index,
"retrieval": True,
},
)
)
part_index += 1
current = list(prefix)
current_len = sum(len(b.render()) + 2 for b in current)
for block in group.blocks[len(prefix) :]:
block_len = len(block.render()) + 2
if current_len + block_len > config.max_chunk_size and len(current) > len(prefix):
flush()
current.append(block)
current_len += block_len
flush()
return chunks or [_render_group(group, start_index, config)]
def _build_groups(blocks: list[Block]) -> list[_Group]:
running_headers = _running_header_texts(blocks)
heading_stack: list[tuple[int, str, Block]] = []
groups: list[_Group] = []
current = _Group()
def flush() -> None:
nonlocal current
if _group_has_body(current):
groups.append(current)
elif current.blocks:
# Explicit parser headings with no body are empty document sections,
# not catalog items to fold into the preceding section.
if not any(block.meta.get("source_heading") for block in current.blocks):
_attach_heading_only_to_previous(groups, current)
current = _Group()
for raw in blocks:
if _is_noise(raw, running_headers):
continue
signal = _heading_signal(raw, current_top_level=bool(heading_stack))
if signal:
flush()
while heading_stack and heading_stack[-1][0] >= signal.level:
heading_stack.pop()
heading = _heading_block(raw, signal)
heading_stack.append((signal.level, signal.text, heading))
path = [item[1] for item in heading_stack]
current = _Group(heading_path=path)
for _, _, h_block in heading_stack:
current.blocks.append(_enrich_block(h_block, path, current.blocks))
continue
path = [item[1] for item in heading_stack]
if not current.blocks and heading_stack:
current.heading_path = path
for _, _, h_block in heading_stack:
current.blocks.append(_enrich_block(h_block, path, current.blocks))
elif not current.heading_path:
current.heading_path = path
current.blocks.append(_enrich_block(raw, current.heading_path, current.blocks))
flush()
return groups
def split_pdf_semantic(blocks: list[Block], config: SplitConfig) -> list[Chunk]:
"""Split PDF/Word-derived blocks into heading/layout-preserving multimodal chunks."""
groups = _build_groups(blocks)
chunks: list[Chunk] = []
for group in groups:
for chunk in _split_oversized_group(group, len(chunks), config):
chunks.append(chunk.model_copy(update={"index": len(chunks)}))
return chunks
+82
View File
@@ -0,0 +1,82 @@
"""Choose a PDF chunking strategy from parsed document signals."""
from __future__ import annotations
import re
from pathlib import Path
from rag_cut.models import Block, BlockType
FEATURE_TITLE_RE = re.compile(r"^\s*\d{1,2}[..、]\s*[A-Za-z0-9\u4e00-\u9fff&/ -]{2,24}")
OPERATION_TERMS = (
"点击",
"點擊",
"选择",
"選擇",
"输入",
"輸入",
"打开",
"打開",
"登入",
"用戶可",
"用户可",
"按<",
"點撃",
)
REPORT_TERMS = (
"年度报告",
"年报",
"財務報表",
"财务报表",
"公司治理",
"董事会",
"董事會",
"审计报告",
"審計報告",
"合并资产负债表",
"合併資產負債表",
"经营情况",
"經營情況",
"营业收入",
"營業收入",
"现金流量",
"現金流量",
"股东",
"股東",
)
REPORT_NAME_TERMS = ("annual", "report", "年度", "年报", "年報", "研报", "研報")
PDF_STRATEGY_FEATURE_STEPS = "pdf_feature_step_screenshot"
PDF_STRATEGY_OUTLINE_REPORT = "pdf_outline_report"
def _text(block: Block) -> str:
return (block.text or block.markdown or "").strip()
def _term_count(text: str, terms: tuple[str, ...]) -> int:
return sum(text.count(term) for term in terms)
def choose_pdf_chunk_strategy(path: Path, blocks: list[Block]) -> str:
"""Classify PDFs as operation manuals or report-like documents."""
text_blocks = [block for block in blocks if block.type != BlockType.IMAGE and _text(block)]
image_count = sum(1 for block in blocks if block.type == BlockType.IMAGE)
sample = "\n".join(_text(block) for block in text_blocks[:240])
filename = path.name.lower()
report_score = _term_count(sample, REPORT_TERMS)
if any(term in filename for term in REPORT_NAME_TERMS):
report_score += 3
feature_title_count = sum(1 for block in text_blocks if FEATURE_TITLE_RE.match(_text(block)))
operation_score = _term_count(sample, OPERATION_TERMS)
screenshot_density = image_count / max(len(text_blocks), 1)
if report_score >= 3 and operation_score < 18:
return PDF_STRATEGY_OUTLINE_REPORT
if feature_title_count >= 3 and operation_score >= 6 and image_count >= 3 and screenshot_density >= 0.12:
return PDF_STRATEGY_FEATURE_STEPS
return PDF_STRATEGY_OUTLINE_REPORT
+53
View File
@@ -0,0 +1,53 @@
"""Table-aware chunk grouping helpers."""
from __future__ import annotations
from rag_cut.models import Block, BlockType
def group_tables_with_context(blocks: list[Block]) -> list[list[Block]]:
"""
Group each table with its contextual title paragraph and footnote.
Keeps: heading/intro → table → footnote as one atomic group when adjacent.
"""
if not blocks:
return []
groups: list[list[Block]] = []
current: list[Block] = []
i = 0
while i < len(blocks):
block = blocks[i]
if block.type != BlockType.TABLE:
current.append(block)
i += 1
continue
group: list[Block] = []
if current:
tail = current[-1]
if tail.type in {BlockType.HEADING, BlockType.PARAGRAPH}:
group.append(tail)
current = current[:-1]
if current:
groups.append(current)
current = []
group.append(block)
j = i + 1
while j < len(blocks):
nxt = blocks[j]
if nxt.type == BlockType.PARAGRAPH and len((nxt.text or "")) <= 320:
if any(k in (nxt.text or "") for k in ("注", "备注", "说明", "Note", "※")):
group.append(nxt)
j += 1
break
break
groups.append(group)
i = j
if current:
groups.append(current)
return groups
+40
View File
@@ -0,0 +1,40 @@
"""Start the RAG-cut backend API server."""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
BACKEND_DIR = Path(__file__).resolve().parent
def main() -> None:
parser = argparse.ArgumentParser(description="Start RAG-cut backend API")
parser.add_argument("--host", default="127.0.0.1", help="Bind host (default: 127.0.0.1)")
parser.add_argument("--port", type=int, default=8000, help="Bind port (default: 8000)")
parser.add_argument(
"--reload",
action=argparse.BooleanOptionalAction,
default=True,
help="Auto-reload on code changes (default: on)",
)
args = parser.parse_args()
backend = str(BACKEND_DIR)
if backend not in sys.path:
sys.path.insert(0, backend)
import uvicorn
uvicorn.run(
"api.main:app",
host=args.host,
port=args.port,
reload=args.reload,
reload_dirs=[backend] if args.reload else None,
)
if __name__ == "__main__":
main()
View File
+244
View File
@@ -0,0 +1,244 @@
"""API integration tests for chunk persistence and recall."""
from __future__ import annotations
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
from fastapi.testclient import TestClient
from api.main import app
from rag_cut.models import Chunk, ChunkResult, SplitMode
def result(filename: str = "manual.txt") -> ChunkResult:
content = "交易密码可在账户安全页面修改"
return ChunkResult(
filename=filename,
doc_id="abcdef123456",
split_mode=SplitMode.DEFAULT,
block_count=2,
chunk_count=2,
chunks=[
Chunk(index=0, content="![logo](logo.png)", char_count=18, block_types=["image"], meta={"retrieval": False}),
Chunk(index=1, content=content, char_count=len(content), block_types=["paragraph"], meta={"retrieval": True}),
],
)
class ApiTest(unittest.TestCase):
def setUp(self) -> None:
self.client = TestClient(app)
def test_chunk_preserves_upload_filename_and_persists_result(self) -> None:
with tempfile.TemporaryDirectory() as tmp, patch("api.main.RESULTS_DIR", Path(tmp)):
with patch("api.main.chunk_document", side_effect=lambda path, config=None: result(path.name)):
response = self.client.post("/api/chunk", files={"file": ("用户手册.txt", b"body", "text/plain")})
self.assertEqual(response.status_code, 200)
self.assertEqual(response.json()["filename"], "用户手册.txt")
self.assertTrue((Path(tmp) / "abcdef123456.json").exists())
def test_recall_endpoint_excludes_preview_only_chunk(self) -> None:
with tempfile.TemporaryDirectory() as tmp, patch("api.main.RESULTS_DIR", Path(tmp)):
(Path(tmp) / "abcdef123456.json").write_text(result().model_dump_json(), encoding="utf-8")
response = self.client.post(
"/api/recall",
json={"doc_id": "abcdef123456", "query": "交易密码", "top_k": 5},
)
self.assertEqual(response.status_code, 200)
payload = response.json()
self.assertEqual(payload["candidate_count"], 1)
self.assertEqual(payload["results"][0]["chunk_index"], 1)
def test_recall_rejects_invalid_doc_id(self) -> None:
response = self.client.post("/api/recall", json={"doc_id": "../secret", "query": "test"})
self.assertEqual(response.status_code, 400)
def test_list_and_get_historical_results(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
results_dir = Path(tmp) / "results"
uploads_dir = Path(tmp) / "uploads"
results_dir.mkdir()
uploads_dir.mkdir()
(results_dir / "abcdef123456.json").write_text(result().model_dump_json(), encoding="utf-8")
upload_doc = uploads_dir / "abcdef123456"
upload_doc.mkdir()
(upload_doc / "manual.txt").write_text("hello original", encoding="utf-8")
with (
patch("api.main.RESULTS_DIR", results_dir),
patch("api.main.UPLOADS_DIR", uploads_dir),
):
listed = self.client.get("/api/results")
self.assertEqual(listed.status_code, 200)
body = listed.json()
self.assertEqual(body["count"], 1)
self.assertEqual(body["results"][0]["doc_id"], "abcdef123456")
self.assertEqual(body["results"][0]["filename"], "manual.txt")
self.assertTrue(body["results"][0]["has_original"])
detail = self.client.get("/api/results/abcdef123456")
self.assertEqual(detail.status_code, 200)
self.assertEqual(detail.json()["chunk_count"], 2)
self.assertTrue(detail.json()["has_original"])
original = self.client.get("/api/results/abcdef123456/original")
self.assertEqual(original.status_code, 200)
self.assertEqual(original.content, b"hello original")
def test_get_result_not_found(self) -> None:
with tempfile.TemporaryDirectory() as tmp, patch("api.main.RESULTS_DIR", Path(tmp)):
response = self.client.get("/api/results/abcdef123456")
self.assertEqual(response.status_code, 404)
def test_word_preview_serves_converted_pdf(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp)
uploads_dir = root / "uploads"
doc_dir = uploads_dir / "abcdef123456"
conversion_dir = root / "assets" / "abcdef123456" / "_conversion"
doc_dir.mkdir(parents=True)
conversion_dir.mkdir(parents=True)
(doc_dir / "manual.docx").write_bytes(b"word")
preview = conversion_dir / "manual.pdf"
preview.write_bytes(b"%PDF-1.7 preview")
with (
patch("api.main.UPLOADS_DIR", uploads_dir),
patch("api.main.STORAGE_DIR", root),
):
response = self.client.get("/api/results/abcdef123456/preview")
self.assertEqual(response.status_code, 200)
self.assertEqual(response.headers["content-type"], "application/pdf")
self.assertEqual(response.content, b"%PDF-1.7 preview")
def test_word_preview_returns_not_found_without_conversion(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp)
uploads_dir = root / "uploads"
doc_dir = uploads_dir / "abcdef123456"
doc_dir.mkdir(parents=True)
(doc_dir / "manual.doc").write_bytes(b"word")
with (
patch("api.main.UPLOADS_DIR", uploads_dir),
patch("api.main.STORAGE_DIR", root),
):
response = self.client.get("/api/results/abcdef123456/preview")
self.assertEqual(response.status_code, 404)
def test_delete_historical_result_removes_files(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp)
results_dir = root / "results"
uploads_dir = root / "uploads"
assets_dir = root / "assets" / "abcdef123456"
results_dir.mkdir()
uploads_dir.mkdir()
assets_dir.mkdir(parents=True)
(results_dir / "abcdef123456.json").write_text(result().model_dump_json(), encoding="utf-8")
upload_doc = uploads_dir / "abcdef123456"
upload_doc.mkdir()
(upload_doc / "manual.txt").write_text("hello", encoding="utf-8")
(assets_dir / "img.png").write_bytes(b"png")
with (
patch("api.main.RESULTS_DIR", results_dir),
patch("api.main.UPLOADS_DIR", uploads_dir),
patch("api.main.STORAGE_DIR", root),
):
deleted = self.client.delete("/api/results/abcdef123456")
self.assertEqual(deleted.status_code, 200)
body = deleted.json()
self.assertTrue(body["deleted"])
self.assertTrue(body["removed"]["result"])
self.assertTrue(body["removed"]["uploads"])
self.assertTrue(body["removed"]["assets"])
self.assertFalse((results_dir / "abcdef123456.json").exists())
self.assertFalse(upload_doc.exists())
self.assertFalse(assets_dir.exists())
listed = self.client.get("/api/results")
self.assertEqual(listed.json()["count"], 0)
missing = self.client.delete("/api/results/abcdef123456")
self.assertEqual(missing.status_code, 404)
def test_delete_rejects_invalid_doc_id(self) -> None:
response = self.client.delete("/api/results/not-a-valid")
self.assertEqual(response.status_code, 400)
def test_get_original_rejects_invalid_doc_id(self) -> None:
response = self.client.get("/api/results/../secret/original")
self.assertIn(response.status_code, (400, 404))
def test_chunk_accepts_manual_mode_and_preserves_zero_overlap(self) -> None:
captured: dict = {}
def fake_chunk(path, config=None):
captured["config"] = config
return result(path.name)
with tempfile.TemporaryDirectory() as tmp, patch("api.main.RESULTS_DIR", Path(tmp)):
with patch("api.main.chunk_document", side_effect=fake_chunk):
response = self.client.post(
"/api/chunk",
files={"file": ("sheet.csv", b"a,b\n1,2\n", "text/csv")},
data={
"mode": "by_row",
"overlap": "0",
"max_chunk_size": "2400",
"header_row_start": "1",
"header_row_end": "1",
"start_row": "2",
"rows_per_chunk": "5",
},
)
self.assertEqual(response.status_code, 200)
config = captured["config"]
self.assertIsNotNone(config)
self.assertEqual(config.mode, SplitMode.BY_ROW)
self.assertEqual(config.overlap, 0)
self.assertEqual(config.rows_per_chunk, 5)
def test_chunk_accepts_parent_child_mode_fields(self) -> None:
captured: dict = {}
def fake_chunk(path, config=None):
captured["config"] = config
return result(path.name)
with tempfile.TemporaryDirectory() as tmp, patch("api.main.RESULTS_DIR", Path(tmp)):
with patch("api.main.chunk_document", side_effect=fake_chunk):
response = self.client.post(
"/api/chunk",
files={"file": ("guide.md", b"a##b###c", "text/markdown")},
data={
"mode": "parent_child",
"parent_delimiter": "##",
"child_delimiter": "###",
"max_chunk_size": "2000",
"child_max_size": "512",
"overlap": "0",
},
)
self.assertEqual(response.status_code, 200)
config = captured["config"]
self.assertEqual(config.mode, SplitMode.PARENT_CHILD)
self.assertEqual(config.parent_delimiter, "##")
self.assertEqual(config.child_delimiter, "###")
self.assertEqual(config.child_max_size, 512)
self.assertEqual(config.overlap, 0)
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,290 @@
"""Tests for generic heading/layout multimodal splitting."""
from __future__ import annotations
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
from rag_cut.models import Block, BlockType, SplitConfig, SplitMode
from rag_cut.pipeline import chunk_document
from rag_cut.splitters.pdf_semantic import CHUNK_STRATEGY, split_pdf_semantic
def h(text: str, page: int = 1, level: int = 0, y: int = 100) -> Block:
return Block(
type=BlockType.HEADING if level else BlockType.PARAGRAPH,
text=text,
level=level,
meta={"page": page, "bbox": [72, y, 520, y + 20], "page_height": 800},
)
def p(text: str, page: int = 1, y: int = 130) -> Block:
return Block(
type=BlockType.PARAGRAPH,
text=text,
meta={"page": page, "bbox": [72, y, 520, y + 20], "page_height": 800},
)
def img(name: str, page: int = 1, y: int = 180) -> Block:
return Block(
type=BlockType.IMAGE,
image_id=name,
image_path=f"assets/test/{name}.png",
ocr_text="Login Submit",
meta={"page": page, "bbox": [100, y, 500, y + 180], "page_height": 800},
)
def table(page: int = 1, y: int = 220) -> Block:
return Block(
type=BlockType.TABLE,
markdown="| Field | Description |\n| --- | --- |\n| status | Account status |",
image_id="table1",
image_path="assets/test/table1.png",
meta={
"page": page,
"pages": [page],
"bbox": [72, y, 520, y + 120],
"bboxes": [{"page": page, "bbox": [72, y, 520, y + 120]}],
"row_count": 2,
"col_count": 2,
"cross_page": False,
"crop_path": "assets/test/table1.png",
},
)
class HeadingLayoutMultimodalTest(unittest.TestCase):
def test_numbered_headings_create_same_level_boundaries(self) -> None:
chunks = split_pdf_semantic(
[
h("1 Overview"),
p("Intro text."),
h("1.1 Account Status", y=180),
p("Account status details.", y=210),
h("1.2 Trade Detail", y=260),
p("Trade detail body.", y=290),
],
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
)
headings = [c.meta["heading"] for c in chunks]
self.assertEqual(headings, ["1 Overview", "1.1 Account Status", "1.2 Trade Detail"])
self.assertTrue(all(c.meta["chunk_strategy"] == CHUNK_STRATEGY for c in chunks))
def test_letter_subheadings_bind_nearest_images(self) -> None:
chunks = split_pdf_semantic(
[
h("1 Getting Started"),
h("A. Log in", y=140),
p("Open the login page.", y=170),
img("login", y=210),
h("B) Authentication", y=430),
p("Approve the authentication request.", y=460),
img("auth", y=500),
],
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
)
self.assertEqual([c.meta["heading"] for c in chunks], ["A. Log in", "B) Authentication"])
self.assertEqual(chunks[0].meta["heading_path"], ["1 Getting Started", "A. Log in"])
self.assertEqual(chunks[0].meta["images"][0]["image_id"], "login")
self.assertEqual(chunks[1].meta["images"][0]["image_id"], "auth")
def test_table_stays_atomic_with_markdown_and_position(self) -> None:
chunks = split_pdf_semantic(
[h("2 Settings"), p("The fields are listed below."), table(y=180)],
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
)
self.assertEqual(len(chunks), 1)
self.assertIn("| Field | Description |", chunks[0].content)
self.assertEqual(chunks[0].meta["tables"][0]["row_count"], 2)
self.assertEqual(chunks[0].meta["tables"][0]["bbox"], [72, 180, 520, 300])
def test_unlabelled_image_only_chunk_is_preview_only(self) -> None:
cover_logo = img("cover-logo")
cover_logo.ocr_text = ""
chunks = split_pdf_semantic(
[cover_logo, h("1 Overview", page=2), p("Useful body.", page=2)],
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
)
self.assertEqual(len(chunks), 2)
self.assertFalse(chunks[0].meta["retrieval"])
self.assertTrue(chunks[1].meta["retrieval"])
self.assertIn("cover-logo", chunks[0].content)
def test_image_only_chunk_with_ocr_remains_retrievable(self) -> None:
chunks = split_pdf_semantic(
[img("workflow-screenshot")],
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
)
self.assertEqual(len(chunks), 1)
self.assertTrue(chunks[0].meta["retrieval"])
def test_noise_blocks_are_skipped(self) -> None:
blocks = []
for page in range(1, 5):
blocks.append(p("Product Manual", page=page, y=20))
blocks.append(p(str(page), page=page, y=760))
blocks.extend(
[
p("Contents"),
p("1 Intro ........ 3"),
p("2 Setup ........ 8"),
h("1 Intro", page=2),
p("Useful body.", page=2),
]
)
chunks = split_pdf_semantic(blocks, SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000))
content = "\n".join(c.content for c in chunks)
self.assertIn("Useful body.", content)
self.assertNotIn("Product Manual", content)
self.assertNotIn("Contents", content)
self.assertNotIn("........", content)
self.assertNotIn("2 Setup", content)
def test_cross_page_same_heading_is_one_chunk(self) -> None:
chunks = split_pdf_semantic(
[h("3 Client"), p("Page one text.", page=1), p("Page two continuation.", page=2)],
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
)
self.assertEqual(len(chunks), 1)
self.assertEqual(chunks[0].meta["pages"], [1, 2])
self.assertIn("Page two continuation.", chunks[0].content)
def test_instructional_steps_stay_in_parent_section(self) -> None:
"""Step 1/2/3 lines must not become headings that get flushed away."""
chunks = split_pdf_semantic(
[
h("5.7 CRS SUBMISSION TO IRD", level=2),
h("5.7.1 Submission at 1st time", level=3, y=140),
p(
"For CRS submission at first time by G3SB, prior consent has to be obtained "
"from IRD by submitting test data file to the AEOI Portal for validation.",
y=170,
),
img("workflow", y=210),
p("Step 1 Users prepare the XML as per sections 5.1– 5.5.", y=420),
p('Step 2 Set parameter "Export in test data format?" value to "Y".', y=450),
p("Step 3 Check XML file tag DocTypeIndic was using OECD11.", y=480),
p(
"Step 4 Go to Registration and login page of AEOI Portal: "
"https://aeoi1.ird.gov.hk/portal/landing/",
y=510,
),
p("Download the encryption tools to encrypt the XML file.", y=540),
],
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=4000),
)
content = "\n".join(c.content for c in chunks)
self.assertIn("Step 1 Users prepare the XML", content)
self.assertIn("Export in test data format", content)
self.assertIn("OECD11", content)
self.assertIn("Step 4 Go to Registration", content)
self.assertTrue(any(c.meta["heading"] == "5.7.1 Submission at 1st time" for c in chunks))
def test_oversized_section_splits_by_length_without_parent(self) -> None:
blocks = [h("4 Reports")] + [p(f"Long paragraph {i} " + "x" * 80, y=130 + i) for i in range(8)]
chunks = split_pdf_semantic(blocks, SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=260))
self.assertGreater(len(chunks), 1)
self.assertTrue(all(c.meta.get("retrieval", True) for c in chunks))
self.assertFalse(any(c.meta.get("is_section_parent") for c in chunks))
self.assertFalse(any(c.meta.get("is_sub_chunk") for c in chunks))
self.assertFalse(any(c.meta.get("parent_chunk_id") is not None for c in chunks))
def test_catalog_titles_fold_into_section_not_dropped(self) -> None:
"""TOC-style '1.xxx' lines between catalog headings must not vanish on flush."""
chunks = split_pdf_semantic(
[
h("14.W. Whitney George", level=1, y=100),
p("Famous fund manager bio.", y=130),
Block(
type=BlockType.HEADING,
text="形態指標",
level=1,
meta={
"page": 1,
"bbox": [72, 200, 200, 220],
"page_height": 800,
"font_size": 16,
"body_font_size": 12,
},
),
p("共包含以下11個形態指標說明:", y=240),
h("1.頭肩頂形態", level=1, y=260),
h("2.頭肩底形態", level=1, y=280),
h("1.頭肩頂形態", level=1, page=2, y=120),
p("整體介紹:頂部反轉形態說明。", page=2, y=150),
],
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=4000),
)
morph = next(c for c in chunks if c.meta.get("heading") == "形態指標")
self.assertEqual(morph.meta.get("heading_path"), ["形態指標"])
self.assertIn("共包含以下11個形態指標說明:", morph.content)
self.assertIn("1.頭肩頂形態", morph.content)
self.assertIn("2.頭肩底形態", morph.content)
detail = next(
c for c in chunks if c.meta.get("heading") == "1.頭肩頂形態" and "整體介紹" in c.content
)
self.assertIn("整體介紹", detail.content)
def test_empty_explicit_section_does_not_pollute_previous_chunk(self) -> None:
empty_heading = h("2 REFERENCES", level=1, page=3, y=100)
empty_heading.meta["source_heading"] = True
chunks = split_pdf_semantic(
[
h("1 GLOSSARY", level=1, page=2, y=100),
table(page=2, y=140),
empty_heading,
h("3 DOCUMENT HISTORY", level=1, page=3, y=180),
p("History body.", page=3, y=220),
],
SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=4000),
)
glossary = next(c for c in chunks if c.meta.get("heading") == "1 GLOSSARY")
self.assertNotIn("2 REFERENCES", glossary.content)
self.assertFalse(any("2 REFERENCES" in c.content for c in chunks))
class DummyParser:
def parse(self, path: Path, assets_dir: Path) -> list[Block]:
return [h("1 Overview"), p("Pipeline body.")]
class PipelineSemanticIntegrationTest(unittest.TestCase):
def test_pdf_default_mode_uses_heading_layout_multimodal(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
tmp_root = Path(tmp)
source = tmp_root / "sample.pdf"
source.write_bytes(b"%PDF-1.4\n")
with patch("rag_cut.pipeline.get_parser", return_value=DummyParser()):
result = chunk_document(
source,
config=SplitConfig(mode=SplitMode.DEFAULT, max_chunk_size=2000),
storage_root=tmp_root / "storage",
)
self.assertEqual(result.chunk_count, 1)
self.assertEqual(result.split_config["chunk_strategy"], CHUNK_STRATEGY)
self.assertEqual(result.chunks[0].meta["chunk_strategy"], CHUNK_STRATEGY)
if __name__ == "__main__":
unittest.main()
+98
View File
@@ -0,0 +1,98 @@
"""Tests for PDF text merge and heading classification."""
from __future__ import annotations
import unittest
from rag_cut.layout_meta import enrich_layout_metadata
from rag_cut.models import Block, BlockType
from rag_cut.parsers.pdf.layout import TextLine, _merge_lines_to_paragraphs
from rag_cut.parsers.pdf.text_extract import MAX_HEADING_CHARS, _font_heading_level, _looks_like_ui_label
from rag_cut.parsers.pdf.noise_filter import is_margin_noise_block
from rag_cut.splitters.heading_splitter import is_heading_block, normalize_heading_block, numbered_heading_level
class ParagraphMergeTest(unittest.TestCase):
def test_does_not_merge_title_and_body_with_different_font_sizes(self) -> None:
lines = [
TextLine(x0=90, y0=90, x1=400, y1=110, text="Master Collection/Pattern Scanning", font_size=18),
TextLine(x0=90, y0=115, x1=480, y1=132, text="及所有大師/形態/策略指標的位置", font_size=17.5),
TextLine(
x0=90,
y0=145,
x1=500,
y1=200,
text="頁面左側【篩選】欄下方即是【Master Collection】文件夾。點擊此檔夾即可出現下拉頁面。",
font_size=11,
),
]
merged = _merge_lines_to_paragraphs(lines, page_width=600)
texts = [ln.text for ln in merged]
self.assertEqual(len(merged), 2)
self.assertIn("Master Collection", texts[0])
self.assertTrue(texts[1].startswith("頁面左側"))
self.assertNotIn("頁面左側", texts[0])
class FontHeadingTest(unittest.TestCase):
def test_short_large_text_is_heading(self) -> None:
level = _font_heading_level(18, 11, "Master Collection 位置說明")
self.assertIsNotNone(level)
def test_title_plus_body_blob_is_not_heading(self) -> None:
blob = (
"Master Collection/Pattern Scanning/Strategy Scanning 及所有大師/形態/策略指標的位置 "
"頁面左側【篩選】欄下方即是【Master Collection】文件夾。點擊此檔夾即可出現下拉頁面,"
"其中包含了十五個大師的選股策略。如圖所示:"
)
self.assertGreater(len(blob), MAX_HEADING_CHARS)
self.assertIsNone(_font_heading_level(18, 11, blob))
def test_compact_numbered_chinese_title_is_heading(self) -> None:
# FAQ style: "10.上升三角形態" with no space after the dot.
level = _font_heading_level(16, 12, "10.上升三角形態")
self.assertEqual(level, 1)
self.assertFalse(_looks_like_ui_label("10.上升三角形態"))
def test_cjk_section_banner_is_level1_not_ui_label(self) -> None:
self.assertFalse(_looks_like_ui_label("形態指標"))
self.assertEqual(_font_heading_level(16, 12, "形態指標"), 1)
class HeadingDemoteTest(unittest.TestCase):
def test_normalize_demotes_overlong_heading(self) -> None:
blob = "T" * (MAX_HEADING_CHARS + 20)
block = Block(type=BlockType.HEADING, text=blob, level=1, meta={"page": 1})
normalized = normalize_heading_block(block)
self.assertEqual(normalized.type, BlockType.PARAGRAPH)
self.assertFalse(is_heading_block(normalized))
def test_enrich_demotes_overlong_heading_block(self) -> None:
blob = (
"Master Collection/Pattern Scanning/Strategy Scanning Master Collection/"
"Pattern Scanning/Strategy Scanning 及所有大師 /形態/策略指標的位置 "
"頁面左側【篩選】欄下方即是【Master Collection】文件夾。點擊此檔夾即可出 "
"現下拉頁面,其中包含了十五個大師的選股策略。"
)
blocks = [
Block(type=BlockType.HEADING, text=blob, level=1, meta={"page": 1, "bbox": [90, 90, 500, 240]}),
Block(type=BlockType.IMAGE, image_id="a.png", image_path="a.png", meta={"page": 1, "bbox": [90, 250, 500, 480]}),
]
enriched = enrich_layout_metadata(blocks)
self.assertEqual(enriched[0].type, BlockType.PARAGRAPH)
self.assertNotEqual(enriched[1].meta.get("bound_heading"), blob)
def test_top_of_page_numbered_title_not_margin_noise(self) -> None:
block = Block(
type=BlockType.PARAGRAPH,
text="10.上升三角形態",
meta={"page": 50, "bbox": [90, 74, 209, 90], "page_height": 842},
)
self.assertFalse(is_margin_noise_block(block, 842))
self.assertEqual(numbered_heading_level("10.上升三角形態"), 1)
self.assertTrue(is_heading_block(normalize_heading_block(block)))
if __name__ == "__main__":
unittest.main()
+166
View File
@@ -0,0 +1,166 @@
"""Tests for MinerU content_list → Block mapping."""
from __future__ import annotations
import tempfile
import subprocess
import unittest
from pathlib import Path
from unittest.mock import Mock, patch
from rag_cut.models import BlockType
from rag_cut.parsers.mineru_adapter import (
_blocks_from_content_list,
_find_content_list,
_run_mineru,
)
class MineruAdapterProcessTest(unittest.TestCase):
def test_default_timeout_allows_large_manual_to_finish(self) -> None:
with patch.dict("os.environ", {}, clear=True):
from rag_cut.parsers.mineru_adapter import _mineru_timeout_sec
self.assertEqual(_mineru_timeout_sec(), 540.0)
def test_configured_api_url_is_passed_to_mineru(self) -> None:
process = Mock()
process.communicate.return_value = ("", "")
process.returncode = 0
with (
tempfile.TemporaryDirectory() as tmp,
patch("rag_cut.parsers.mineru_adapter._mineru_command", return_value="mineru"),
patch.dict("os.environ", {"RAG_CUT_MINERU_API_URL": "http://127.0.0.1:30000"}),
patch("rag_cut.parsers.mineru_adapter.subprocess.Popen", return_value=process) as popen,
):
ok = _run_mineru(Path(tmp) / "manual.pdf", Path(tmp) / "output")
self.assertTrue(ok)
args = popen.call_args.args[0]
self.assertEqual(args[-2:], ["--api-url", "http://127.0.0.1:30000"])
def test_timeout_terminates_windows_process_tree(self) -> None:
process = Mock()
process.pid = 4321
process.communicate.side_effect = subprocess.TimeoutExpired("mineru", 30)
process.poll.return_value = None
with (
tempfile.TemporaryDirectory() as tmp,
patch("rag_cut.parsers.mineru_adapter._mineru_command", return_value="mineru"),
patch("rag_cut.parsers.mineru_adapter._mineru_timeout_sec", return_value=30),
patch("rag_cut.parsers.mineru_adapter.os.name", "nt"),
patch("rag_cut.parsers.mineru_adapter.subprocess.Popen", return_value=process),
patch("rag_cut.parsers.mineru_adapter.subprocess.run") as run,
):
ok = _run_mineru(Path(tmp) / "manual.pdf", Path(tmp) / "output")
self.assertFalse(ok)
run.assert_called_once_with(
["taskkill", "/PID", "4321", "/T", "/F"],
check=False,
capture_output=True,
timeout=10,
)
class MineruAdapterMappingTest(unittest.TestCase):
def test_chart_becomes_image_block(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
assets = Path(tmp)
img = assets / "chart.jpg"
img.write_bytes(b"fake")
items = [
{
"type": "text",
"text": "第二十節 TradingView",
"text_level": 2,
"page_idx": 30,
"bbox": [127, 86, 363, 104],
},
{
"type": "chart",
"img_path": "chart.jpg",
"page_idx": 30,
"bbox": [129, 171, 878, 428],
"image_caption": [],
},
]
with patch("rag_cut.parsers.mineru_adapter._mineru_ocr_enabled", return_value=False):
blocks = _blocks_from_content_list(items, assets / "content_list.json", assets)
types = [b.type for b in blocks]
self.assertEqual(types, [BlockType.HEADING, BlockType.IMAGE])
self.assertEqual(blocks[1].image_id, "chart.jpg")
self.assertEqual(blocks[1].meta.get("mineru_type"), "chart")
self.assertEqual(blocks[1].meta.get("page"), 31)
def test_table_keeps_markdown_and_screenshot(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
assets = Path(tmp)
img = assets / "table.png"
img.write_bytes(b"png")
items = [
{
"type": "table",
"img_path": "table.png",
"page_idx": 1,
"bbox": [100, 100, 400, 300],
"table_body": "<table><tr><td>A</td><td>B</td></tr></table>",
}
]
with patch("rag_cut.parsers.mineru_adapter._mineru_ocr_enabled", return_value=False):
blocks = _blocks_from_content_list(items, assets / "content_list.json", assets)
self.assertEqual(len(blocks), 1)
self.assertEqual(blocks[0].type, BlockType.TABLE)
self.assertIn("<table>", blocks[0].markdown or "")
self.assertEqual(blocks[0].image_id, "table.png")
self.assertTrue(blocks[0].image_path)
def test_page_footnote_kept_as_paragraph(self) -> None:
items = [
{
"type": "page_footnote",
"text": "<sup>1</sup> See Market Master.",
"page_idx": 8,
"bbox": [85, 889, 912, 916],
}
]
with tempfile.TemporaryDirectory() as tmp:
blocks = _blocks_from_content_list(items, Path(tmp) / "content_list.json", Path(tmp))
self.assertEqual(len(blocks), 1)
self.assertEqual(blocks[0].type, BlockType.PARAGRAPH)
self.assertTrue(blocks[0].meta.get("is_footnote"))
def test_list_items_become_paragraph(self) -> None:
items = [
{
"type": "list",
"list_items": ["第一點說明", {"text": "第二點說明"}],
"page_idx": 2,
"bbox": [100, 200, 400, 260],
}
]
with tempfile.TemporaryDirectory() as tmp:
blocks = _blocks_from_content_list(items, Path(tmp) / "content_list.json", Path(tmp))
self.assertEqual(len(blocks), 1)
self.assertEqual(blocks[0].type, BlockType.PARAGRAPH)
self.assertIn("第一點說明", blocks[0].text or "")
self.assertIn("第二點說明", blocks[0].text or "")
def test_find_content_list_prefers_non_v2(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp) / "doc" / "auto"
root.mkdir(parents=True)
v1 = root / "doc_content_list.json"
v2 = root / "doc_content_list_v2.json"
v1.write_text("[]", encoding="utf-8")
v2.write_text("[]", encoding="utf-8")
found = _find_content_list(Path(tmp))
self.assertEqual(found, v1)
if __name__ == "__main__":
unittest.main()
+80
View File
@@ -0,0 +1,80 @@
"""Tests for parent/child delimiter splitting."""
from __future__ import annotations
import tempfile
import unittest
from pathlib import Path
from rag_cut.models import Block, BlockType, SplitConfig, SplitMode
from rag_cut.pipeline import chunk_document
from rag_cut.splitters.parent_child import split_by_parent_child
def p(text: str) -> Block:
return Block(type=BlockType.PARAGRAPH, text=text)
class ParentChildSplitTest(unittest.TestCase):
def test_parent_and_child_chunks_are_linked(self) -> None:
blocks = [p("Intro A##Detail A1###Detail A2##Intro B###Detail B1")]
groups = split_by_parent_child(
blocks,
SplitConfig(
mode=SplitMode.PARENT_CHILD,
parent_delimiter="##",
child_delimiter="###",
max_chunk_size=2000,
child_max_size=500,
overlap=0,
),
)
parents = [g for g in groups if g.meta.get("is_section_parent")]
children = [g for g in groups if g.meta.get("is_sub_chunk")]
self.assertGreaterEqual(len(parents), 2)
self.assertGreaterEqual(len(children), 2)
self.assertTrue(all(g.meta.get("retrieval") is False for g in parents))
self.assertTrue(all(g.meta.get("retrieval") is True for g in children))
self.assertTrue(all(g.meta.get("parent_section_id") for g in children))
def test_requires_parent_delimiter(self) -> None:
with self.assertRaises(ValueError):
split_by_parent_child(
[p("hello")],
SplitConfig(mode=SplitMode.PARENT_CHILD, parent_delimiter=None),
)
def test_pipeline_persists_parent_chunk_ids(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
source = Path(tmp) / "pc.md"
source.write_text(
"Alpha parent##A1 child###A2 child##Beta parent###B1 child",
encoding="utf-8",
)
result = chunk_document(
source,
config=SplitConfig(
mode=SplitMode.PARENT_CHILD,
parent_delimiter="##",
child_delimiter="###",
max_chunk_size=2000,
child_max_size=400,
overlap=0,
),
storage_root=Path(tmp) / "storage",
)
self.assertEqual(result.split_mode, SplitMode.PARENT_CHILD)
self.assertEqual(result.split_config.get("chunk_strategy"), "parent_child_delimiter")
parents = [c for c in result.chunks if c.meta.get("is_section_parent")]
children = [c for c in result.chunks if c.meta.get("is_sub_chunk")]
self.assertTrue(parents)
self.assertTrue(children)
for child in children:
self.assertIn("parent_chunk_id", child.meta)
self.assertIsInstance(child.meta["parent_chunk_id"], int)
if __name__ == "__main__":
unittest.main()
+315
View File
@@ -0,0 +1,315 @@
"""Tests for PDF margin noise filtering."""
from __future__ import annotations
import unittest
from pathlib import Path
from rag_cut.layout_meta import enrich_layout_metadata, sort_blocks_reading_order
from rag_cut.models import Block, BlockType
from rag_cut.parsers.mineru_adapter import _blocks_from_content_list
from rag_cut.parsers.pdf.noise_filter import (
detect_running_header_texts,
filter_noise_blocks,
filter_toc_blocks,
is_margin_noise_block,
is_tiny_image_block,
is_toc_entry_line,
is_toc_title_text,
)
class MineruNoiseFilterTest(unittest.TestCase):
def test_skips_header_and_page_number_types(self) -> None:
items = [
{"type": "text", "text": "目錄", "text_level": 2, "page_idx": 1},
{
"type": "text",
"text": "1. 登入 ...3\n2. 股票報價. ....4",
"page_idx": 1,
"bbox": [114, 190, 883, 517],
},
{"type": "header", "text": "H5 i-Trade 用戶使用手冊", "page_idx": 1},
{"type": "header", "text": "NEN WAFE:SOLUTIONS", "page_idx": 1},
{"type": "page_number", "text": "2", "page_idx": 1},
]
blocks = _blocks_from_content_list(items, Path("content_list.json"), Path("assets"))
# MinerU adapter may still emit 目录 text; pipeline noise filter removes it.
texts = [block.text for block in blocks]
self.assertIn("目錄", texts)
self.assertEqual(texts.count("H5 i-Trade 用戶使用手冊"), 0)
self.assertEqual(texts.count("2"), 0)
filtered = filter_noise_blocks(blocks)
filtered_texts = [block.text for block in filtered]
self.assertNotIn("目錄", filtered_texts)
self.assertFalse(any("登入" in (t or "") for t in filtered_texts))
class NoiseFilterHeuristicTest(unittest.TestCase):
def test_margin_page_number_is_noise(self) -> None:
block = Block(
type=BlockType.PARAGRAPH,
text="14",
meta={"page": 8, "bbox": [867, 942, 882, 954], "page_height": 1000},
)
self.assertTrue(is_margin_noise_block(block, 1000))
def test_body_paragraph_is_kept(self) -> None:
block = Block(
type=BlockType.PARAGRAPH,
text="開啟AFEH5 i-Trade網站,用戶無需登入即可查閱各項延遲15分鐘資訊。",
meta={"page": 3, "bbox": [116, 200, 880, 260], "page_height": 1000},
)
self.assertFalse(is_margin_noise_block(block, 1000))
def test_running_header_detection(self) -> None:
blocks = []
for page in range(1, 11):
blocks.append(
Block(
type=BlockType.PARAGRAPH,
text="H5 i-Trade 用戶使用手冊",
meta={"page": page, "bbox": [651, 79, 882, 97], "page_height": 1000},
)
)
blocks.append(
Block(
type=BlockType.HEADING,
text=f"{page}. 章節",
level=2,
meta={"page": page, "bbox": [116, 122, 191, 146], "page_height": 1000},
)
)
running = detect_running_header_texts(blocks)
self.assertIn("H5 i-Trade 用戶使用手冊", running)
self.assertNotIn("1. 章節", running)
def test_filter_noise_blocks_removes_margin_and_running_text(self) -> None:
blocks = [
Block(
type=BlockType.HEADING,
text="目錄",
level=2,
meta={"page": 2, "bbox": [119, 123, 176, 146], "page_height": 1000},
),
Block(
type=BlockType.PARAGRAPH,
text="1. 登入 ...3",
meta={"page": 2, "bbox": [114, 190, 883, 517], "page_height": 1000},
),
Block(
type=BlockType.PARAGRAPH,
text="2",
meta={"page": 2, "bbox": [868, 942, 880, 954], "page_height": 1000},
),
]
for page in range(1, 6):
blocks.append(
Block(
type=BlockType.PARAGRAPH,
text="H5 i-Trade 用戶使用手冊",
meta={"page": page, "bbox": [651, 79, 882, 97], "page_height": 1000},
)
)
filtered = filter_noise_blocks(blocks)
texts = [block.text for block in filtered]
self.assertNotIn("目錄", texts)
self.assertNotIn("1. 登入 ...3", texts)
self.assertNotIn("2", texts)
self.assertNotIn("H5 i-Trade 用戶使用手冊", texts)
def test_filter_toc_drops_contents_keeps_in_section_catalog(self) -> None:
blocks = [
Block(type=BlockType.HEADING, text="Contents", level=1, meta={"page": 1, "bbox": [72, 80, 200, 100]}),
Block(
type=BlockType.PARAGRAPH,
text="1 Intro ........ 3",
meta={"page": 1, "bbox": [72, 120, 400, 140]},
),
Block(
type=BlockType.PARAGRAPH,
text="2 Setup ........ 8",
meta={"page": 1, "bbox": [72, 150, 400, 170]},
),
Block(
type=BlockType.HEADING,
text="1 Intro",
level=1,
meta={"page": 2, "bbox": [72, 100, 200, 120]},
),
Block(
type=BlockType.PARAGRAPH,
text="Useful body paragraph describing the product workflow in detail.",
meta={"page": 2, "bbox": [72, 140, 500, 200]},
),
Block(
type=BlockType.HEADING,
text="形態指標",
level=1,
meta={"page": 45, "bbox": [72, 200, 200, 220]},
),
Block(
type=BlockType.PARAGRAPH,
text="共包含以下11個形態指標說明:",
meta={"page": 45, "bbox": [72, 240, 400, 260]},
),
Block(
type=BlockType.HEADING,
text="1.頭肩頂形態",
level=1,
meta={"page": 45, "bbox": [72, 270, 220, 290]},
),
]
self.assertTrue(is_toc_title_text("目錄"))
self.assertTrue(is_toc_entry_line("1 Intro ........ 3"))
self.assertFalse(is_toc_entry_line("1.頭肩頂形態"))
filtered = filter_toc_blocks(blocks)
texts = [b.text for b in filtered]
self.assertNotIn("Contents", texts)
self.assertNotIn("1 Intro ........ 3", texts)
self.assertIn("1 Intro", texts)
self.assertIn("Useful body paragraph describing the product workflow in detail.", texts)
self.assertIn("形態指標", texts)
self.assertIn("1.頭肩頂形態", texts)
def test_filter_toc_drops_mineru_compacted_contents_block(self) -> None:
contents = Block(
type=BlockType.PARAGRAPH,
text=(
"CONTENTS..\n"
"1 GLOSSARY ..2\n"
"5.1 AEOI ID SETTING.. 4"
),
meta={"page": 2, "mineru_type": "text"},
)
self.assertTrue(is_toc_title_text("CONTENTS.."))
self.assertTrue(is_toc_entry_line("1 GLOSSARY ..2"))
self.assertTrue(is_toc_entry_line("5.1 AEOI ID SETTING.. 4"))
self.assertEqual(filter_toc_blocks([contents]), [])
def test_filter_toc_keeps_real_section_below_directory_on_same_page(self) -> None:
blocks = [
Block(
type=BlockType.HEADING,
text="CONTENTS",
level=2,
meta={"page": 2, "bbox": [115, 129, 284, 152]},
),
Block(
type=BlockType.PARAGRAPH,
text="CONTENTS..\n1 GLOSSARY ..2\n2 REFERENCES 3\n3 DOCUMENT HISTORY.. .3",
meta={"page": 2, "bbox": [117, 167, 884, 718]},
),
Block(
type=BlockType.HEADING,
text="1 GLOSSARY",
level=2,
meta={"page": 2, "bbox": [117, 771, 321, 793]},
),
Block(
type=BlockType.TABLE,
markdown="| Abbreviation | Description |\n| --- | --- |",
meta={"page": 2, "bbox": [127, 809, 885, 891]},
),
]
filtered = filter_toc_blocks(blocks)
self.assertEqual([block.type for block in filtered], [BlockType.HEADING, BlockType.TABLE])
self.assertEqual(filtered[0].text, "1 GLOSSARY")
def test_tiny_logo_image_is_filtered_wide_flowchart_kept(self) -> None:
logo = Block(
type=BlockType.IMAGE,
image_id="logo",
image_path="assets/logo.png",
meta={"page": 14, "bbox": [121, 420, 156, 445], "page_height": 1000},
)
flowchart = Block(
type=BlockType.IMAGE,
image_id="flow",
image_path="assets/flow.png",
meta={"page": 13, "bbox": [119, 614, 949, 689], "page_height": 1000},
)
self.assertTrue(is_tiny_image_block(logo))
self.assertFalse(is_tiny_image_block(flowchart))
filtered = filter_noise_blocks([logo, flowchart])
ids = [block.image_id for block in filtered]
self.assertEqual(ids, ["flow"])
def test_nested_image_fragment_dropped(self) -> None:
outer = Block(
type=BlockType.IMAGE,
image_id="screenshot",
image_path="assets/shot.png",
meta={"page": 14, "bbox": [100, 400, 500, 700], "page_height": 1000},
)
nested = Block(
type=BlockType.IMAGE,
image_id="icon",
image_path="assets/icon.png",
meta={"page": 14, "bbox": [200, 450, 280, 530], "page_height": 1000},
)
filtered = filter_noise_blocks([outer, nested])
ids = [block.image_id for block in filtered]
self.assertEqual(ids, ["screenshot"])
class LayoutOrderAndBindTest(unittest.TestCase):
def test_sort_puts_page_top_before_lower_fragment(self) -> None:
misordered = [
Block(
type=BlockType.IMAGE,
image_id="logo",
image_path="assets/logo.png",
meta={"page": 14, "bbox": [121, 420, 156, 445]},
),
Block(
type=BlockType.PARAGRAPH,
text="5.8 APPENDIX",
meta={"page": 14, "bbox": [114, 85, 324, 107]},
),
Block(
type=BlockType.PARAGRAPH,
text="5.9 FAQ",
meta={"page": 14, "bbox": [114, 334, 231, 355]},
),
]
ordered = sort_blocks_reading_order(misordered)
texts = [(b.image_id or b.text) for b in ordered]
self.assertEqual(texts, ["5.8 APPENDIX", "5.9 FAQ", "logo"])
def test_enrich_binds_image_to_spatial_heading(self) -> None:
blocks = [
Block(
type=BlockType.HEADING,
text="5.7.2Subsequent submissions",
level=3,
meta={"page": 13, "bbox": [115, 468, 465, 488]},
),
Block(
type=BlockType.IMAGE,
image_id="wrong_early",
image_path="assets/x.png",
meta={"page": 14, "bbox": [341, 513, 431, 648]},
),
Block(
type=BlockType.PARAGRAPH,
text="5.9 FAQ",
meta={"page": 14, "bbox": [114, 334, 231, 355]},
),
]
enriched = enrich_layout_metadata(blocks)
faq_img = next(b for b in enriched if b.type == BlockType.IMAGE)
self.assertEqual(faq_img.meta.get("bound_heading"), "5.9 FAQ")
self.assertLess(
next(b for b in enriched if b.text == "5.9 FAQ").meta["order_index"],
faq_img.meta["order_index"],
)
if __name__ == "__main__":
unittest.main()
+51
View File
@@ -0,0 +1,51 @@
"""Tests for recall ranking and retrieval eligibility."""
from __future__ import annotations
import unittest
from rag_cut.models import Chunk
from rag_cut.retrieval import recall_chunks
def chunk(index: int, content: str, *, retrieval: bool = True, heading: str = "") -> Chunk:
return Chunk(
index=index,
content=content,
char_count=len(content),
block_types=["paragraph"],
meta={"retrieval": retrieval, "heading": heading},
)
class RecallChunksTest(unittest.TestCase):
def test_relevant_chunk_ranks_first(self) -> None:
results, candidate_count = recall_chunks(
"股票交易密码",
[
chunk(0, "登录后可以修改股票交易密码", heading="账户安全"),
chunk(1, "年度报告及公司治理", heading="公司资料"),
],
top_k=2,
)
self.assertEqual(candidate_count, 2)
self.assertEqual(results[0]["chunk_index"], 0)
self.assertGreater(results[0]["score"], 0)
def test_preview_only_chunks_are_not_candidates(self) -> None:
results, candidate_count = recall_chunks(
"logo",
[chunk(0, "![logo](logo.png)", retrieval=False), chunk(1, "Useful body")],
)
self.assertEqual(candidate_count, 1)
self.assertEqual(results, [])
def test_top_k_is_respected(self) -> None:
results, _ = recall_chunks("account", [chunk(i, f"account details {i}") for i in range(5)], top_k=2)
self.assertEqual(len(results), 2)
if __name__ == "__main__":
unittest.main()
+80
View File
@@ -0,0 +1,80 @@
"""Tests for spreadsheet layout detection and row splitting."""
from __future__ import annotations
import unittest
from rag_cut.parsers.pdf.tables import detect_spreadsheet_layout, is_qa_style_table
from rag_cut.pipeline import chunk_document
from rag_cut.splitters.by_row import split_by_row
from rag_cut.models import Block, BlockType, SplitConfig, SplitMode
BACKMAN_HEADER = [
[
"Session会话(必填):用于标识1个对话",
"query 用户输入(必填):消息内容",
"用户ID(必填)",
"使用的大语言模型(必填)",
"要求AI回复的语言(必填)",
"reference_output 标准答案(可选)",
],
["session", "query", "userid", "model", "lang", "reference_output"],
["1", "账户余额是多少?", "SUPPORT", "gpt-4o-mini", "zh", "SELECT 1"],
["2", "今日有哪些账户透支?", "SUPPORT", "gpt-4o-mini", "zh", "SELECT 2"],
]
class SpreadsheetLayoutTest(unittest.TestCase):
def test_detect_template_description_and_header_rows(self) -> None:
layout = detect_spreadsheet_layout(BACKMAN_HEADER)
self.assertEqual(layout["preamble_rows"], 1)
self.assertEqual(layout["header_row_start"], 2)
self.assertEqual(layout["header_row_end"], 2)
self.assertEqual(layout["data_start_row"], 3)
self.assertTrue(is_qa_style_table(BACKMAN_HEADER, layout))
def test_split_by_row_uses_real_header_and_one_row_per_chunk(self) -> None:
layout = detect_spreadsheet_layout(BACKMAN_HEADER)
block = Block(
type=BlockType.TABLE,
markdown="",
meta={"rows": BACKMAN_HEADER, "table_title": "testset", **layout},
)
groups = split_by_row(
[block],
SplitConfig(
mode=SplitMode.BY_ROW,
header_row_start=layout["header_row_start"],
header_row_end=layout["header_row_end"],
start_row=layout["data_start_row"],
rows_per_chunk=1,
),
)
self.assertEqual(len(groups), 2)
first_md = groups[0][0].markdown
self.assertIn("| session | query | userid | model | lang | reference_output |", first_md)
self.assertIn("账户余额是多少?", first_md)
self.assertNotIn("Session会话(必填)", first_md)
class BackmanFixtureTest(unittest.TestCase):
def test_backman_upload_chunks_one_row_each(self) -> None:
from pathlib import Path
uploads = list((Path(__file__).resolve().parents[2] / "storage" / "uploads").rglob("BackmanAI*.xlsx"))
if not uploads:
self.skipTest("BackmanAI fixture not uploaded")
result = chunk_document(uploads[0])
self.assertEqual(result.split_config["rows_per_chunk"], 1)
self.assertEqual(result.split_config["header_row_start"], 2)
self.assertEqual(result.split_config["start_row"], 3)
self.assertGreaterEqual(result.chunk_count, 40)
first = result.chunks[0].content
self.assertIn("| session | query |", first)
self.assertNotIn("Session会话(必填)", first)
self.assertLess(result.chunks[0].char_count, 1200)
if __name__ == "__main__":
unittest.main()