Add sentence- and heading-aware chunking with a pluggable tokenizer (#26)

Replace the fixed-window-only chunker with selectable strategies behind a
ChunkingStrategy protocol:

- fixed_tokens (default, unchanged behavior), sentence_aware (packs whole
  sentences, never splits mid-sentence, with sentence-level overlap), and
  heading_aware (splits on parser-emitted heading boundaries then sentence-packs
  each section, tagging chunks with their heading).
- Pluggable tokenizer: real BPE counts via the optional ingestion-tokenizer
  extra (tiktoken), whitespace word-count fallback otherwise.
- Orchestrator passes parser sections and a per-run tokenizer into chunking.
- Chunk IDs remain deterministic across strategies. Docs updated.

Ingestion coverage 87% -> 89%.

Author: DiogoRibeiro7

README.md

@@ -64,6 +64,7 @@ pip install -e ".[dev]"
 | `observability-cli` | `jsonschema`, `numpy`, `pandas` | `--validate`, `validate-config`, `--detect-anomalies`, `--export-csv` |
 | `parquet` | `pyarrow` | `--export-parquet` (combine with `observability-cli`) |
 | `ingestion-pdf` | `pypdf` | Ingesting `.pdf` source documents |
+| `ingestion-tokenizer` | `tiktoken` | Real token counts for chunking (whitespace fallback otherwise) |
 
 Example: `pip install -e ".[dev,observability-cli,parquet]"`.
 

docs/mvp-ingestion.md

@@ -32,7 +32,7 @@ ingestion:
   input_path: ./examples         # directory scanned recursively
   max_documents: 100             # cap on processed files
 chunking:
-  strategy: fixed_tokens
+  strategy: fixed_tokens      # fixed_tokens | sentence_aware | heading_aware
   target_tokens: 400
   overlap_tokens: 40
 output:
@@ -45,11 +45,32 @@ output:
 
 Relative paths are resolved against the config file's directory.
 
+## Chunking strategies
+
+Selectable via `chunking.strategy`:
+
+- `fixed_tokens` (default) — sliding window over whitespace tokens. Fast and
+  format-agnostic; may split mid-sentence.
+- `sentence_aware` — packs whole sentences up to `target_tokens`, carrying
+  `overlap_tokens` worth of trailing sentences into the next chunk. Never splits
+  mid-sentence. Best for prose.
+- `heading_aware` — splits on heading boundaries emitted by the Markdown/HTML
+  parsers, then sentence-packs within each section; each chunk records its
+  originating heading in `section`. Falls back to `sentence_aware` when the
+  document has no headings.
+
+Token counts use a real tokenizer when the optional `ingestion-tokenizer` extra
+(`tiktoken`) is installed, and fall back to a whitespace word count otherwise.
+Chunk IDs are deterministic under every strategy (keyed by document, index, and
+offsets), so artifacts stay diffable across runs.
+
 ## Supported inputs
 
 - `.txt`
-- `.md`
-- `.json`
+- `.md` / `.markdown` (frontmatter stripped, headings preserved)
+- `.html` / `.htm` (reduced to clean text)
+- `.json` (deterministic re-serialization, or configurable text-field selection)
+- `.pdf` (requires the optional `ingestion-pdf` extra)
 
 Files are discovered recursively under `ingestion.input_path` and processed in stable sorted order.
 

pyproject.toml

@@ -35,6 +35,9 @@ parquet = [
 ingestion-pdf = [
   "pypdf>=4.0.0",
 ]
+ingestion-tokenizer = [
+  "tiktoken>=0.7.0",
+]
 dev = [
   "mypy>=1.11.0",
   "pytest>=8.3.0",

src/llm_knowledge_ingestion/chunking/strategies.py

@@ -2,9 +2,18 @@
 
 import re
 from dataclasses import dataclass
+from typing import Protocol
 
+from llm_knowledge_ingestion.chunking.tokenizer import Tokenizer, get_tokenizer
 from llm_knowledge_ingestion.contracts.models import Chunk
 from llm_knowledge_ingestion.dedup.hashing import sha256_text
+from llm_knowledge_ingestion.parsers.base import Section
+
+SUPPORTED_STRATEGIES = {"fixed_tokens", "sentence_aware", "heading_aware"}
+
+TOKEN_RE = re.compile(r"\S+")
+# Split after sentence-ending punctuation followed by whitespace.
+SENTENCE_BOUNDARY_RE = re.compile(r"(?<=[.!?])\s+")
 
 
 @dataclass(frozen=True, slots=True)
@@ -14,8 +23,10 @@ class ChunkingConfig:
     overlap_tokens: int = 40
 
     def __post_init__(self) -> None:
-        if self.strategy != "fixed_tokens":
-            raise ValueError("Only fixed_tokens strategy is supported in MVP")
+        if self.strategy not in SUPPORTED_STRATEGIES:
+            raise ValueError(
+                f"strategy must be one of {sorted(SUPPORTED_STRATEGIES)}, got {self.strategy!r}"
+            )
         if self.target_tokens <= 0:
             raise ValueError("target_tokens must be > 0")
         if self.overlap_tokens < 0:
@@ -24,44 +35,218 @@ def __post_init__(self) -> None:
             raise ValueError("overlap_tokens must be < target_tokens")
 
 
-TOKEN_RE = re.compile(r"\S+")
+class ChunkingStrategy(Protocol):
+    def split(
+        self,
+        text: str,
+        document_id: str,
+        config: ChunkingConfig,
+        sections: list[Section] | None,
+        tokenizer: Tokenizer,
+    ) -> list[Chunk]:
+        """Split text into deterministic chunks."""
+        ...
 
 
 def _chunk_id(document_id: str, chunk_index: int, text: str, start: int, end: int) -> str:
     digest = sha256_text(f"{document_id}|{chunk_index}|{start}|{end}|{text}")
     return f"chk_{digest[:28]}"
 
 
-def chunk_document(content: str, document_id: str, config: ChunkingConfig) -> list[Chunk]:
-    """Split content into deterministic token-window chunks."""
-    tokens = list(TOKEN_RE.finditer(content))
-    if not tokens:
-        return []
+def _make_chunk(
+    document_id: str,
+    chunk_index: int,
+    content: str,
+    start: int,
+    end: int,
+    token_count: int,
+    section: str | None,
+) -> Chunk:
+    text = content[start:end]
+    return Chunk(
+        chunk_id=_chunk_id(document_id, chunk_index, text, start, end),
+        document_id=document_id,
+        chunk_index=chunk_index,
+        text=text,
+        token_count_estimate=token_count,
+        start_offset=start,
+        end_offset=end,
+        section=section,
+        metadata={},
+    )
 
-    step = config.target_tokens - config.overlap_tokens
-    chunks: list[Chunk] = []
-    chunk_index = 0
 
-    for token_start in range(0, len(tokens), step):
-        token_end = min(len(tokens), token_start + config.target_tokens)
-        start_offset = tokens[token_start].start()
-        end_offset = tokens[token_end - 1].end()
-        text = content[start_offset:end_offset]
-        chunks.append(
-            Chunk(
-                chunk_id=_chunk_id(document_id, chunk_index, text, start_offset, end_offset),
-                document_id=document_id,
-                chunk_index=chunk_index,
-                text=text,
-                token_count_estimate=token_end - token_start,
-                start_offset=start_offset,
-                end_offset=end_offset,
-                section=None,
-                metadata={},
+class FixedTokenStrategy:
+    """Whitespace-token sliding window — the backwards-compatible default."""
+
+    def split(
+        self,
+        text: str,
+        document_id: str,
+        config: ChunkingConfig,
+        sections: list[Section] | None,
+        tokenizer: Tokenizer,
+    ) -> list[Chunk]:
+        tokens = list(TOKEN_RE.finditer(text))
+        if not tokens:
+            return []
+        step = config.target_tokens - config.overlap_tokens
+        chunks: list[Chunk] = []
+        chunk_index = 0
+        for token_start in range(0, len(tokens), step):
+            token_end = min(len(tokens), token_start + config.target_tokens)
+            start = tokens[token_start].start()
+            end = tokens[token_end - 1].end()
+            chunks.append(
+                _make_chunk(
+                    document_id,
+                    chunk_index,
+                    text,
+                    start,
+                    end,
+                    tokenizer.count(text[start:end]),
+                    None,
+                )
             )
-        )
-        chunk_index += 1
-        if token_end == len(tokens):
-            break
+            chunk_index += 1
+            if token_end == len(tokens):
+                break
+        return chunks
+
+
+def _sentence_spans(text: str, start: int, end: int) -> list[tuple[int, int]]:
+    """Return (start, end) char spans of sentences within text[start:end]."""
+    segment = text[start:end]
+    spans: list[tuple[int, int]] = []
+    cursor = 0
+    for piece in SENTENCE_BOUNDARY_RE.split(segment):
+        if not piece:
+            continue
+        idx = segment.find(piece, cursor)
+        if idx < 0:
+            continue
+        spans.append((start + idx, start + idx + len(piece)))
+        cursor = idx + len(piece)
+    return spans
 
+
+def _pack_sentences(
+    text: str,
+    document_id: str,
+    config: ChunkingConfig,
+    tokenizer: Tokenizer,
+    region_start: int,
+    region_end: int,
+    section: str | None,
+    start_index: int,
+) -> list[Chunk]:
+    """Greedily pack sentences into chunks under the token budget."""
+    spans = _sentence_spans(text, region_start, region_end)
+    if not spans:
+        return []
+    chunks: list[Chunk] = []
+    chunk_index = start_index
+    window: list[tuple[int, int]] = []
+    window_tokens = 0
+    for span in spans:
+        span_tokens = tokenizer.count(text[span[0] : span[1]])
+        # Flush when the window is non-empty and adding the next sentence overflows.
+        if window and window_tokens + span_tokens > config.target_tokens:
+            start, end = window[0][0], window[-1][1]
+            chunks.append(
+                _make_chunk(document_id, chunk_index, text, start, end, window_tokens, section)
+            )
+            chunk_index += 1
+            # Carry trailing sentences as overlap.
+            overlap: list[tuple[int, int]] = []
+            overlap_tokens = 0
+            for prev in reversed(window):
+                prev_tokens = tokenizer.count(text[prev[0] : prev[1]])
+                if overlap_tokens + prev_tokens > config.overlap_tokens:
+                    break
+                overlap.insert(0, prev)
+                overlap_tokens += prev_tokens
+            window = overlap
+            window_tokens = overlap_tokens
+        window.append(span)
+        window_tokens += span_tokens
+    if window:
+        start, end = window[0][0], window[-1][1]
+        chunks.append(
+            _make_chunk(document_id, chunk_index, text, start, end, window_tokens, section)
+        )
     return chunks
+
+
+class SentenceAwareStrategy:
+    """Pack whole sentences into chunks, preferring sentence boundaries."""
+
+    def split(
+        self,
+        text: str,
+        document_id: str,
+        config: ChunkingConfig,
+        sections: list[Section] | None,
+        tokenizer: Tokenizer,
+    ) -> list[Chunk]:
+        return _pack_sentences(text, document_id, config, tokenizer, 0, len(text), None, 0)
+
+
+class HeadingAwareStrategy:
+    """Split on parser-emitted heading boundaries, then sentence-pack each region."""
+
+    def split(
+        self,
+        text: str,
+        document_id: str,
+        config: ChunkingConfig,
+        sections: list[Section] | None,
+        tokenizer: Tokenizer,
+    ) -> list[Chunk]:
+        if not sections:
+            # No structure available: degrade to sentence packing.
+            return SentenceAwareStrategy().split(text, document_id, config, sections, tokenizer)
+
+        # Build (region_start, region_end, heading) tuples covering the whole text.
+        ordered = sorted(sections, key=lambda s: s.start_offset)
+        boundaries: list[tuple[int, int, str | None]] = []
+        if ordered[0].start_offset > 0:
+            boundaries.append((0, ordered[0].start_offset, None))
+        for index, section in enumerate(ordered):
+            region_end = ordered[index + 1].start_offset if index + 1 < len(ordered) else len(text)
+            boundaries.append((section.start_offset, region_end, section.heading))
+
+        chunks: list[Chunk] = []
+        for region_start, region_end, heading in boundaries:
+            chunks.extend(
+                _pack_sentences(
+                    text,
+                    document_id,
+                    config,
+                    tokenizer,
+                    region_start,
+                    region_end,
+                    heading,
+                    len(chunks),
+                )
+            )
+        return chunks
+
+
+_STRATEGIES: dict[str, ChunkingStrategy] = {
+    "fixed_tokens": FixedTokenStrategy(),
+    "sentence_aware": SentenceAwareStrategy(),
+    "heading_aware": HeadingAwareStrategy(),
+}
+
+
+def chunk_document(
+    content: str,
+    document_id: str,
+    config: ChunkingConfig,
+    sections: list[Section] | None = None,
+    tokenizer: Tokenizer | None = None,
+) -> list[Chunk]:
+    """Split content into deterministic chunks using the configured strategy."""
+    strategy = _STRATEGIES[config.strategy]
+    return strategy.split(content, document_id, config, sections, tokenizer or get_tokenizer())

src/llm_knowledge_ingestion/chunking/tokenizer.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+import re
+from typing import Protocol
+
+TOKEN_RE = re.compile(r"\S+")
+
+
+class Tokenizer(Protocol):
+    name: str
+
+    def count(self, text: str) -> int:
+        """Return the number of tokens in ``text``."""
+        ...
+
+
+class WhitespaceTokenizer:
+    """Dependency-free fallback: counts whitespace-separated words."""
+
+    name = "whitespace"
+
+    def count(self, text: str) -> int:
+        return len(TOKEN_RE.findall(text))
+
+
+class TiktokenTokenizer:
+    """Real BPE token counts via the optional ``tiktoken`` dependency."""
+
+    def __init__(self, encoding_name: str = "cl100k_base") -> None:
+        import tiktoken
+
+        self._encoding = tiktoken.get_encoding(encoding_name)
+        self.name = f"tiktoken:{encoding_name}"
+
+    def count(self, text: str) -> int:
+        return len(self._encoding.encode(text))
+
+
+def get_tokenizer() -> Tokenizer:
+    """Return a real tokenizer when ``tiktoken`` is installed, else the fallback."""
+    try:
+        return TiktokenTokenizer()
+    except ImportError:
+        return WhitespaceTokenizer()