Tree caching to JSON

Parsing a PDF with PyMuPDF4LLM takes ten to thirty seconds for a typical research paper. That is fine for the first run, but painful if every question triggers a re-parse. We solve this by serializing the tree to JSON on first build and loading it from disk on every subsequent run.

Tree caching flow

First run parses and caches. Every subsequent run loads from disk in milliseconds.

main.py
python
from dataclasses import asdict
import json
from pathlib import Path
from tree import parse_pdf, TreeNode

TREE_CACHE_PATH = Path("results/document_tree.json")

def get_tree() -> TreeNode:
    TREE_CACHE_PATH.parent.mkdir(parents=True, exist_ok=True)

    if TREE_CACHE_PATH.exists():
        with open(TREE_CACHE_PATH) as f:
            data = json.load(f)
        tree = dict_to_treenode(data.get("root", data))
        return tree

    tree = parse_pdf(str(PDF_PATH))
    with open(TREE_CACHE_PATH, "w") as f:
        json.dump(asdict(tree), f, indent=2, default=str)
    return tree

asdict from dataclasses walks the entire tree and produces a plain dict. json.dump writes it. On the next run, we detect the cache file and skip parsing entirely.

main.py
python
def dict_to_treenode(data: dict) -> TreeNode:
    """Recursively reconstruct TreeNode from dictionary."""
    children = [
        dict_to_treenode(child) for child in data.get("children", [])
    ]
    data_copy = data.copy()
    data_copy["children"] = children
    return TreeNode(**data_copy)

Reconstruction walks the dict top down, recursively rebuilding children first, then instantiating the parent with the reconstructed list. This preserves the exact runtime shape of the original tree.

JSON is human readable, which matters when you are debugging why a particular section got picked. You can open the cache file, grep for a section title, and see exactly what the agent is working with. Pickle is faster to write but opaque, and it ties the cache format to a specific Python version.

Ordering exercise: Put the caching logic in the right order

Loading practice…