TreeNode data model

The TreeNode is the central data structure of the entire pipeline. Every retrieval decision and every citation flows through it. We want something simple enough to serialize to JSON but rich enough to drive navigation.

tree.py
python
from dataclasses import dataclass, field
from typing import List, Optional

@dataclass
class TreeNode:
    """Hierarchical document node"""
    id: str
    title: str
    level: int                # 0=root, 1=chapter, 2=section, 3=subsection
    page_start: int
    page_end: int
    content: str
    children: List["TreeNode"] = field(default_factory=list)
    heading_type: Optional[str] = None  # "numbered", "unnumbered", "roman", ...
    summary: str = ""

The node carries everything the agent needs: an id for routing, a title and summary for relevance judgment, page boundaries for citations, and children for descent.

Notice what is NOT on the node: no embeddings, no vector ids, no similarity scores. The tree is pure structure. That is what makes it cheap to build, cheap to store, and cheap to reason about.

tree.py
python
def _build_tree_from_markdown(self, markdown, page_contents, doc_name):
    lines = markdown.split("\n")

    root = TreeNode(
        id="root", title=doc_name, level=0,
        page_start=1, page_end=max(page_contents.keys()),
        content="", heading_type="root",
    )

    # Stack tracks the current path: (level, node)
    stack = [(0, root)]

    for i, line in enumerate(lines):
        stripped = line.strip()
        if stripped.startswith("#"):
            level = len(stripped.split()[0])
            title = stripped.lstrip("#").strip()
            node = TreeNode(
                id=f"{title[:20]}_{i}", title=title, level=level,
                page_start=1, page_end=1, content="",
            )
            # Pop until we find the right parent
            while stack and stack[-1][0] >= level:
                stack.pop()
            stack[-1][1].children.append(node)
            stack.append((level, node))

    return root

The builder is a simple markdown header parser with a stack. Each hash count becomes a level. Deeper levels are pushed as children of the current frame, shallower levels pop the stack.

A stack lets us walk the markdown top to bottom in one pass without re-scanning. Recursion would need to know in advance where each section ends, which means scanning ahead for the next same-level heading. The stack approach is linear time and dead simple to reason about.

Quiz: Quiz

Loading practice…