haiku.rag-slim 0.16.0__py3-none-any.whl → 0.24.0__py3-none-any.whl

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}")

haiku/rag/chunkers/base.py

@@ -0,0 +1,31 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from docling_core.types.doc.document import DoclingDocument
+
+    from haiku.rag.store.models.chunk import Chunk
+
+
+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["Chunk"]:
+        """Split a document into chunks with metadata.
+
+        Args:
+            document: The DoclingDocument to chunk.
+
+        Returns:
+            List of Chunk with content and structured metadata in the metadata dict
+            (doc_item_refs, headings, labels, page_numbers).
+
+        Raises:
+            ValueError: If chunking fails.
+        """
+        pass

haiku/rag/chunkers/docling_local.py

@@ -0,0 +1,164 @@
+from typing import TYPE_CHECKING, cast
+
+from haiku.rag.chunkers.base import DocumentChunker
+from haiku.rag.config import AppConfig, Config
+from haiku.rag.store.models.chunk import Chunk, ChunkMetadata
+
+if TYPE_CHECKING:
+    from docling_core.transforms.chunker.doc_chunk import DocMeta
+    from docling_core.types.doc.document import DoclingDocument
+
+
+def _create_markdown_serializer_provider(use_markdown_tables: bool = True):
+    """Create a markdown serializer provider with configurable table rendering.
+
+    This function creates a custom serializer provider that extends ChunkingSerializerProvider
+    from docling-core. It's implemented as a factory function to avoid importing
+    docling-core at module level.
+
+    Args:
+        use_markdown_tables: If True, use MarkdownTableSerializer for rendering tables as
+            markdown. If False, use default TripletTableSerializer for narrative format.
+    """
+    from docling_core.transforms.chunker.hierarchical_chunker import (
+        ChunkingDocSerializer,
+        ChunkingSerializerProvider,
+    )
+    from docling_core.transforms.serializer.markdown import MarkdownTableSerializer
+
+    class MDTableSerializerProvider(ChunkingSerializerProvider):
+        """Serializer provider for markdown table output."""
+
+        def __init__(self, use_markdown_tables: bool = True):
+            self.use_markdown_tables = use_markdown_tables
+
+        def get_serializer(self, doc):
+            if self.use_markdown_tables:
+                return ChunkingDocSerializer(
+                    doc=doc,
+                    table_serializer=MarkdownTableSerializer(),
+                )
+            else:
+                # Use default ChunkingDocSerializer (TripletTableSerializer)
+                return ChunkingDocSerializer(doc=doc)
+
+    return MDTableSerializerProvider(use_markdown_tables=use_markdown_tables)
+
+
+class DoclingLocalChunker(DocumentChunker):
+    """Local document chunker using docling's chunkers.
+
+    Supports both hybrid (structure-aware) and hierarchical chunking strategies.
+    Chunking is performed locally using the HuggingFace tokenizer specified in
+    configuration.
+
+    Args:
+        config: Application configuration.
+    """
+
+    def __init__(self, config: AppConfig = Config):
+        from docling_core.transforms.chunker.hierarchical_chunker import (
+            HierarchicalChunker,
+        )
+        from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
+        from docling_core.transforms.chunker.tokenizer.huggingface import (
+            HuggingFaceTokenizer,
+        )
+        from transformers import AutoTokenizer
+
+        self.config = config
+        self.chunk_size = config.processing.chunk_size
+        self.chunker_type = config.processing.chunker_type
+        self.tokenizer_name = config.processing.chunking_tokenizer
+
+        if self.chunker_type == "hybrid":
+            hf_tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)
+            tokenizer = HuggingFaceTokenizer(
+                tokenizer=hf_tokenizer, max_tokens=self.chunk_size
+            )
+            serializer_provider = _create_markdown_serializer_provider(
+                use_markdown_tables=config.processing.chunking_use_markdown_tables
+            )
+            self.chunker = HybridChunker(
+                tokenizer=tokenizer,
+                merge_peers=config.processing.chunking_merge_peers,
+                serializer_provider=serializer_provider,
+            )
+        elif self.chunker_type == "hierarchical":
+            serializer_provider = _create_markdown_serializer_provider(
+                use_markdown_tables=config.processing.chunking_use_markdown_tables
+            )
+            self.chunker = HierarchicalChunker(serializer_provider=serializer_provider)
+        else:
+            raise ValueError(
+                f"Unsupported chunker_type: {self.chunker_type}. "
+                "Must be 'hybrid' or 'hierarchical'."
+            )
+
+    async def chunk(self, document: "DoclingDocument") -> list[Chunk]:
+        """Split the document into chunks with metadata.
+
+        Extracts structured metadata from each DocChunk including:
+        - doc_item_refs: JSON pointer references to DocItems (e.g., "#/texts/5")
+        - headings: Section heading hierarchy
+        - labels: Semantic labels for each doc_item (e.g., "paragraph", "table")
+        - page_numbers: Page numbers where content appears
+
+        Args:
+            document: The DoclingDocument to be split into chunks.
+
+        Returns:
+            List of Chunk containing content and structured metadata.
+        """
+        if document is None:
+            return []
+
+        raw_chunks = list(self.chunker.chunk(document))
+        result: list[Chunk] = []
+
+        for chunk in raw_chunks:
+            text = chunk.text
+
+            # Extract metadata from DocChunk.meta (cast to DocMeta for type safety)
+            doc_item_refs: list[str] = []
+            labels: list[str] = []
+            page_numbers: list[int] = []
+            headings: list[str] | None = None
+
+            meta = cast("DocMeta | None", chunk.meta)
+            if meta and meta.doc_items:
+                for doc_item in meta.doc_items:
+                    # Get JSON pointer reference
+                    if doc_item.self_ref:
+                        doc_item_refs.append(doc_item.self_ref)
+                    # Get label
+                    if doc_item.label:
+                        labels.append(doc_item.label)
+                    # Get page numbers from provenance
+                    if doc_item.prov:
+                        for prov in doc_item.prov:
+                            if (
+                                prov.page_no is not None
+                                and prov.page_no not in page_numbers
+                            ):
+                                page_numbers.append(prov.page_no)
+
+            # Get headings from chunk metadata
+            if meta and meta.headings:
+                headings = list(meta.headings)
+
+            chunk_metadata = ChunkMetadata(
+                doc_item_refs=doc_item_refs,
+                headings=headings,
+                labels=labels,
+                page_numbers=sorted(page_numbers),
+            )
+            result.append(
+                Chunk(
+                    content=text,
+                    metadata=chunk_metadata.model_dump(),
+                    order=len(result),
+                )
+            )
+
+        return result

haiku/rag/chunkers/docling_serve.py

@@ -0,0 +1,179 @@
+import re
+from io import BytesIO
+from typing import TYPE_CHECKING
+
+from haiku.rag.chunkers.base import DocumentChunker
+from haiku.rag.config import AppConfig, Config
+from haiku.rag.providers.docling_serve import DoclingServeClient
+from haiku.rag.store.models.chunk import Chunk, ChunkMetadata
+
+if TYPE_CHECKING:
+    from docling_core.types.doc.document import DoclingDocument
+
+# Pattern to parse refs like "#/texts/5" or "#/tables/0"
+REF_PATTERN = re.compile(r"^#/(\w+)/(\d+)$")
+
+
+def _resolve_label_from_document(ref: str, document: "DoclingDocument") -> str | None:
+    """Resolve the label for a doc_item ref by looking it up in the document.
+
+    The docling-serve API only returns ref strings in doc_items, not labels.
+    This function resolves actual labels from the DoclingDocument.
+    See: https://github.com/docling-project/docling-serve/issues/448
+
+    Args:
+        ref: JSON pointer reference like "#/texts/5" or "#/tables/0"
+        document: The DoclingDocument to look up the item in
+
+    Returns:
+        The label string if found, None otherwise
+    """
+    match = REF_PATTERN.match(ref)
+    if not match:
+        return None
+
+    collection_name = match.group(1)
+    index = int(match.group(2))
+
+    collection = getattr(document, collection_name, None)
+    if collection is None or index >= len(collection):
+        return None
+
+    item = collection[index]
+    return getattr(item, "label", None)
+
+
+class DoclingServeChunker(DocumentChunker):
+    """Remote document chunker using docling-serve API.
+
+    Sends DoclingDocument JSON to docling-serve for chunking. Supports both hybrid
+    and hierarchical chunking strategies via remote API.
+
+    Args:
+        config: Application configuration containing docling-serve settings.
+    """
+
+    def __init__(self, config: AppConfig = Config):
+        self.config = config
+        self.client = DoclingServeClient(
+            base_url=config.providers.docling_serve.base_url,
+            api_key=config.providers.docling_serve.api_key,
+        )
+        self.chunker_type = config.processing.chunker_type
+
+    def _build_chunking_data(self) -> dict[str, str]:
+        """Build form data for chunking request."""
+        return {
+            "chunking_max_tokens": str(self.config.processing.chunk_size),
+            "chunking_tokenizer": self.config.processing.chunking_tokenizer,
+            "chunking_merge_peers": str(
+                self.config.processing.chunking_merge_peers
+            ).lower(),
+            "chunking_use_markdown_tables": str(
+                self.config.processing.chunking_use_markdown_tables
+            ).lower(),
+        }
+
+    async def _call_chunk_api(self, document: "DoclingDocument") -> list[dict]:
+        """Call docling-serve chunking API and return raw chunk data.
+
+        Args:
+            document: The DoclingDocument to be split into chunks.
+
+        Returns:
+            List of chunk dictionaries from API response.
+
+        Raises:
+            ValueError: If chunking fails or service is unavailable.
+        """
+        # Determine endpoint based on chunker_type
+        if self.chunker_type == "hierarchical":
+            endpoint = "/v1/chunk/hierarchical/file/async"
+        else:
+            endpoint = "/v1/chunk/hybrid/file/async"
+
+        # Export document to JSON
+        doc_json = document.model_dump_json()
+        doc_bytes = doc_json.encode("utf-8")
+
+        # Prepare multipart request with DoclingDocument JSON
+        files = {"files": ("document.json", BytesIO(doc_bytes), "application/json")}
+        data = self._build_chunking_data()
+
+        result = await self.client.submit_and_poll(
+            endpoint=endpoint,
+            files=files,
+            data=data,
+            name="document",
+        )
+
+        return result.get("chunks", [])
+
+    async def chunk(self, document: "DoclingDocument") -> list[Chunk]:
+        """Split the document into chunks with metadata via docling-serve.
+
+        Extracts structured metadata from the API response including:
+        - doc_item_refs: JSON pointer references to DocItems (e.g., "#/texts/5")
+        - headings: Section heading hierarchy
+        - labels: Semantic labels for each doc_item
+        - page_numbers: Page numbers where content appears
+
+        Args:
+            document: The DoclingDocument to be split into chunks.
+
+        Returns:
+            List of Chunk containing content and structured metadata.
+
+        Raises:
+            ValueError: If chunking fails or service is unavailable.
+        """
+        if document is None:
+            return []
+
+        raw_chunks = await self._call_chunk_api(document)
+        result: list[Chunk] = []
+
+        for chunk in raw_chunks:
+            text = chunk.get("text", "")
+
+            # doc_items from docling-serve is a list of ref strings like ["#/texts/1", "#/tables/0"]
+            doc_items = chunk.get("doc_items", [])
+            doc_item_refs: list[str] = []
+            labels: list[str] = []
+
+            for item in doc_items:
+                if isinstance(item, str):
+                    # docling-serve returns refs as strings directly
+                    doc_item_refs.append(item)
+                    # Resolve label from the document using the ref
+                    label = _resolve_label_from_document(item, document)
+                    if label:
+                        labels.append(label)
+                elif isinstance(item, dict):
+                    # Handle dict format if API ever returns it
+                    if "self_ref" in item:
+                        doc_item_refs.append(item["self_ref"])
+                    if "label" in item:
+                        labels.append(item["label"])
+
+            # Get headings directly from chunk
+            headings = chunk.get("headings")
+
+            # Get page numbers directly from chunk
+            page_numbers = chunk.get("page_numbers", [])
+
+            chunk_metadata = ChunkMetadata(
+                doc_item_refs=doc_item_refs,
+                headings=headings,
+                labels=labels,
+                page_numbers=sorted(page_numbers) if page_numbers else [],
+            )
+            result.append(
+                Chunk(
+                    content=text,
+                    metadata=chunk_metadata.model_dump(),
+                    order=len(result),
+                )
+            )
+
+        return result