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  # uma cerca de código, escrita assim para este arquivo poder viver dentro de 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:
    # Heurística grosseira para inglês (cerca de 0,75 palavra por token). Troque por um tokenizador de verdade.
    return max(1, round(len(text.split()) * 4 / 3))
 
 
@dataclass
class Block:
    kind: str  # heading, parágrafo, lista, tabela, código
    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
            # Linhas de continuação (indentadas) ficam com a lista delas
            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]] = []  # (nível, título) dos headings abertos
    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]:
    """Empacota unidades de forma gulosa em grupos dentro do orçamento, repetindo o prefixo em cada grupo."""
    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]  # tabelas, listas e código ficam inteiros quando cabem
    lines = block.text.splitlines()
    if block.kind == "table":
        header = "\n".join(lines[:2])  # linha de cabeçalho + separador, repetidos em cada pedaço
        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
 
# Prefira o tokenizador do modelo que vai consumir o 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:
    # Substituto do seu índice vetorial ou híbrido, para o exemplo rodar sem modelo
    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:  # deduplica, mantendo a posição do melhor filho
            parent_ids.append(chunk.parent_id)
    return [parents[pid] for pid in parent_ids[:max_parents]]
 
 
if __name__ == "__main__":
    # Filhos pequenos para o match preciso, seções inteiras para o modelo
    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]:
    """A baseline ingênua: janelas de caracteres que ignoram a estrutura."""
    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
 
 
# Cada pergunta lista os fatos que precisam TODOS chegar ao modelo para uma resposta correta
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)))