diff --git a/haiku_rag_slim/haiku/rag/chunkers/__init__.py b/haiku_rag_slim/haiku/rag/chunkers/__init__.py new file mode 100644 index 00000000..38d43f22 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/chunkers/__init__.py @@ -0,0 +1,31 @@ +"""Document chunker abstraction for haiku.rag.""" + +from haiku.rag.chunkers.base import DocumentChunker +from haiku.rag.config import AppConfig, Config + +__all__ = ["DocumentChunker", "get_chunker"] + + +def get_chunker(config: AppConfig = Config) -> DocumentChunker: + """Get a document chunker instance based on configuration. + + Args: + config: Configuration to use. Defaults to global Config. + + Returns: + DocumentChunker instance configured according to the config. + + Raises: + ValueError: If the chunker provider is not recognized. + """ + if config.processing.chunker == "docling-local": + from haiku.rag.chunkers.docling_local import DoclingLocalChunker + + return DoclingLocalChunker(config) + + if config.processing.chunker == "docling-serve": + from haiku.rag.chunkers.docling_serve import DoclingServeChunker + + return DoclingServeChunker(config) + + raise ValueError(f"Unsupported chunker: {config.processing.chunker}") diff --git a/haiku_rag_slim/haiku/rag/chunkers/base.py b/haiku_rag_slim/haiku/rag/chunkers/base.py new file mode 100644 index 00000000..01df87f6 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/chunkers/base.py @@ -0,0 +1,28 @@ +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from docling_core.types.doc.document import DoclingDocument + + +class DocumentChunker(ABC): + """Abstract base class for document chunkers. + + Document chunkers split DoclingDocuments into smaller text chunks suitable + for embedding and retrieval, respecting document structure and semantic boundaries. + """ + + @abstractmethod + async def chunk(self, document: "DoclingDocument") -> list[str]: + """Split a document into chunks. + + Args: + document: The DoclingDocument to chunk. + + Returns: + List of text chunks with semantic boundaries preserved. + + Raises: + ValueError: If chunking fails. + """ + pass diff --git a/haiku_rag_slim/haiku/rag/chunker.py b/haiku_rag_slim/haiku/rag/chunkers/docling_local.py similarity index 58% rename from haiku_rag_slim/haiku/rag/chunker.py rename to haiku_rag_slim/haiku/rag/chunkers/docling_local.py index 7a45c7bf..40dd020c 100644 --- a/haiku_rag_slim/haiku/rag/chunker.py +++ b/haiku_rag_slim/haiku/rag/chunkers/docling_local.py @@ -1,38 +1,38 @@ from typing import TYPE_CHECKING -from haiku.rag.config import Config +from haiku.rag.chunkers.base import DocumentChunker +from haiku.rag.config import AppConfig, Config if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument -class Chunker: - """A class that chunks text into smaller pieces for embedding and retrieval. +class DoclingLocalChunker(DocumentChunker): + """Local document chunker using docling's HybridChunker. Uses docling's structure-aware chunking to create semantically meaningful chunks - that respect document boundaries. + that respect document boundaries. Chunking is performed locally using the + HuggingFace tokenizer specified in configuration. Args: - chunk_size: The maximum size of a chunk in tokens. - tokenizer_name: HuggingFace model name for tokenization. + config: Application configuration. """ - def __init__( - self, - chunk_size: int = Config.processing.chunk_size, - tokenizer_name: str = Config.processing.chunking_tokenizer, - ): + def __init__(self, config: AppConfig = Config): from docling_core.transforms.chunker.hybrid_chunker import HybridChunker from docling_core.transforms.chunker.tokenizer.huggingface import ( HuggingFaceTokenizer, ) from transformers import AutoTokenizer - self.chunk_size = chunk_size - self.tokenizer_name = tokenizer_name + self.config = config + self.chunk_size = config.processing.chunk_size + self.tokenizer_name = config.processing.chunking_tokenizer - hf_tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) - tokenizer = HuggingFaceTokenizer(tokenizer=hf_tokenizer, max_tokens=chunk_size) + hf_tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name) + tokenizer = HuggingFaceTokenizer( + tokenizer=hf_tokenizer, max_tokens=self.chunk_size + ) self.chunker = HybridChunker(tokenizer=tokenizer) @@ -51,6 +51,3 @@ class Chunker: # Chunk using docling's hybrid chunker chunks = list(self.chunker.chunk(document)) return [self.chunker.contextualize(chunk) for chunk in chunks] - - -chunker = Chunker() diff --git a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py new file mode 100644 index 00000000..9547cf73 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py @@ -0,0 +1,31 @@ +from typing import TYPE_CHECKING + +from haiku.rag.chunkers.base import DocumentChunker +from haiku.rag.config import AppConfig + +if TYPE_CHECKING: + from docling_core.types.doc.document import DoclingDocument + + +class DoclingServeChunker(DocumentChunker): + """Remote document chunker using docling-serve API. + + Placeholder - will be implemented in a future commit. + """ + + def __init__(self, config: AppConfig): + raise NotImplementedError("DoclingServeChunker not yet implemented") + + async def chunk(self, document: "DoclingDocument") -> list[str]: + """Split the document into chunks via docling-serve. + + Args: + document: The DoclingDocument to be split into chunks. + + Returns: + A list of text chunks with semantic boundaries. + + Raises: + NotImplementedError: This chunker is not yet implemented. + """ + raise NotImplementedError("DoclingServeChunker not yet implemented") diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 84544454..8fb114d4 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -55,6 +55,7 @@ class ProcessingConfig(BaseModel): context_chunk_radius: int = 0 markdown_preprocessor: str = "" converter: str = "docling-local" + chunker: str = "docling-local" chunking_tokenizer: str = "Qwen/Qwen3-Embedding-0.6B" diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 370662ca..5818ef14 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -150,9 +150,11 @@ class ChunkRepository: ) -> list[Chunk]: """Create chunks and embeddings for a document from DoclingDocument.""" # Lazy imports to avoid loading docling during module import - from haiku.rag.chunker import chunker + from haiku.rag.chunkers import get_chunker from haiku.rag.converters import get_converter + chunker = get_chunker(self.store._config) + # Optionally preprocess markdown before chunking processed_document = document preprocessor_path = self.store._config.processing.markdown_preprocessor diff --git a/tests/test_chunker.py b/tests/test_chunker.py index 40c8085e..c65f617d 100644 --- a/tests/test_chunker.py +++ b/tests/test_chunker.py @@ -2,14 +2,16 @@ import pytest from datasets import Dataset from transformers import AutoTokenizer -from haiku.rag.chunker import Chunker -from haiku.rag.config import Config +from haiku.rag.chunkers import get_chunker +from haiku.rag.chunkers.docling_local import DoclingLocalChunker +from haiku.rag.config import AppConfig, Config from haiku.rag.converters import get_converter @pytest.mark.asyncio -async def test_chunker(qa_corpus: Dataset): - chunker = Chunker() +async def test_local_chunker(qa_corpus: Dataset): + """Test DoclingLocalChunker with real document.""" + chunker = DoclingLocalChunker() doc_text = qa_corpus[0]["document_extracted"] # Convert text to DoclingDocument @@ -43,3 +45,31 @@ async def test_chunker(qa_corpus: Dataset): # Due to structure-aware chunking, we might have some variation in token count # but it should be reasonable assert abs(total_tokens - original_tokens) <= original_tokens * 0.1 + + +@pytest.mark.asyncio +async def test_local_chunker_custom_config(): + """Test DoclingLocalChunker with custom configuration.""" + config = AppConfig() + config.processing.chunk_size = 128 + config.processing.chunking_tokenizer = "Qwen/Qwen3-Embedding-0.6B" + + chunker = DoclingLocalChunker(config) + assert chunker.chunk_size == 128 + assert chunker.tokenizer_name == "Qwen/Qwen3-Embedding-0.6B" + + +def test_get_chunker_docling_local(): + """Test factory returns DoclingLocalChunker for docling-local.""" + config = AppConfig() + config.processing.chunker = "docling-local" + chunker = get_chunker(config) + assert isinstance(chunker, DoclingLocalChunker) + + +def test_get_chunker_invalid(): + """Test factory raises error for invalid chunker.""" + config = AppConfig() + config.processing.chunker = "invalid-chunker" + with pytest.raises(ValueError, match="Unsupported chunker"): + get_chunker(config)