feat(documents): add text pdf ingestion with heading reconstruction
Why: - Support Text PDFs in the same DocumentTree/markdown contract as DOCX. Changes: - PyMuPDF parser, text-layer gate, shared heading heuristics, upload dispatch for .pdf. Impact: - Scanned/image PDFs are rejected at upload; requires pymupdf installed. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -26,3 +26,6 @@ SEMANTIC_MIN_CHUNK_SIZE=3
|
||||
|
||||
# SQLite database path
|
||||
DATABASE_URL=sqlite:///./data/chunking_benchmark.db
|
||||
# Text PDF gate (reject scanned/image PDFs)
|
||||
PDF_MIN_TOTAL_CHARS=100
|
||||
PDF_MIN_MEDIAN_CHARS_PER_PAGE=40
|
||||
|
||||
16
docs/adr/0016-pymupdf-for-text-pdf-extraction.md
Normal file
16
docs/adr/0016-pymupdf-for-text-pdf-extraction.md
Normal file
@@ -0,0 +1,16 @@
|
||||
# PyMuPDF for Text PDF extraction
|
||||
|
||||
Text PDF ingestion uses PyMuPDF (`fitz`) so Heading Reconstruction can read font size/weight, text blocks, and outline bookmarks — the same structural signals `python-docx` gives us for DOCX. We rejected flat string extractors (`pypdf`), table-first libraries as the primary path (`pdfplumber`), and LibreOffice PDF→DOCX conversion (lossy styles, slow, fights ADR 0006’s “prefer native structure”).
|
||||
|
||||
## Considered Options
|
||||
|
||||
- **PyMuPDF** — chosen; best fit for Heading Reconstruction; AGPL acceptable while this stays an internal benchmarker
|
||||
- **pdfplumber** — strong tables, weaker hierarchy; deferred to backlog if tables are a measured failure
|
||||
- **pypdf** — too little layout/font signal
|
||||
- **LibreOffice PDF→DOCX → existing parser** — reuses DOCX path but conversion quality is unreliable
|
||||
|
||||
## Consequences
|
||||
|
||||
- Add `pymupdf` dependency; keep a single PDF code path in the documents parser seam
|
||||
- If we later ship the parser as a distributed service, revisit AGPL vs a permissive alternative
|
||||
- Table-heavy and Scanned PDF work stays out of this ADR (see `backlog/`)
|
||||
@@ -35,5 +35,9 @@ class Settings(BaseSettings):
|
||||
# Database
|
||||
database_url: str = "sqlite:///./data/chunking_benchmark.db"
|
||||
|
||||
# Text PDF gate (reject Scanned PDFs with near-empty text layer)
|
||||
pdf_min_total_chars: int = 100
|
||||
pdf_min_median_chars_per_page: int = 40
|
||||
|
||||
|
||||
settings = Settings()
|
||||
79
src/documents/heading_heuristics.py
Normal file
79
src/documents/heading_heuristics.py
Normal file
@@ -0,0 +1,79 @@
|
||||
"""Shared heading detection heuristics for DOCX and Text PDF parsers.
|
||||
|
||||
Used when native styles/fonts/outlines are missing or incomplete
|
||||
(common in table-heavy Farsi documents).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Protocol
|
||||
|
||||
|
||||
class StyledBlock(Protocol):
|
||||
"""Minimal block interface: mutable style_name + text."""
|
||||
|
||||
style_name: str
|
||||
text: str
|
||||
|
||||
|
||||
_HEADING_PATTERNS: list[tuple[re.Pattern[str], int]] = [
|
||||
(re.compile(r"^بخش\s+"), 1),
|
||||
(re.compile(r"^\d+[\-\.]\d+"), 2),
|
||||
(re.compile(r"^\*\s*بخش\s+"), 1),
|
||||
(re.compile(r"^\*\s*\S"), 2),
|
||||
]
|
||||
|
||||
|
||||
def heading_level_from_style(style_name: str) -> int | None:
|
||||
"""Return heading level from a style name like 'Heading 1' or 'Heading1'."""
|
||||
m = re.match(r"^heading\s*(\d+)$", style_name.strip(), re.IGNORECASE)
|
||||
if m:
|
||||
return int(m.group(1))
|
||||
return None
|
||||
|
||||
|
||||
def is_likely_heading(text: str) -> int | None:
|
||||
"""Detect heading-like patterns in body text.
|
||||
|
||||
Returns heading level (1 or 2) if detected, else None.
|
||||
"""
|
||||
text = text.strip()
|
||||
if len(text) > 150:
|
||||
return None
|
||||
for pattern, level in _HEADING_PATTERNS:
|
||||
if pattern.match(text):
|
||||
return level
|
||||
return None
|
||||
|
||||
|
||||
def apply_heading_heuristics(
|
||||
blocks: list[StyledBlock],
|
||||
*,
|
||||
only_when_no_headings: bool = True,
|
||||
) -> bool:
|
||||
"""Upgrade Normal blocks to HeadingN via text patterns.
|
||||
|
||||
Args:
|
||||
blocks: Mutable text blocks with style_name / text.
|
||||
only_when_no_headings: If True (DOCX default), skip when any
|
||||
Heading styles already exist. If False, only upgrade remaining
|
||||
Normal blocks (PDF gap-fill after outline/font).
|
||||
|
||||
Returns:
|
||||
True if any blocks were upgraded.
|
||||
"""
|
||||
if only_when_no_headings:
|
||||
if any(heading_level_from_style(b.style_name) is not None for b in blocks):
|
||||
return False
|
||||
|
||||
upgraded = False
|
||||
for block in blocks:
|
||||
if block.style_name != "Normal":
|
||||
continue
|
||||
level = is_likely_heading(block.text)
|
||||
if level is not None:
|
||||
block.style_name = f"Heading{level}"
|
||||
upgraded = True
|
||||
|
||||
return upgraded
|
||||
@@ -1,18 +1,13 @@
|
||||
"""DOCX parser: extracts document tree + flat markdown.
|
||||
"""Document parser: DOCX/DOC + dispatcher for all supported formats.
|
||||
|
||||
Uses python-docx to read paragraph styles and build a hierarchical
|
||||
DocumentTree (Document > Section > Article > Paragraph). Also produces
|
||||
a markdown representation consumed by chunking strategies.
|
||||
Produces a hierarchical DocumentTree (Document > Section > Article > Paragraph)
|
||||
and markdown consumed by chunking strategies.
|
||||
|
||||
Supports both .docx and .doc formats. .doc files are converted to
|
||||
.docx via LibreOffice headless mode before parsing.
|
||||
|
||||
Handles documents where content is in tables (not just paragraphs).
|
||||
Supports .docx, .doc (via LibreOffice), and .pdf (via pdf_parser).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import subprocess
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
@@ -27,6 +22,10 @@ from src.core.models import (
|
||||
NodeType,
|
||||
)
|
||||
from src.core.exceptions import DocumentProcessingError
|
||||
from src.documents.heading_heuristics import (
|
||||
apply_heading_heuristics,
|
||||
heading_level_from_style,
|
||||
)
|
||||
|
||||
|
||||
# ── Heading level → node type mapping ──────────────────────────────
|
||||
@@ -34,89 +33,41 @@ from src.core.exceptions import DocumentProcessingError
|
||||
_HEADING_MAP: dict[int, NodeType] = {
|
||||
1: NodeType.SECTION,
|
||||
2: NodeType.ARTICLE,
|
||||
# 3+ also ARTICLE (nesting depth determines hierarchy)
|
||||
}
|
||||
|
||||
|
||||
def _heading_level(style_name: str) -> int | None:
|
||||
"""Return the heading level from a style name, or None if not a heading.
|
||||
|
||||
Handles both "Heading 1" (display name) and "Heading1" (style ID from XML).
|
||||
"""
|
||||
m = re.match(r"^heading\s*(\d+)$", style_name.strip(), re.IGNORECASE)
|
||||
if m:
|
||||
return int(m.group(1))
|
||||
return None
|
||||
|
||||
|
||||
def _node_type_for_level(level: int) -> NodeType:
|
||||
return _HEADING_MAP.get(level, NodeType.ARTICLE)
|
||||
|
||||
|
||||
# ── Heuristic heading detection for table-heavy documents ──────────
|
||||
# ── Text block ─────────────────────────────────────────────────────
|
||||
|
||||
# Patterns that look like section/article headings in Farsi/English docs
|
||||
_HEADING_PATTERNS: list[tuple[str, int]] = [
|
||||
# Farsi section markers: "بخش اول", "بخش دوم", etc.
|
||||
(re.compile(r"^بخش\s+"), 1),
|
||||
# Numbered sections: "1-1", "2-1", "10-1", "1-1-8"
|
||||
(re.compile(r"^\d+[\-\.]\d+"), 2),
|
||||
# Starred sections: "*بخش اول", "*تعریف"
|
||||
(re.compile(r"^\*\s*بخش\s+"), 1),
|
||||
(re.compile(r"^\*\s*\S"), 2),
|
||||
]
|
||||
class TextBlock:
|
||||
"""A unit of text extracted from a document, preserving reading order."""
|
||||
|
||||
__slots__ = ("style_name", "text", "font_size", "bold", "page", "locked")
|
||||
|
||||
def _is_likely_heading(text: str) -> int | None:
|
||||
"""Detect heading-like patterns in table-extracted text.
|
||||
|
||||
Returns the heading level (1 or 2) if detected, else None.
|
||||
Used when the document has no real heading styles (table-only content).
|
||||
"""
|
||||
text = text.strip()
|
||||
if len(text) > 150: # headings are short
|
||||
return None
|
||||
for pattern, level in _HEADING_PATTERNS:
|
||||
if pattern.match(text):
|
||||
return level
|
||||
return None
|
||||
|
||||
|
||||
def _detect_heading_blocks(blocks: list[_TextBlock]) -> bool:
|
||||
"""Upgrade paragraph blocks to heading blocks based on content patterns.
|
||||
|
||||
Returns True if any blocks were upgraded.
|
||||
Only activates when no real headings exist in the document.
|
||||
"""
|
||||
# Check if there are already real headings
|
||||
has_headings = any(_heading_level(b.style_name) is not None for b in blocks)
|
||||
if has_headings:
|
||||
return False
|
||||
|
||||
upgraded = False
|
||||
for block in blocks:
|
||||
if block.style_name != "Normal":
|
||||
continue
|
||||
level = _is_likely_heading(block.text)
|
||||
if level is not None:
|
||||
block.style_name = f"Heading{level}"
|
||||
upgraded = True
|
||||
|
||||
return upgraded
|
||||
def __init__(
|
||||
self,
|
||||
style_name: str,
|
||||
text: str,
|
||||
*,
|
||||
font_size: float = 0.0,
|
||||
bold: bool = False,
|
||||
page: int = 0,
|
||||
locked: bool = False,
|
||||
) -> None:
|
||||
self.style_name = style_name
|
||||
self.text = text
|
||||
self.font_size = font_size
|
||||
self.bold = bold
|
||||
self.page = page
|
||||
self.locked = locked
|
||||
|
||||
|
||||
# ── Text block extraction (paragraphs + tables) ───────────────────
|
||||
|
||||
class _TextBlock:
|
||||
"""A unit of text extracted from the document, preserving reading order."""
|
||||
__slots__ = ("style_name", "text")
|
||||
|
||||
def __init__(self, style_name: str, text: str) -> None:
|
||||
self.style_name = style_name
|
||||
self.text = text
|
||||
|
||||
|
||||
def _extract_text_blocks(doc: DocxDocumentType) -> list[_TextBlock]:
|
||||
def _extract_text_blocks(doc: DocxDocumentType) -> list[TextBlock]:
|
||||
"""Walk the document body in reading order, extracting paragraphs and tables.
|
||||
|
||||
This handles documents where content lives inside table cells
|
||||
@@ -124,20 +75,17 @@ def _extract_text_blocks(doc: DocxDocumentType) -> list[_TextBlock]:
|
||||
"""
|
||||
from docx.oxml.ns import qn
|
||||
|
||||
blocks: list[_TextBlock] = []
|
||||
blocks: list[TextBlock] = []
|
||||
|
||||
for element in doc.element.body:
|
||||
tag = element.tag.split("}")[-1] if "}" in element.tag else element.tag
|
||||
|
||||
if tag == "p":
|
||||
# Paragraph — extract text and style
|
||||
text = element.text or ""
|
||||
# Also check for runs (text split across formatting)
|
||||
if not text.strip():
|
||||
runs = element.findall(qn("w:r"))
|
||||
text = "".join(r.text or "" for r in runs)
|
||||
|
||||
# Get style name
|
||||
ppr = element.find(qn("w:pPr"))
|
||||
style_name = "Normal"
|
||||
if ppr is not None:
|
||||
@@ -146,19 +94,15 @@ def _extract_text_blocks(doc: DocxDocumentType) -> list[_TextBlock]:
|
||||
style_name = pstyle.get(qn("w:val"), "Normal")
|
||||
|
||||
if text.strip():
|
||||
blocks.append(_TextBlock(style_name, text.strip()))
|
||||
blocks.append(TextBlock(style_name, text.strip()))
|
||||
|
||||
elif tag == "tbl":
|
||||
# Table — extract all cell text as paragraph blocks
|
||||
for row in element.findall(qn("w:tr")):
|
||||
for cell in row.findall(qn("w:tc")):
|
||||
for para in cell.findall(qn("w:p")):
|
||||
# Get paragraph text
|
||||
text = ""
|
||||
runs = para.findall(qn("w:r"))
|
||||
text = "".join(r.text or "" for r in runs)
|
||||
|
||||
# Get style
|
||||
ppr = para.find(qn("w:pPr"))
|
||||
style_name = "Normal"
|
||||
if ppr is not None:
|
||||
@@ -167,39 +111,36 @@ def _extract_text_blocks(doc: DocxDocumentType) -> list[_TextBlock]:
|
||||
style_name = pstyle.get(qn("w:val"), "Normal")
|
||||
|
||||
if text.strip():
|
||||
blocks.append(_TextBlock(style_name, text.strip()))
|
||||
blocks.append(TextBlock(style_name, text.strip()))
|
||||
|
||||
return blocks
|
||||
|
||||
|
||||
# ── Tree builder ───────────────────────────────────────────────────
|
||||
|
||||
def build_document_tree(paragraphs: list[Paragraph] | list[_TextBlock]) -> DocumentTreeNode:
|
||||
def build_document_tree(paragraphs: list[Paragraph] | list[TextBlock]) -> DocumentTreeNode:
|
||||
"""Build a DocumentTreeNode tree from a list of text blocks.
|
||||
|
||||
Accepts either python-docx Paragraph objects or _TextBlock objects.
|
||||
Accepts either python-docx Paragraph objects or TextBlock objects.
|
||||
"""
|
||||
root = DocumentTreeNode(node_type=NodeType.DOCUMENT, text="", heading=None, heading_level=None)
|
||||
|
||||
# Stack of (level, node) for current nesting. level 0 = root.
|
||||
stack: list[tuple[int, DocumentTreeNode]] = [(0, root)]
|
||||
|
||||
for block in paragraphs:
|
||||
# Get style name and text from either type
|
||||
if isinstance(block, _TextBlock):
|
||||
if isinstance(block, TextBlock):
|
||||
style_name = block.style_name
|
||||
text = block.text
|
||||
else:
|
||||
style_name = block.style.name
|
||||
text = block.text.strip()
|
||||
|
||||
level = _heading_level(style_name)
|
||||
level = heading_level_from_style(style_name)
|
||||
|
||||
if not text:
|
||||
continue # skip blank paragraphs
|
||||
continue
|
||||
|
||||
if level is not None:
|
||||
# Pop back to parent level
|
||||
while len(stack) > 1 and stack[-1][0] >= level:
|
||||
stack.pop()
|
||||
|
||||
@@ -212,7 +153,6 @@ def build_document_tree(paragraphs: list[Paragraph] | list[_TextBlock]) -> Docum
|
||||
stack[-1][1].children.append(node)
|
||||
stack.append((level, node))
|
||||
else:
|
||||
# Body text — add as paragraph child of current heading
|
||||
node = DocumentTreeNode(
|
||||
node_type=NodeType.PARAGRAPH,
|
||||
text=text,
|
||||
@@ -244,7 +184,7 @@ def tree_to_markdown(node: DocumentTreeNode, depth: int = 0) -> str:
|
||||
# ── Public API ─────────────────────────────────────────────────────
|
||||
|
||||
class ParseResult:
|
||||
"""Output of parse_docx(): tree + markdown + raw text."""
|
||||
"""Output of parse_document() / parse_docx() / parse_pdf()."""
|
||||
|
||||
__slots__ = ("tree", "markdown", "plain_text", "paragraph_count")
|
||||
|
||||
@@ -261,6 +201,31 @@ class ParseResult:
|
||||
self.paragraph_count = paragraph_count
|
||||
|
||||
|
||||
SUPPORTED_SUFFIXES = (".docx", ".doc", ".pdf")
|
||||
|
||||
|
||||
def parse_document(file_path: str | Path) -> ParseResult:
|
||||
"""Parse a supported document into DocumentTree + markdown.
|
||||
|
||||
Dispatches by suffix: .pdf → pdf_parser; .docx/.doc → parse_docx.
|
||||
"""
|
||||
path = Path(file_path)
|
||||
if not path.exists():
|
||||
raise DocumentProcessingError(f"File not found: {path}")
|
||||
|
||||
suffix = path.suffix.lower()
|
||||
if suffix == ".pdf":
|
||||
from src.documents.pdf_parser import parse_pdf
|
||||
|
||||
return parse_pdf(path)
|
||||
if suffix in (".docx", ".doc"):
|
||||
return parse_docx(path)
|
||||
|
||||
raise DocumentProcessingError(
|
||||
f"Not a supported file format: {suffix} (expected {', '.join(SUPPORTED_SUFFIXES)})"
|
||||
)
|
||||
|
||||
|
||||
# ── .doc → .docx conversion ───────────────────────────────────────
|
||||
|
||||
def _convert_doc_to_docx(doc_path: Path) -> Path:
|
||||
@@ -268,9 +233,6 @@ def _convert_doc_to_docx(doc_path: Path) -> Path:
|
||||
|
||||
Returns the path to the converted .docx file (in a temp directory).
|
||||
The caller is responsible for cleanup.
|
||||
|
||||
Raises:
|
||||
DocumentProcessingError: If conversion fails.
|
||||
"""
|
||||
out_dir = Path(tempfile.mkdtemp(prefix="docconv_"))
|
||||
try:
|
||||
@@ -291,7 +253,6 @@ def _convert_doc_to_docx(doc_path: Path) -> Path:
|
||||
f"LibreOffice conversion failed: {result.stderr}"
|
||||
)
|
||||
|
||||
# Find the converted file
|
||||
converted = out_dir / doc_path.with_suffix(".docx").name
|
||||
if not converted.exists():
|
||||
raise DocumentProcessingError(
|
||||
@@ -312,15 +273,6 @@ def parse_docx(file_path: str | Path) -> ParseResult:
|
||||
"""Parse a .docx or .doc file into a DocumentTree + markdown.
|
||||
|
||||
.doc files are automatically converted to .docx via LibreOffice.
|
||||
|
||||
Args:
|
||||
file_path: Path to the .docx or .doc file.
|
||||
|
||||
Returns:
|
||||
ParseResult with tree, markdown, plain_text, and paragraph_count.
|
||||
|
||||
Raises:
|
||||
DocumentProcessingError: If the file cannot be parsed.
|
||||
"""
|
||||
path = Path(file_path)
|
||||
if not path.exists():
|
||||
@@ -328,9 +280,10 @@ def parse_docx(file_path: str | Path) -> ParseResult:
|
||||
|
||||
suffix = path.suffix.lower()
|
||||
if suffix not in (".docx", ".doc"):
|
||||
raise DocumentProcessingError(f"Not a supported file format: {suffix} (expected .docx or .doc)")
|
||||
raise DocumentProcessingError(
|
||||
f"Not a supported file format: {suffix} (expected .docx or .doc)"
|
||||
)
|
||||
|
||||
# Convert .doc to .docx if needed
|
||||
if suffix == ".doc":
|
||||
path = _convert_doc_to_docx(path)
|
||||
|
||||
@@ -339,14 +292,12 @@ def parse_docx(file_path: str | Path) -> ParseResult:
|
||||
except Exception as exc:
|
||||
raise DocumentProcessingError(f"Failed to open DOCX: {exc}") from exc
|
||||
|
||||
# Extract text blocks from paragraphs + tables (preserves reading order)
|
||||
blocks = _extract_text_blocks(doc)
|
||||
|
||||
if not blocks:
|
||||
raise DocumentProcessingError("Document contains no text content")
|
||||
|
||||
# Detect heading patterns in table-only documents
|
||||
_detect_heading_blocks(blocks)
|
||||
apply_heading_heuristics(blocks, only_when_no_headings=True)
|
||||
|
||||
root = build_document_tree(blocks)
|
||||
tree = DocumentTree(root=root)
|
||||
|
||||
228
src/documents/pdf_parser.py
Normal file
228
src/documents/pdf_parser.py
Normal file
@@ -0,0 +1,228 @@
|
||||
"""Text PDF parser via PyMuPDF with Heading Reconstruction.
|
||||
|
||||
Pipeline:
|
||||
1. Text-layer Gate — reject Scanned PDFs (near-empty text)
|
||||
2. Extract reading-order blocks (Table Flattening via block order)
|
||||
3. Heading Reconstruction: outline → font size → shared heuristics
|
||||
4. Build DocumentTree + markdown (same contract as DOCX)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import statistics
|
||||
from pathlib import Path
|
||||
|
||||
import fitz # PyMuPDF
|
||||
|
||||
from src.core.config import settings
|
||||
from src.core.exceptions import DocumentProcessingError
|
||||
from src.core.models import DocumentTree
|
||||
from src.documents.heading_heuristics import apply_heading_heuristics
|
||||
from src.documents.parser import (
|
||||
ParseResult,
|
||||
TextBlock,
|
||||
build_document_tree,
|
||||
tree_to_markdown,
|
||||
)
|
||||
|
||||
|
||||
def _median_chars_per_page(page_char_counts: list[int]) -> float:
|
||||
if not page_char_counts:
|
||||
return 0.0
|
||||
return float(statistics.median(page_char_counts))
|
||||
|
||||
|
||||
def _enforce_text_layer_gate(doc: fitz.Document) -> None:
|
||||
"""Hard-reject PDFs without enough extractable text (Scanned PDFs)."""
|
||||
page_counts: list[int] = []
|
||||
for page in doc:
|
||||
page_counts.append(len(page.get_text("text").strip()))
|
||||
|
||||
total = sum(page_counts)
|
||||
median = _median_chars_per_page(page_counts)
|
||||
|
||||
if total < settings.pdf_min_total_chars or median < settings.pdf_min_median_chars_per_page:
|
||||
raise DocumentProcessingError(
|
||||
"PDF has no usable text layer (likely a scanned/image PDF). "
|
||||
"OCR is not enabled for this benchmarker — only Text PDFs are supported. "
|
||||
f"(total_chars={total}, median_chars_per_page={median:.0f}; "
|
||||
f"need total>={settings.pdf_min_total_chars} and "
|
||||
f"median>={settings.pdf_min_median_chars_per_page})"
|
||||
)
|
||||
|
||||
|
||||
def _line_font_meta(line: dict) -> tuple[str, float, bool]:
|
||||
"""Return (text, max_font_size, any_bold) for a dict-line."""
|
||||
spans = line.get("spans") or []
|
||||
parts: list[str] = []
|
||||
max_size = 0.0
|
||||
any_bold = False
|
||||
for span in spans:
|
||||
text = (span.get("text") or "").strip()
|
||||
if text:
|
||||
parts.append(span.get("text") or "")
|
||||
size = float(span.get("size") or 0)
|
||||
if size > max_size:
|
||||
max_size = size
|
||||
# flags bit 4 (16) = bold in PyMuPDF
|
||||
if int(span.get("flags") or 0) & 2**4:
|
||||
any_bold = True
|
||||
return ("".join(parts).strip(), max_size, any_bold)
|
||||
|
||||
|
||||
def _extract_blocks(doc: fitz.Document) -> list[TextBlock]:
|
||||
"""Extract reading-order text blocks with font metadata (tables flattened)."""
|
||||
blocks: list[TextBlock] = []
|
||||
for page_index, page in enumerate(doc):
|
||||
page_dict = page.get_text("dict", flags=fitz.TEXT_PRESERVE_WHITESPACE)
|
||||
for block in page_dict.get("blocks") or []:
|
||||
if block.get("type") != 0: # 0 = text
|
||||
continue
|
||||
for line in block.get("lines") or []:
|
||||
text, font_size, bold = _line_font_meta(line)
|
||||
if not text:
|
||||
continue
|
||||
blocks.append(
|
||||
TextBlock(
|
||||
style_name="Normal",
|
||||
text=text,
|
||||
font_size=font_size,
|
||||
bold=bold,
|
||||
page=page_index + 1,
|
||||
locked=False,
|
||||
)
|
||||
)
|
||||
return blocks
|
||||
|
||||
|
||||
def _normalize(s: str) -> str:
|
||||
return " ".join(s.split()).casefold()
|
||||
|
||||
|
||||
def _apply_outline_headings(doc: fitz.Document, blocks: list[TextBlock]) -> int:
|
||||
"""Mark blocks that match PDF outline/bookmark titles. Returns match count."""
|
||||
toc = doc.get_toc(simple=True)
|
||||
if not toc:
|
||||
return 0
|
||||
|
||||
matched = 0
|
||||
used: set[int] = set()
|
||||
|
||||
for level, title, page in toc:
|
||||
title_norm = _normalize(title)
|
||||
if not title_norm:
|
||||
continue
|
||||
level = max(1, min(int(level), 6))
|
||||
|
||||
# Prefer same page; fall back to any unused block with matching title
|
||||
candidates = [
|
||||
(i, b)
|
||||
for i, b in enumerate(blocks)
|
||||
if i not in used and not b.locked
|
||||
]
|
||||
same_page = [(i, b) for i, b in candidates if b.page == page]
|
||||
search_order = same_page + [(i, b) for i, b in candidates if b.page != page]
|
||||
|
||||
for i, block in search_order:
|
||||
block_norm = _normalize(block.text)
|
||||
if block_norm == title_norm or title_norm in block_norm or block_norm in title_norm:
|
||||
block.style_name = f"Heading{level}"
|
||||
block.locked = True
|
||||
used.add(i)
|
||||
matched += 1
|
||||
break
|
||||
|
||||
return matched
|
||||
|
||||
|
||||
def _apply_font_headings(blocks: list[TextBlock]) -> int:
|
||||
"""Mark unlocked blocks as headings from font size clusters vs body size."""
|
||||
sizes = [b.font_size for b in blocks if b.font_size > 0]
|
||||
if len(sizes) < 3:
|
||||
return 0
|
||||
|
||||
body_size = statistics.median(sizes)
|
||||
if body_size <= 0:
|
||||
return 0
|
||||
|
||||
# Distinct sizes clearly above body text
|
||||
heading_sizes = sorted(
|
||||
{round(s, 1) for s in sizes if s >= body_size * 1.2},
|
||||
reverse=True,
|
||||
)
|
||||
if not heading_sizes:
|
||||
# Bold + short lines at/near body size as weak H2
|
||||
upgraded = 0
|
||||
for block in blocks:
|
||||
if block.locked or block.style_name != "Normal":
|
||||
continue
|
||||
if block.bold and len(block.text) <= 150:
|
||||
block.style_name = "Heading2"
|
||||
upgraded += 1
|
||||
return upgraded
|
||||
|
||||
size_to_level: dict[float, int] = {}
|
||||
for idx, size in enumerate(heading_sizes):
|
||||
size_to_level[size] = 1 if idx == 0 else 2
|
||||
|
||||
upgraded = 0
|
||||
for block in blocks:
|
||||
if block.locked or block.style_name != "Normal":
|
||||
continue
|
||||
if len(block.text) > 150:
|
||||
continue
|
||||
rounded = round(block.font_size, 1)
|
||||
level = size_to_level.get(rounded)
|
||||
if level is None:
|
||||
continue
|
||||
block.style_name = f"Heading{level}"
|
||||
upgraded += 1
|
||||
|
||||
return upgraded
|
||||
|
||||
|
||||
def parse_pdf(file_path: str | Path) -> ParseResult:
|
||||
"""Parse a Text PDF into DocumentTree + markdown.
|
||||
|
||||
Raises DocumentProcessingError for missing files, Scanned PDFs, or empty text.
|
||||
"""
|
||||
path = Path(file_path)
|
||||
if not path.exists():
|
||||
raise DocumentProcessingError(f"File not found: {path}")
|
||||
if path.suffix.lower() != ".pdf":
|
||||
raise DocumentProcessingError(f"Not a PDF: {path.suffix}")
|
||||
|
||||
try:
|
||||
doc = fitz.open(path)
|
||||
except Exception as exc:
|
||||
raise DocumentProcessingError(f"Failed to open PDF: {exc}") from exc
|
||||
|
||||
try:
|
||||
if doc.page_count < 1:
|
||||
raise DocumentProcessingError("PDF has no pages")
|
||||
|
||||
_enforce_text_layer_gate(doc)
|
||||
blocks = _extract_blocks(doc)
|
||||
if not blocks:
|
||||
raise DocumentProcessingError("PDF contains no extractable text blocks")
|
||||
|
||||
_apply_outline_headings(doc, blocks)
|
||||
_apply_font_headings(blocks)
|
||||
|
||||
# Gap-fill remaining Normal blocks (outline titles stay locked).
|
||||
# When outline/font found nothing, this is a full heuristic pass.
|
||||
apply_heading_heuristics(blocks, only_when_no_headings=False)
|
||||
|
||||
root = build_document_tree(blocks)
|
||||
tree = DocumentTree(root=root)
|
||||
markdown = tree_to_markdown(root)
|
||||
plain_text = "\n".join(b.text for b in blocks if b.text)
|
||||
|
||||
return ParseResult(
|
||||
tree=tree,
|
||||
markdown=markdown,
|
||||
plain_text=plain_text,
|
||||
paragraph_count=len(blocks),
|
||||
)
|
||||
finally:
|
||||
doc.close()
|
||||
@@ -1,8 +1,8 @@
|
||||
"""Document and strategy API routes.
|
||||
|
||||
Endpoints:
|
||||
POST /documents Upload a .docx file
|
||||
POST /documents/{id}/process Run chunking strategies (stub until Phase 2)
|
||||
POST /documents Upload a .docx / .doc / .pdf file
|
||||
POST /documents/{id}/process Run chunking strategies
|
||||
DELETE /documents/{id} Remove document + vectors
|
||||
GET /strategies List available chunking strategies
|
||||
"""
|
||||
@@ -10,7 +10,6 @@ Endpoints:
|
||||
from fastapi import APIRouter, File, UploadFile
|
||||
|
||||
from src.core.exceptions import DocumentProcessingError
|
||||
from src.core.models import PaginatedResponse, StrategyName
|
||||
from src.documents.models import (
|
||||
DeleteResponse,
|
||||
DocumentResponse,
|
||||
@@ -20,6 +19,7 @@ from src.documents.models import (
|
||||
StrategiesResponse,
|
||||
StrategyInfo,
|
||||
)
|
||||
from src.documents.parser import SUPPORTED_SUFFIXES
|
||||
from src.documents import service
|
||||
|
||||
router = APIRouter()
|
||||
@@ -39,11 +39,15 @@ async def list_documents(offset: int = 0, limit: int = 50):
|
||||
|
||||
@router.post("/documents", response_model=DocumentResponse, status_code=201)
|
||||
async def upload_document(file: UploadFile = File(...)):
|
||||
"""Upload a .docx file. Parses it, stores the document tree in SQLite."""
|
||||
"""Upload a document. Parses it, stores the document tree in SQLite."""
|
||||
if not file.filename:
|
||||
raise DocumentProcessingError("No filename provided")
|
||||
if not file.filename.lower().endswith((".docx", ".doc")):
|
||||
raise DocumentProcessingError("Only .docx and .doc files are supported")
|
||||
|
||||
lower = file.filename.lower()
|
||||
if not any(lower.endswith(s) for s in SUPPORTED_SUFFIXES):
|
||||
raise DocumentProcessingError(
|
||||
f"Only {', '.join(SUPPORTED_SUFFIXES)} files are supported"
|
||||
)
|
||||
|
||||
content = await file.read()
|
||||
if not content:
|
||||
@@ -61,10 +65,7 @@ async def upload_document(file: UploadFile = File(...)):
|
||||
|
||||
@router.post("/documents/{doc_id}/process", response_model=ProcessResponse)
|
||||
async def process_document(doc_id: str, request: ProcessRequest):
|
||||
"""Run selected chunking strategies on an uploaded document.
|
||||
|
||||
Currently returns stub results until Phase 2 implements the strategies.
|
||||
"""
|
||||
"""Run selected chunking strategies on an uploaded document."""
|
||||
return service.process_document(doc_id, request)
|
||||
|
||||
|
||||
|
||||
@@ -16,7 +16,7 @@ from src.documents.models import (
|
||||
ProcessResponse,
|
||||
StrategyResult,
|
||||
)
|
||||
from src.documents.parser import parse_docx
|
||||
from src.documents.parser import parse_document, SUPPORTED_SUFFIXES
|
||||
from src.storage import sqlite as db
|
||||
from src.storage import qdrant as qdr
|
||||
|
||||
@@ -49,17 +49,22 @@ STRATEGY_DESCRIPTIONS: dict[StrategyName, str] = {
|
||||
# ── Upload ─────────────────────────────────────────────────────────
|
||||
|
||||
def upload_document(filename: str, file_bytes: bytes) -> dict[str, Any]:
|
||||
"""Parse a .docx upload, store in SQLite, return document record."""
|
||||
"""Parse an uploaded document, store in SQLite, return document record."""
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
suffix = Path(filename).suffix or ".docx"
|
||||
suffix = Path(filename).suffix.lower()
|
||||
if suffix not in SUPPORTED_SUFFIXES:
|
||||
raise DocumentProcessingError(
|
||||
f"Only {', '.join(SUPPORTED_SUFFIXES)} files are supported"
|
||||
)
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
|
||||
tmp.write(file_bytes)
|
||||
tmp_path = tmp.name
|
||||
|
||||
try:
|
||||
result = parse_docx(tmp_path)
|
||||
result = parse_document(tmp_path)
|
||||
finally:
|
||||
Path(tmp_path).unlink(missing_ok=True)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user