Chunking Enterprise Documents Without Losing Meaning

 PDF / DOCX / HTML
        │
        ▼
 ┌────────────────────────┐
 │ Layout-aware extraction│  (its own problem: OCR, reading order, tables)
 └───────────┬────────────┘
             │ Markdown
             ▼
 ┌────────────────────────┐     ┌──────────────────────────────┐
 │ Block parser           │ ──▶ │ Section tree (heading paths) │
 └────────────────────────┘     └──────────────┬───────────────┘
                                               │
                     ┌─────────────────────────┴─────────────────────┐
                     ▼                                               ▼
        ┌─────────────────────────┐                    ┌──────────────────────────┐
        │ Child chunks            │                    │ Parent sections          │
        │ breadcrumb + blocks     │                    │ full text, keyed by id   │
        │ + metadata + parent_id  │                    │ (document store)         │
        └────────────┬────────────┘                    └────────────▲─────────────┘
                     │ embed + index                                │ fetch by parent_id
                     ▼                                              │
 query ──▶ search children ──▶ top-k child hits ──▶ dedupe parent ids

import re
from dataclasses import dataclass, field
from typing import Callable
 
TokenCounter = Callable[[str], int]
 
HEADING = re.compile(r"^(#{1,6})\s+(.+?)\s*#*\s*$")
LIST_ITEM = re.compile(r"^\s*(?:[-*+]|\d+[.)])\s+")
FENCE = "`" * 3  # a code fence, spelled this way so this file can live inside Markdown
 
SAMPLE_POLICY = """# Travel Policy
 
Applies to all employees and contractors traveling on company business.
 
## Domestic Travel
 
### Booking
 
Book flights and hotels through the corporate travel portal at least 7 days
in advance. Economy class is the default for flights under 6 hours.
 
## International Travel
 
### Per diem
 
Per diem covers meals and incidentals. Lodging is reimbursed separately
against receipts, up to the cap for the destination city.
 
| City | Daily limit (USD) | Lodging cap (USD) |
|---|---|---|
| London | 95 | 320 |
| Tokyo | 90 | 280 |
| Sao Paulo | 70 | 210 |
| New York | 105 | 350 |
 
Exception: trips longer than 30 consecutive days receive 75% of the daily
limit from day 31 onward.
 
### Visa and approvals
 
Before booking any international trip, follow this procedure:
 
1. Check the entry requirements for your passport on the official government site.
2. Request a business invitation letter from the host office.
3. Get written approval from your director and attach it to the trip request.
4. Submit the visa application and upload the receipt to the travel portal.
5. Book flights only after the visa is issued.
 
## Expense Reporting
 
Submit expenses within 15 days of returning. Receipts are mandatory for any
item above 25 USD.
"""
 
 
def approx_tokens(text: str) -> int:
    # Rough English heuristic (about 0.75 words per token). Swap in a real tokenizer.
    return max(1, round(len(text.split()) * 4 / 3))
 
 
@dataclass
class Block:
    kind: str  # heading, paragraph, list, table, code
    text: str
    level: int = 0
 
 
@dataclass
class Section:
    section_id: str
    heading_path: list[str]
    blocks: list[Block] = field(default_factory=list)
 
    @property
    def breadcrumb(self) -> str:
        return " > ".join(self.heading_path)
 
    @property
    def text(self) -> str:
        body = "\n\n".join(block.text for block in self.blocks)
        return f"{self.breadcrumb}\n\n{body}".strip()
 
 
@dataclass
class Chunk:
    chunk_id: str
    parent_id: str
    heading_path: list[str]
    text: str
    token_count: int
    metadata: dict[str, str]
 
 
def _starts_block(line: str) -> bool:
    return bool(
        HEADING.match(line)
        or line.startswith(FENCE)
        or line.lstrip().startswith("|")
        or LIST_ITEM.match(line)
    )
 
 
def parse_blocks(markdown: str) -> list[Block]:
    lines = markdown.splitlines()
    blocks: list[Block] = []
    i = 0
    while i < len(lines):
        line = lines[i]
        if not line.strip():
            i += 1
        elif line.startswith(FENCE):
            j = i + 1
            while j < len(lines) and not lines[j].startswith(FENCE):
                j += 1
            blocks.append(Block("code", "\n".join(lines[i : j + 1])))
            i = j + 1
        elif match := HEADING.match(line):
            blocks.append(Block("heading", match.group(2), level=len(match.group(1))))
            i += 1
        elif line.lstrip().startswith("|"):
            j = i
            while j < len(lines) and lines[j].lstrip().startswith("|"):
                j += 1
            blocks.append(Block("table", "\n".join(lines[i:j])))
            i = j
        elif LIST_ITEM.match(line):
            j = i + 1
            # Continuation lines (indented) stay with their list
            while j < len(lines) and lines[j].strip() and (
                LIST_ITEM.match(lines[j]) or lines[j].startswith((" ", "\t"))
            ):
                j += 1
            blocks.append(Block("list", "\n".join(lines[i:j])))
            i = j
        else:
            j = i + 1
            while j < len(lines) and lines[j].strip() and not _starts_block(lines[j]):
                j += 1
            blocks.append(Block("paragraph", " ".join(l.strip() for l in lines[i:j])))
            i = j
    return blocks
 
 
def build_sections(blocks: list[Block], doc_id: str) -> list[Section]:
    sections: list[Section] = []
    stack: list[tuple[int, str]] = []  # (level, title) of the open headings
    current = Section(f"{doc_id}#s0", [doc_id])
    for block in blocks:
        if block.kind == "heading":
            if current.blocks:
                sections.append(current)
            while stack and stack[-1][0] >= block.level:
                stack.pop()
            stack.append((block.level, block.text))
            path = [title for _, title in stack]
            current = Section(f"{doc_id}#s{len(sections) + 1}", path)
        else:
            current.blocks.append(block)
    if current.blocks:
        sections.append(current)
    return sections
 
 
def _group(units: list[str], budget: int, count: TokenCounter, prefix: str = "") -> list[str]:
    """Greedily packs units into groups under budget, repeating prefix in each group."""
    groups: list[str] = []
    current: list[str] = []
    for unit in units:
        candidate = "\n".join([prefix, *current, unit]).strip()
        if current and count(candidate) > budget:
            groups.append("\n".join([prefix, *current]).strip())
            current = []
        current.append(unit)
    if current:
        groups.append("\n".join([prefix, *current]).strip())
    return groups
 
 
def split_block(block: Block, budget: int, count: TokenCounter) -> list[str]:
    if count(block.text) <= budget:
        return [block.text]  # tables, lists and code stay intact when they fit
    lines = block.text.splitlines()
    if block.kind == "table":
        header = "\n".join(lines[:2])  # header row + separator, repeated in every piece
        return _group(lines[2:], budget, count, prefix=header)
    if block.kind == "list":
        items: list[str] = []
        for line in lines:
            if LIST_ITEM.match(line) or not items:
                items.append(line)
            else:
                items[-1] += "\n" + line
        return _group(items, budget, count)
    if block.kind == "code":
        fence, body = lines[0], lines[1:-1]
        return [f"{fence}\n{group}\n{FENCE}" for group in _group(body, budget, count)]
    sentences = re.split(r"(?<=[.!?])\s+", block.text)
    return [group.replace("\n", " ") for group in _group(sentences, budget, count)]
 
 
def chunk_document(
    markdown: str,
    doc_id: str,
    max_tokens: int = 200,
    count_tokens: TokenCounter = approx_tokens,
    metadata: dict[str, str] | None = None,
) -> tuple[list[Chunk], dict[str, Section]]:
    sections = build_sections(parse_blocks(markdown), doc_id)
    chunks: list[Chunk] = []
    for section in sections:
        header = f"{section.breadcrumb}\n\n"
        budget = max_tokens - count_tokens(header)
        units = [piece for block in section.blocks for piece in split_block(block, budget, count_tokens)]
        bodies: list[str] = []
        current: list[str] = []
        for unit in units:
            if current and count_tokens(header + "\n\n".join([*current, unit])) > max_tokens:
                bodies.append("\n\n".join(current))
                current = []
            current.append(unit)
        if current:
            bodies.append("\n\n".join(current))
        for n, body in enumerate(bodies):
            text = header + body
            chunks.append(
                Chunk(
                    chunk_id=f"{section.section_id}-c{n}",
                    parent_id=section.section_id,
                    heading_path=section.heading_path,
                    text=text,
                    token_count=count_tokens(text),
                    metadata={"doc_id": doc_id, "section": section.breadcrumb, **(metadata or {})},
                )
            )
    return chunks, {section.section_id: section for section in sections}
 
 
if __name__ == "__main__":
    chunks, parents = chunk_document(
        SAMPLE_POLICY,
        doc_id="travel-policy",
        max_tokens=120,
        metadata={"version": "3.2", "tenant": "acme"},
    )
    for chunk in chunks:
        print(f"--- {chunk.chunk_id} (parent {chunk.parent_id}, ~{chunk.token_count} tokens)")
        print(chunk.text, end="\n\n")

import tiktoken
 
from chunker import SAMPLE_POLICY, chunk_document
 
# Prefer the tokenizer of the model that will consume the chunk
encoding = tiktoken.get_encoding("cl100k_base")
 
 
def count_tokens(text: str) -> int:
    return len(encoding.encode(text))
 
 
chunks, parents = chunk_document(
    SAMPLE_POLICY, doc_id="travel-policy", max_tokens=120, count_tokens=count_tokens
)
for chunk in chunks:
    print(chunk.chunk_id, chunk.token_count)

import math
import re
from collections import Counter
 
from chunker import SAMPLE_POLICY, Chunk, Section, chunk_document
 
STOP_WORDS = {"a", "an", "the", "is", "are", "i", "do", "what", "for", "in", "of", "to", "my", "on", "how"}
 
 
def terms(text: str) -> list[str]:
    return [t for t in re.findall(r"[a-z0-9]+", text.lower()) if t not in STOP_WORDS]
 
 
def score(query: str, text: str) -> float:
    # Stand-in for your vector or hybrid index, so the example runs without a model
    counts = Counter(terms(text))
    return sum(1 + math.log(counts[t]) for t in set(terms(query)) if counts[t])
 
 
def search_children(
    query: str,
    chunks: list[Chunk],
    k: int = 4,
    filters: dict[str, str] | None = None,
) -> list[Chunk]:
    filters = filters or {}
    allowed = [c for c in chunks if all(c.metadata.get(key) == value for key, value in filters.items())]
    scored = [(score(query, c.text), c) for c in allowed]
    ranked = sorted((pair for pair in scored if pair[0] > 0), key=lambda pair: pair[0], reverse=True)
    return [chunk for _, chunk in ranked[:k]]
 
 
def small_to_big(
    query: str,
    chunks: list[Chunk],
    parents: dict[str, Section],
    k: int = 4,
    max_parents: int = 2,
    filters: dict[str, str] | None = None,
) -> list[Section]:
    parent_ids: list[str] = []
    for chunk in search_children(query, chunks, k=k, filters=filters):
        if chunk.parent_id not in parent_ids:  # dedupe, keep the rank of the best child
            parent_ids.append(chunk.parent_id)
    return [parents[pid] for pid in parent_ids[:max_parents]]
 
 
if __name__ == "__main__":
    # Small children for precise matching, whole sections for the model
    chunks, parents = chunk_document(
        SAMPLE_POLICY, doc_id="travel-policy", max_tokens=45, metadata={"tenant": "acme"}
    )
    for question in ["What is the per diem in Tokyo?", "What comes after director approval?"]:
        hits = search_children(question, chunks, filters={"tenant": "acme"})
        print(f"Q: {question}")
        print("   matched children:", [c.chunk_id for c in hits])
        for section in small_to_big(question, chunks, parents, filters={"tenant": "acme"}):
            print(f"\n[{section.section_id}]\n{section.text}\n")

from chunker import SAMPLE_POLICY, Chunk, chunk_document
from retrieval import search_children, small_to_big
 
 
def fixed_size_chunks(text: str, size: int = 200, overlap: int = 40) -> list[Chunk]:
    """The naive baseline: character windows that ignore structure."""
    chunks = []
    for n, start in enumerate(range(0, len(text), size - overlap)):
        window = text[start : start + size]
        chunks.append(Chunk(f"fixed-{n}", f"fixed-{n}", [], window, len(window.split()), {}))
    return chunks
 
 
# Each question lists the facts that must ALL reach the model for a correct answer
EVAL_SET = {
    "What is the daily limit in Tokyo for a 45 day trip?": [
        "Daily limit (USD)", "| Tokyo | 90", "75% of the daily",
    ],
    "What is the lodging cap in New York?": ["Lodging cap (USD)", "| New York | 105 | 350"],
    "What are the steps before I book an international flight?": [
        "1. Check the entry", "3. Get written approval", "5. Book flights only after",
    ],
    "How long do I have to submit expenses?": ["within 15 days"],
}
 
 
def fact_recall(context: str, facts: list[str]) -> float:
    return sum(fact in context for fact in facts) / len(facts)
 
 
def evaluate(name: str, retrieve) -> None:
    scores = [fact_recall(retrieve(q), facts) for q, facts in EVAL_SET.items()]
    print(f"{name:<22} fact recall = {sum(scores) / len(scores):.2f}")
 
 
if __name__ == "__main__":
    k = 2
    fixed = fixed_size_chunks(SAMPLE_POLICY)
    children, parents = chunk_document(SAMPLE_POLICY, doc_id="travel-policy", max_tokens=45)
 
    evaluate("fixed-size windows", lambda q: "\n".join(c.text for c in search_children(q, fixed, k=k)))
    evaluate("structural children", lambda q: "\n".join(c.text for c in search_children(q, children, k=k)))
    evaluate("small-to-big parents", lambda q: "\n".join(s.text for s in small_to_big(q, children, parents, k=k)))