@@ -0,0 +1 @@
|
||||
"""RAG-cut FastAPI app. Start with: python backend/run.py"""
|
||||
@@ -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")
|
||||
@@ -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"]
|
||||
@@ -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)
|
||||
@@ -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"")
|
||||
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""
|
||||
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
|
||||
@@ -0,0 +1,5 @@
|
||||
"""Document parsers."""
|
||||
|
||||
from rag_cut.parsers.registry import get_parser
|
||||
|
||||
__all__ = ["get_parser"]
|
||||
@@ -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
|
||||
@@ -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),
|
||||
)
|
||||
]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -0,0 +1,5 @@
|
||||
"""PyMuPDF-based PDF parsing pipeline."""
|
||||
|
||||
from rag_cut.parsers.pdf.pipeline import parse_pdf
|
||||
|
||||
__all__ = ["parse_pdf"]
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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))
|
||||
@@ -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}"
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
@@ -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, "", [], {})}
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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,
|
||||
},
|
||||
)
|
||||
]
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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"]
|
||||
@@ -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]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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="", 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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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, "", 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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user