Skip to content

chunker.md

Latest commit

 

History

History
154 lines (118 loc) · 5.53 KB

chunker.md

File metadata and controls

154 lines (118 loc) · 5.53 KB
Raw

rag-chunker — reference

Structure-aware chunking for AgentContextOS. Step 1.5 introduces the Chunker SPI and ships HeadingAwareChunker as the default implementation. The chunker sits between parsers (Step 1.3 / 1.4) and the metadata enricher / embedder (Steps 1.6 / 1.8) in the ingest pipeline.

Overview

A Chunker consumes a ParsedDocument and returns an ordered list of Chunk instances. Each chunk carries:

  • content — the chunk text
  • position — 0-based monotonic index within the document
  • parent_id — ChunkId of the enclosing heading chunk (or another parent in hierarchical schemes); None at the top level
  • token_count — under the chunker's configured tokenizer
  • acl_labels / tenant_id — propagated from ctx.principal and ctx.tenant_id so PolicyEngine has what it needs at retrieval time
  • metadata — chunker-specific, includes {"kind": "heading"|"body", "level": N}

The chunker contract (rag_core.spi.chunker.Chunker):

  • Chunks share parsed.document.id as document_id and ctx.tenant_id as tenant_id.
  • position is 0-based and monotonically increasing.
  • If parent_id is set, the parent appears earlier in the returned list (no forward references).
  • Each returned chunk has content or content_ref set (the Chunk validator enforces this).

HeadingAwareChunker

from rag_chunker import HeadingAwareChunker

chunker = HeadingAwareChunker(
    max_tokens=512,        # chunk token budget
    overlap_tokens=64,     # tail bytes seeded into the next chunk
    token_counter=None,    # default: TiktokenCounter("cl100k_base")
)
chunks = await chunker.chunk(ctx, parsed)

Behaviour:

  1. Walk parsed.blocks in order.
  2. Each BlockType.heading emits one chunk; its parent_id points to the closest enclosing heading at a strictly higher level (smaller level number), or None at top level.
  3. Body blocks (paragraph, list_item, table_row, code, quote, caption, other) accumulate in a buffer until adding the next block would exceed max_tokens, at which point the buffer flushes as a chunk under the current heading.
  4. When flushing inside a section, the trailing ~overlap_tokens of text seed the next chunk's buffer. Overlap is suppressed across heading boundaries to avoid bleeding section content.
  5. Blocks larger than max_tokens are split on sentence boundaries via rag_chunker.sentence.split_sentences; sentences still too large are hard-split by token count so no body chunk ever exceeds the budget.
  6. Whitespace-only blocks are dropped.

HeadingAwareChunker.metadata on emitted chunks:

  • Heading chunks: {"kind": "heading", "level": N} where N ∈ {1..6}.
  • Body chunks: {"kind": "body"}.

TokenCounter

from rag_chunker import TiktokenCounter, TokenCounter

class TokenCounter(Protocol):
    def count(self, text: str) -> int: ...
    def name(self) -> str: ...

Default is TiktokenCounter("cl100k_base") — the encoder used by GPT-4, GPT-3.5-turbo, and most OpenAI embedding models, and a reasonable estimator for Claude as well. Pass TiktokenCounter("o200k_base") for GPT-4o-style budgets, or supply any object that implements the TokenCounter protocol (HuggingFace tokenizer wrapper, etc.).

OCR adapter

ocr_result_to_parsed_document() turns an OCRResult (Step 1.4) into a synthetic ParsedDocument so OCR'd images flow through the same chunker as native parser output. Each OCRRegion becomes one paragraph block with attributes carrying engine, ocr_confidence, bbox, and optional page. The Step 1.10 ingest pipeline will use this when stitching parser blocks with OCR'd image-page blocks.

from rag_chunker import ocr_result_to_parsed_document

parsed = ocr_result_to_parsed_document(
    ctx, ocr_result,
    document_id=DocumentId("scan-001"),
    corpus_id=CorpusId("invoices"),
    source_uri="s3://bucket/scan.png",
)
chunks = await chunker.chunk(ctx, parsed)

CLI

ragctl chunk path/to/doc.md
ragctl chunk doc.txt --max-tokens 256 --overlap-tokens 32
ragctl chunk doc.md -n 10        # show first 10 chunks

ragctl chunk parses the file (via the Step 1.3 parser registry) and runs the chunker, printing block/chunk counts, per-kind tally, and a preview of the first N chunks with parent-link arrows.

Errors

All chunker failures surface as rag_core.errors.ChunkError:

  • Tokenizer crash (unwrapped errors are re-raised as ChunkError).
  • Malformed block input where offset arithmetic fails.

HeadingAwareChunker.__init__ raises ValueError for invalid max_tokens / overlap_tokens configuration.

Internals

HeadingAwareChunker keeps a small heading stack [(level, chunk_id), …] in ascending level order. On each new heading:

  1. Flush any pending body buffer (without overlap — overlap doesn't cross heading boundaries).
  2. Pop frames with level >= new_level (closes prior sibling / deeper headings).
  3. New heading's parent_id = top of stack after pop.
  4. Push (new_level, new_chunk_id).

Body chunks always link to the top of the stack.

Extension points

Implement a custom chunker:

  1. Subclass rag_core.spi.chunker.Chunker.
  2. Implement chunk(ctx, parsed) -> list[Chunk] and health().
  3. Honour the SPI invariants (tenant_id from ctx, monotonic position, parent_id references earlier chunks, content/content_ref set).
  4. Map unexpected errors to rag_core.errors.ChunkError.

The contract tests in tests/contract/test_chunker.py cover the invariants; new chunkers should add a unit-test file under tests/chunker/ mirroring test_heading_aware.py.