feat: add TokenChunker for token-level chunking via tiktoken (closes #44)

CHANGELOG.md

@@ -8,6 +8,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
+
+- `TokenChunker` for token-level chunking using tiktoken with a new `[tokens]` optional extra (closes #44)
+
 - Structured logging for ingest, query, and retriever stages, including timing and empty-ingest warnings (closes #41).
 - `RecursiveChunker.from_config()` for creating a chunker from pipeline configuration (closes #25).
 

pyproject.toml

@@ -36,6 +36,7 @@ anthropic = ["anthropic>=0.25"]
 huggingface = ["sentence-transformers>=2.7"]
 chromadb = ["chromadb>=0.5"]
 faiss = ["faiss-cpu>=1.8"]
+tokens = ["tiktoken>=0.7"]
 all = [
     "ragframework[pdf]",
     "ragframework[docx]",
@@ -44,6 +45,7 @@ all = [
     "ragframework[huggingface]",
     "ragframework[chromadb]",
     "ragframework[faiss]",
+    "ragframework[tokens]",
 ]
 dev = [
     "pytest>=8.0",

ragframework/document/__init__.py

@@ -1,6 +1,11 @@
 """Document loading and chunking utilities."""
 
-from ragframework.document.chunkers import FixedSizeChunker, RecursiveChunker, SentenceChunker
+from ragframework.document.chunkers import (
+    FixedSizeChunker,
+    RecursiveChunker,
+    SentenceChunker,
+    TokenChunker,
+)
 
 from .html import HTMLLoader
 from .loaders import MarkdownLoader, PDFLoader, TextFileLoader
@@ -16,4 +21,5 @@
     "FixedSizeChunker",
     "RecursiveChunker",
     "SentenceChunker",
+    "TokenChunker",
 ]

ragframework/document/chunkers.py

@@ -359,3 +359,100 @@ def _merge_pieces(self, pieces: list[str], next_separators: list[str]) -> list[s
             final_chunks.append("".join(current_chunk_pieces))
 
         return final_chunks
+
+
+class TokenChunker(TextChunker):
+    """Split text into chunks measured in tokens using tiktoken.
+    Args:
+        chunk_tokens: Maximum number of tokens per chunk.
+        overlap_tokens: Number of overlapping tokens between adjacent chunks.
+        encoding_name: Name of the tiktoken encoding to use.
+    Raises:
+        ImportError: If the ``tokens`` optional dependency is not installed.
+    """
+
+    def __init__(
+        self,
+        chunk_tokens: int = 256,
+        overlap_tokens: int = 32,
+        encoding_name: str = "cl100k_base",
+    ) -> None:
+        # --- Guarded import (same pattern as HuggingFaceEmbedder) ---
+        try:
+            import tiktoken
+        except ImportError as exc:
+            raise ImportError(
+                "Token chunking requires 'ragframework[tokens]'. "
+                "Install it with: pip install ragframework[tokens]"
+            ) from exc
+        # --- Validation (same rules as RecursiveChunker) ---
+        if chunk_tokens <= 0:
+            raise ValueError("chunk_tokens must be positive")
+        if overlap_tokens < 0:
+            raise ValueError("overlap_tokens must be non-negative")
+
+        if overlap_tokens >= chunk_tokens:
+            raise ValueError("overlap_tokens must be less than chunk_tokens")
+
+        self.chunk_tokens = chunk_tokens
+        self.overlap_tokens = overlap_tokens
+        self.encoding_name = encoding_name
+        self._encoding = tiktoken.get_encoding(encoding_name)
+
+    def chunk(self, document: Document) -> list[Chunk]:
+        if not document.content:
+            return []
+        # Treat special tokens like <|endoftext|> as ordinary text
+        all_tokens = self._encoding.encode(document.content, disallowed_special=())
+        if not all_tokens:
+            return []
+        chunks: list[Chunk] = []
+        start_idx = 0
+        chunk_num = 0
+        total_tokens = len(all_tokens)
+        while start_idx < total_tokens:
+            end_idx = min(start_idx + self.chunk_tokens, total_tokens)
+            # Shrink window backwards if it cuts across a multi-byte UTF-8 character
+            chunk_text: str | None = None
+            while end_idx > start_idx:
+                raw_bytes = self._encoding.decode_bytes(all_tokens[start_idx:end_idx])
+                try:
+                    chunk_text = raw_bytes.decode("utf-8")
+                    break
+                except UnicodeDecodeError:
+                    end_idx -= 1
+            # If even 1 token cannot complete a character, it exceeds chunk_tokens
+            if end_idx == start_idx or chunk_text is None:
+                raise ValueError(
+                    f"Character at token index {start_idx} cannot fit within "
+                    f"chunk_tokens={self.chunk_tokens}. Increase chunk_tokens."
+                )
+            token_count = end_idx - start_idx
+            chunks.append(
+                Chunk(
+                    id=f"{document.id}:{chunk_num}",
+                    content=chunk_text,
+                    metadata={
+                        **document.metadata,
+                        "chunk_index": chunk_num,
+                        "token_count": token_count,
+                    },
+                )
+            )
+            chunk_num += 1
+            if end_idx >= total_tokens:
+                break
+            # Calculate next start_idx respecting overlap and character boundaries
+            if self.overlap_tokens == 0:
+                start_idx = end_idx
+            else:
+                target_start = max(start_idx + 1, end_idx - self.overlap_tokens)
+                while target_start < end_idx:
+                    overlap_bytes = self._encoding.decode_bytes(all_tokens[target_start:end_idx])
+                    try:
+                        overlap_bytes.decode("utf-8")
+                        break
+                    except UnicodeDecodeError:
+                        target_start += 1
+                start_idx = target_start
+        return chunks

tests/test_document/test_chunkers.py

@@ -1,5 +1,9 @@
 """Tests for built-in text chunkers."""
 
+import builtins
+import sys
+import types
+
 import pytest
 
 from ragframework.base import Document
@@ -8,6 +12,7 @@
     RecursiveChunker,
     SentenceChunker,
 )
+from typing import Any
 
 
 def test_recursive_chunker_from_config():
@@ -293,3 +298,156 @@ def test_redundant_trailing_chunks(self):
         )
         chunks = chunker.chunk(doc)
         assert [len(c.content) for c in chunks] == [6, 6, 6]
+
+
+# --------------- Fake tiktoken for CI ---------------
+
+
+class FakeEncoding:
+    """Emulates tiktoken with byte-level tokens, special token checks, and decode_bytes."""
+
+    def encode(
+        self,
+        text: str,
+        *,
+        allowed_special: Any = (),
+        disallowed_special: Any = "all",
+    ) -> list[int]:
+        if disallowed_special and "<|endoftext|>" in text:
+            raise ValueError(
+                "Encountered text corresponding to disallowed special token '<|endoftext|>'."
+            )
+        return list(text.encode("utf-8"))
+
+    def decode(self, tokens: list[int], errors: str = "replace") -> str:
+        return bytes(tokens).decode("utf-8", errors=errors)
+
+    def decode_bytes(self, tokens: list[int]) -> bytes:
+        return bytes(tokens)
+
+
+@pytest.fixture
+def fake_tiktoken(monkeypatch):
+    """Inject a fake tiktoken module so tests don't need a real download."""
+    module = types.ModuleType("tiktoken")
+    module.get_encoding = lambda name: FakeEncoding()
+    monkeypatch.setitem(sys.modules, "tiktoken", module)
+
+
+class TestTokenChunker:
+    def test_basic_chunking(self, fake_tiktoken):
+        from ragframework.document.chunkers import TokenChunker
+
+        doc = Document(id="d1", content="A" * 100, metadata={})
+        chunker = TokenChunker(chunk_tokens=30, overlap_tokens=5)
+        chunks = chunker.chunk(doc)
+
+        assert len(chunks) > 1
+        # No chunk should exceed the token limit
+        for c in chunks:
+            assert c.metadata["token_count"] <= 30
+
+    def test_chunk_ids_follow_convention(self, fake_tiktoken):
+        from ragframework.document.chunkers import TokenChunker
+
+        doc = Document(id="doc1", content="Hello world this is a test", metadata={"src": "test"})
+        chunker = TokenChunker(chunk_tokens=10, overlap_tokens=2)
+        chunks = chunker.chunk(doc)
+
+        for i, c in enumerate(chunks):
+            assert c.id == f"doc1:{i}"
+            assert c.metadata["chunk_index"] == i
+            assert c.metadata["src"] == "test"
+            assert "token_count" in c.metadata
+
+    def test_empty_doc_returns_empty(self, fake_tiktoken):
+        from ragframework.document.chunkers import TokenChunker
+
+        doc = Document(id="x", content="", metadata={})
+        chunker = TokenChunker(chunk_tokens=10, overlap_tokens=0)
+        assert chunker.chunk(doc) == []
+
+    def test_short_doc_single_chunk(self, fake_tiktoken):
+        from ragframework.document.chunkers import TokenChunker
+
+        doc = Document(id="x", content="Hi", metadata={})
+        chunker = TokenChunker(chunk_tokens=256, overlap_tokens=0)
+        chunks = chunker.chunk(doc)
+        assert len(chunks) == 1
+        assert chunks[0].content == "Hi"
+        assert chunks[0].metadata["token_count"] == 2  # 'H' and 'i' in our fake
+
+    def test_validation_errors(self, fake_tiktoken):
+        from ragframework.document.chunkers import TokenChunker
+
+        with pytest.raises(ValueError, match="chunk_tokens must be positive"):
+            TokenChunker(chunk_tokens=0)
+        with pytest.raises(ValueError, match="overlap_tokens must be non-negative"):
+            TokenChunker(chunk_tokens=10, overlap_tokens=-1)
+        with pytest.raises(ValueError, match="overlap_tokens must be less than chunk_tokens"):
+            TokenChunker(chunk_tokens=10, overlap_tokens=10)
+
+    def test_no_chunk_exceeds_token_limit(self, fake_tiktoken):
+        from ragframework.document.chunkers import TokenChunker
+
+        doc = Document(id="d1", content="abcdefghijklmnopqrstuvwxyz" * 10, metadata={})
+        chunker = TokenChunker(chunk_tokens=50, overlap_tokens=10)
+        chunks = chunker.chunk(doc)
+
+        for c in chunks:
+            assert c.metadata["token_count"] <= 50
+
+    def test_missing_tiktoken_gives_helpful_message(self, monkeypatch):
+        real_import = builtins.__import__
+
+        def fake_import(name, globals=None, locals=None, fromlist=(), level=0):
+            if name == "tiktoken":
+                raise ImportError("No module named 'tiktoken'")
+            return real_import(name, globals, locals, fromlist, level)
+
+        monkeypatch.delitem(sys.modules, "tiktoken", raising=False)
+        monkeypatch.setattr(builtins, "__import__", fake_import)
+
+        from ragframework.document.chunkers import TokenChunker
+
+        with pytest.raises(ImportError, match=r"ragframework\[tokens\]"):
+            TokenChunker()
+
+    def test_special_tokens_treated_as_ordinary_text(self, fake_tiktoken):
+        from ragframework.document.chunkers import TokenChunker
+
+        doc = Document(id="x", content="Hello <|endoftext|> world", metadata={})
+        chunker = TokenChunker(chunk_tokens=50, overlap_tokens=0)
+        chunks = chunker.chunk(doc)
+        assert len(chunks) == 1
+        assert chunks[0].content == "Hello <|endoftext|> world"
+
+    def test_emoji_preserves_complete_characters(self, fake_tiktoken):
+        from ragframework.document.chunkers import TokenChunker
+
+        doc = Document(id="x", content="A 🙂 B 🚀 C", metadata={})
+        chunker = TokenChunker(chunk_tokens=6, overlap_tokens=0)
+        chunks = chunker.chunk(doc)
+        assert len(chunks) > 1
+        for c in chunks:
+            assert "\ufffd" not in c.content
+        assert "".join(c.content for c in chunks) == "A 🙂 B 🚀 C"
+
+    def test_cjk_preserves_complete_characters(self, fake_tiktoken):
+        from ragframework.document.chunkers import TokenChunker
+
+        doc = Document(id="x", content="你好世界，这是一个测试", metadata={})
+        chunker = TokenChunker(chunk_tokens=9, overlap_tokens=0)
+        chunks = chunker.chunk(doc)
+        assert len(chunks) > 1
+        for c in chunks:
+            assert "\ufffd" not in c.content
+        assert "".join(c.content for c in chunks) == "你好世界，这是一个测试"
+
+    def test_character_cannot_fit_raises_value_error(self, fake_tiktoken):
+        from ragframework.document.chunkers import TokenChunker
+
+        doc = Document(id="x", content="🙂", metadata={})
+        chunker = TokenChunker(chunk_tokens=1, overlap_tokens=0)
+        with pytest.raises(ValueError, match="cannot fit within chunk_tokens"):
+            chunker.chunk(doc)