Keep Page Coordinates with Your RAG Chunks

%pip install -q pymupdf langchain-core llama-index-core

import pymupdf
doc = pymupdf.open()
sections = [
 ("Acme Corporation Annual Report 2025", None),
 ("1. Financial highlights", "Net revenue increased 12% to $4.1 billion, driven by growth in the enterprise software segment. Operating margin improved to 18.4%, compared with 16.9% in the prior year. The company returned $350 million to shareholders through dividends and share repurchases."),
 ("2. Outlook", "Management expects revenue growth of 8% to 10% in fiscal 2026, with continued investment in cloud infrastructure and data centers in Europe and Asia. Capital expenditure is expected to reach $600 million."),
 ("3. Risk factors", "The company depends on a small number of large customers. In fiscal 2025, the three largest customers accounted for 31% of net revenue. The loss of any of these customers could materially affect results."),
 ("4. Supply chain", "Supply chain disruptions in 2024 delayed shipments of hardware appliances by up to six weeks. These delays did not recur in 2025, and inventory levels returned to normal by the second quarter."),
]
page = doc.new_page(); y = 72
for i,(h,body) in enumerate(sections):
    if i == 3:
        page = doc.new_page(); y = 72
    size = 18 if i == 0 else 14
    page.insert_text((72, y+size), h, fontsize=size, fontname="hebo"); y += size + 14
    if body:
        r = pymupdf.Rect(72, y, 400, y+200)
        rc = page.insert_textbox(r, body, fontsize=11, fontname="helv")
        y += 200 - rc + 18
doc.save("acme_report.pdf")

import pymupdf

doc = pymupdf.open("acme_report.pdf")
for page_dict in doc.to_markdown(page_chunks=True):
    text = page_dict["text"]
    for box in page_dict["page_boxes"]:
        start, stop = box["pos"]
        print(page_dict["metadata"]["page_number"], box["class"], box["bbox"], repr(text[start:stop][:45]))

import json
import pymupdf


def chunks_with_coordinates(path, max_chars=1500):
    """Split a PDF into Markdown chunks that remember where each part came from.

    A new chunk starts at each section header, or when a chunk would exceed max_chars.
    Each chunk records the page number and bbox of every layout box it contains.
    """
    doc = pymupdf.open(path)
    chunks, current = [], None
    for page_dict in doc.to_markdown(page_chunks=True):
        page_number = page_dict["metadata"]["page_number"]      # 1-based
        text = page_dict["text"]
        for box in page_dict["page_boxes"]:
            start, stop = box["pos"]
            box_text = text[start:stop]
            is_header = box["class"] in ("title", "section-header")
            new_section = is_header and current and current["has_body"]
            too_long = current and len(current["text"]) + len(box_text) > max_chars
            if current is None or new_section or too_long:
                current = {"text": "", "boxes": [], "has_body": False}
                chunks.append(current)
            current["text"] += box_text
            current["has_body"] |= not is_header
            current["boxes"].append({"page": page_number, "bbox": list(box["bbox"])})
    return chunks

def chunk_metadata(path, chunk):
    """Flat metadata that vector stores accept: strings and numbers only."""
    return {
        "source": path,
        "page": chunk["boxes"][0]["page"],                 # first page, for filtering
        "boxes": json.dumps(chunk["boxes"]),               # every page and bbox
    }

from langchain_core.documents import Document

path = "acme_report.pdf"
documents = [
    Document(page_content=c["text"], metadata=chunk_metadata(path, c))
    for c in chunks_with_coordinates(path)
]

from llama_index.core.schema import TextNode

nodes = [
    TextNode(
        text=c["text"],
        metadata=chunk_metadata(path, c),
        excluded_embed_metadata_keys=["boxes"],   # keep coordinates out of the embedding
        excluded_llm_metadata_keys=["boxes"],     # and out of the prompt
    )
    for c in chunks_with_coordinates(path)
]

def chunk_images(metadata, dpi=150, margin=12):
    """Render the page regions a chunk came from: one image per page it touches."""
    doc = pymupdf.open(metadata["source"])
    regions = {}
    for b in json.loads(metadata["boxes"]):
        regions.setdefault(b["page"], pymupdf.Rect())
        regions[b["page"]] |= pymupdf.Rect(b["bbox"])
    images = []
    for page_number, rect in regions.items():
        page = doc[page_number - 1]
        clip = (rect + (-margin, -margin, margin, margin)) & page.rect
        images.append(page.get_pixmap(clip=clip, dpi=dpi))
    return images

# In a real pipeline this comes from your retriever, e.g. retriever.invoke("How much did revenue grow?")[0]
hit = next(d for d in documents if "Net revenue" in d.page_content)
for i, pix in enumerate(chunk_images(hit.metadata)):
    pix.save(f"source_{i}.png")

from IPython.display import Image, display
display(Image("source_0.png"))

import json

def locate_in_chunk(doc, metadata, quote):
    """Find a quote inside a retrieved chunk's regions only."""
    for b in json.loads(metadata["boxes"]):
        page = doc[b["page"] - 1]
        quads = page.search_for(quote, clip=pymupdf.Rect(b["bbox"]), quads=True)
        if quads:
            return b["page"], quads
    return None, []

doc = pymupdf.open(path)
page_number, quads = locate_in_chunk(doc, hit.metadata, "Net revenue increased 12%")
print(page_number, [q.rect for q in quads])