Scale and measure

So far the tree has one document under its root. Scaling to many documents is a one-line change: make the root a virtual container whose children are document-level trees. The agent learns to pick the right document before walking sections.

Multi-document tree structure

A virtual root holds per-document subtrees. The agent picks the document, then the section.

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

def build_library(pdf_paths: list[str]) -> TreeNode:
    """Wrap multiple document trees under a virtual root."""
    per_doc_trees = [parse_pdf(p) for p in pdf_paths]

    root = TreeNode(
        id="library",
        title="Research Library",
        level=0,
        page_start=1,
        page_end=1,
        content="",
        heading_type="library",
        summary="A multi-document research corpus",
    )
    # Each document tree becomes a top-level child
    root.children = [t.root if hasattr(t, "root") else t for t in per_doc_trees]
    return root

The library root is a TreeNode whose children are full document trees. Every downstream piece (analyze, descend, retrieve) already handles arbitrary depth, so no other code changes.

The agent already logs every LLM call with its latency. Turning that into a cost and latency summary is just arithmetic. For each question, sum the input and output tokens, multiply by the model pricing, and show the total.

retriever.py
python
# Every LLM call appends an entry to state["call_log"]
entry = {
    "call_number":    call_num,
    "call_type":      "navigate",
    "node_id":        node.id,
    "node_title":     node.title,
    "depth":          depth,
    "confidence":     conf,
    "should_descend": descend,
    "target_child":   child_id,
    "reasoning":      reasoning,
    "latency_s":      round(elapsed, 3),
}

This entry is appended to the state via the operator.add reducer. At the end of a run, call_log holds a complete timeline of every decision with timing. Add token counts from the LLM response and you have everything you need for a cost report.

AI prompt: Try it: analyze your own call log

Loading practice…

Quiz: Quiz

Loading practice…