core/chunking.py

"""Section-based chunking of Nemotron Parse output (parsed approach only).

Pure functions, no model dependencies — unit-testable without a GPU.

Input: pages = [{"page": n, "elements": [{"class", "bbox", "text",
"description"?}]}] in reading order, where Picture/Table elements may carry a
"description" generated by MiniCPM at ingest.

Chunking follows the "by title" strategy: an element classified Title or
Section-header closes the current section chunk. Size guards both ways:
  - a section still under min_chars merges forward (its heading folds into
    the running text), so sparse pages don't become useless tiny chunks;
  - a section over max_chars splits at element boundaries with the heading
    repeated on each piece.
Sections span page boundaries; each chunk records every page it draws from
so parent-document retrieval can map chunks back to page images.

Besides section chunks, every described Picture and every Table also becomes
a standalone chunk (sharp matches on one diagram or spec table), and its
description is spliced inline into the enclosing section's text.
"""

from __future__ import annotations

from core.constants import SECTION_MAX_CHARS, SECTION_MIN_CHARS

HEADING_CLASSES = {"Title", "Section-header"}
SKIP_CLASSES = {"Page-header", "Page-footer"}
FIGURE_CLASS = "Picture"
TABLE_CLASS = "Table"
CAPTION_CLASS = "Caption"


def build_chunks(
    pages: list[dict],
    min_chars: int = SECTION_MIN_CHARS,
    max_chars: int = SECTION_MAX_CHARS,
) -> list[dict]:
    """Return chunk dicts: {"type": "section"|"figure"|"table", "heading",
    "pages": [...] (section) or "page"/"bbox" (figure/table), "text"}.
    The embedded text always starts with the section heading when one is known.
    """
    chunks: list[dict] = []
    heading = ""
    parts: list[str] = []
    part_pages: list[int] = []
    size = 0

    def flush() -> None:
        nonlocal parts, part_pages, size
        text = "\n\n".join(parts).strip()
        if text:
            chunks.append(
                {
                    "type": "section",
                    "heading": heading,
                    "pages": sorted(set(part_pages)),
                    "text": f"{heading}\n\n{text}" if heading else text,
                }
            )
        parts, part_pages, size = [], [], 0

    def add_part(page: int, text: str) -> None:
        nonlocal size
        if not text:
            return
        if size and size + len(text) > max_chars:
            flush()  # oversized section: split; heading repeats on the next piece
        parts.append(text)
        part_pages.append(page)
        size += len(text)

    def with_heading(text: str) -> str:
        return f"{heading}\n\n{text}" if heading else text

    for pg in pages:
        n = pg["page"]
        for el in pg["elements"]:
            cls = el.get("class", "")
            text = (el.get("text") or "").strip()
            if cls in SKIP_CLASSES:
                continue

            if cls in HEADING_CLASSES:
                if not text:
                    continue
                if size >= min_chars:
                    flush()
                    heading = text
                    part_pages.append(n)  # heading's own page counts for the section
                elif not parts:
                    # consecutive headings (e.g. Title then Section-header)
                    # combine into a breadcrumb
                    heading = f"{heading} — {text}" if heading else text
                    part_pages.append(n)
                else:
                    add_part(n, text)  # tiny section merges forward
                continue

            if cls == FIGURE_CLASS:
                desc = (el.get("description") or "").strip()
                if desc:
                    chunks.append(
                        {
                            "type": "figure",
                            "heading": heading,
                            "page": n,
                            "bbox": el.get("bbox"),
                            "text": with_heading(f"Figure on page {n}: {desc}"),
                        }
                    )
                    add_part(n, f"[Figure: {desc}]")
                continue

            if cls == TABLE_CLASS:
                desc = (el.get("description") or "").strip()
                if not text and not desc:
                    continue
                body = f"Table on page {n}"
                if desc:
                    body += f": {desc}"
                if text:
                    body += f"\n\n{text}"
                chunks.append(
                    {
                        "type": "table",
                        "heading": heading,
                        "page": n,
                        "bbox": el.get("bbox"),
                        "text": with_heading(body),
                    }
                )
                add_part(n, f"[Table: {desc}]\n{text}" if desc else text)
                continue

            add_part(n, text)

    flush()  # the last section has no next section to merge into
    return chunks