Chunkers return chunks with metadata, i.e. (refs, labels, headings, page_numbers)

.gitignore

@@ -5,6 +5,7 @@ build/
 dist/
 wheels/
 *.egg-info
+**/.DS_Store
 
 # Virtual environments
 .venv

haiku_rag_slim/haiku/rag/chunkers/base.py

@@ -1,10 +1,20 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
 
+from haiku.rag.store.models.chunk import ChunkMetadata
+
 if TYPE_CHECKING:
     from docling_core.types.doc.document import DoclingDocument
 
 
+class ChunkWithMetadata:
+    """A chunk with its text content and structured metadata."""
+
+    def __init__(self, text: str, metadata: ChunkMetadata):
+        self.text = text
+        self.metadata = metadata
+
+
 class DocumentChunker(ABC):
     """Abstract base class for document chunkers.
 
@@ -13,14 +23,15 @@ class DocumentChunker(ABC):
     """
 
     @abstractmethod
-    async def chunk(self, document: "DoclingDocument") -> list[str]:
+    async def chunk(self, document: "DoclingDocument") -> list[ChunkWithMetadata]:
-        """Split a document into chunks.
+        """Split a document into chunks with metadata.
 
         Args:
             document: The DoclingDocument to chunk.
 
         Returns:
-            List of text chunks with semantic boundaries preserved.
+            List of ChunkWithMetadata containing text and structured metadata
+            (doc_item_refs, headings, labels, page_numbers).
 
         Raises:
             ValueError: If chunking fails.

haiku_rag_slim/haiku/rag/chunkers/docling_local.py

@@ -1,9 +1,11 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
-from haiku.rag.chunkers.base import DocumentChunker
+from haiku.rag.chunkers.base import ChunkWithMetadata, DocumentChunker
 from haiku.rag.config import AppConfig, Config
+from haiku.rag.store.models.chunk import ChunkMetadata
 
 if TYPE_CHECKING:
+    from docling_core.transforms.chunker.doc_chunk import DocMeta
     from docling_core.types.doc.document import DoclingDocument
 
 
@@ -93,18 +95,64 @@ class DoclingLocalChunker(DocumentChunker):
                 "Must be 'hybrid' or 'hierarchical'."
             )
 
-    async def chunk(self, document: "DoclingDocument") -> list[str]:
+    async def chunk(self, document: "DoclingDocument") -> list[ChunkWithMetadata]:
-        """Split the document into chunks using docling's structure-aware chunking.
+        """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:
-            A list of text chunks with semantic boundaries.
+            List of ChunkWithMetadata containing text and structured metadata.
         """
         if document is None:
             return []
 
-        # Chunk using docling's hybrid chunker
+        raw_chunks = list(self.chunker.chunk(document))
-        chunks = list(self.chunker.chunk(document))
+        result: list[ChunkWithMetadata] = []
-        return [self.chunker.contextualize(chunk) for chunk in chunks]
+
+        for chunk in raw_chunks:
+            text = self.chunker.contextualize(chunk)
+
+            # 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)
+
+            metadata = ChunkMetadata(
+                doc_item_refs=doc_item_refs,
+                headings=headings,
+                labels=labels,
+                page_numbers=sorted(page_numbers),
+            )
+            result.append(ChunkWithMetadata(text=text, metadata=metadata))
+
+        return result

haiku_rag_slim/haiku/rag/chunkers/docling_serve.py

@@ -1,14 +1,48 @@
+import re
 from io import BytesIO
 from typing import TYPE_CHECKING
 
 import httpx
 
-from haiku.rag.chunkers.base import DocumentChunker
+from haiku.rag.chunkers.base import ChunkWithMetadata, DocumentChunker
 from haiku.rag.config import AppConfig, Config
+from haiku.rag.store.models.chunk import 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.
@@ -27,54 +61,48 @@ class DoclingServeChunker(DocumentChunker):
         self.timeout = config.providers.docling_serve.timeout
         self.chunker_type = config.processing.chunker_type
 
-    async def chunk(self, document: "DoclingDocument") -> list[str]:
+    async def _call_chunk_api(self, document: "DoclingDocument") -> list[dict]:
-        """Split the document into chunks via docling-serve.
+        """Call docling-serve chunking API and return raw chunk data.
-
-        Exports the DoclingDocument to JSON and sends it to docling-serve's chunking
-        endpoint. The API will chunk the document and return the text chunks.
 
         Args:
             document: The DoclingDocument to be split into chunks.
 
         Returns:
-            A list of text chunks with semantic boundaries.
+            List of chunk dictionaries from API response.
 
         Raises:
             ValueError: If chunking fails or service is unavailable.
         """
-        if document is None:
+        # Determine endpoint based on chunker_type
-            return []
+        if self.chunker_type == "hierarchical":
+            url = f"{self.base_url}/v1/chunk/hierarchical/file"
+        else:
+            url = f"{self.base_url}/v1/chunk/hybrid/file"
+
+        # 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")}
+
+        # Build form data with chunking parameters
+        data = {
+            "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(),
+        }
+
+        headers = {}
+        if self.api_key:
+            headers["X-Api-Key"] = self.api_key
 
         try:
-            # Determine endpoint based on chunker_type
-            if self.chunker_type == "hierarchical":
-                url = f"{self.base_url}/v1/chunk/hierarchical/file"
-            else:
-                url = f"{self.base_url}/v1/chunk/hybrid/file"
-
-            # 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")}
-
-            # Build form data with chunking parameters
-            data = {
-                "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(),
-            }
-
-            headers = {}
-            if self.api_key:
-                headers["X-Api-Key"] = self.api_key
-
             async with httpx.AsyncClient(timeout=self.timeout) as client:
                 response = await client.post(
                     url,
@@ -84,10 +112,7 @@ class DoclingServeChunker(DocumentChunker):
                 )
                 response.raise_for_status()
                 result = response.json()
-
+                return result.get("chunks", [])
-            # Extract text from chunks
-            chunks = result.get("chunks", [])
-            return [chunk["text"] for chunk in chunks]
 
         except httpx.ConnectError as e:
             raise ValueError(
@@ -107,3 +132,66 @@ class DoclingServeChunker(DocumentChunker):
             raise ValueError(f"HTTP error from docling-serve: {e}")
         except Exception as e:
             raise ValueError(f"Failed to chunk via docling-serve: {e}")
+
+    async def chunk(self, document: "DoclingDocument") -> list[ChunkWithMetadata]:
+        """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 ChunkWithMetadata containing text 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[ChunkWithMetadata] = []
+
+        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", [])
+
+            metadata = ChunkMetadata(
+                doc_item_refs=doc_item_refs,
+                headings=headings,
+                labels=labels,
+                page_numbers=sorted(page_numbers) if page_numbers else [],
+            )
+            result.append(ChunkWithMetadata(text=text, metadata=metadata))
+
+        return result

haiku_rag_slim/haiku/rag/store/repositories/chunk.py

@@ -225,22 +225,28 @@ class ChunkRepository:
                 )
                 raise e
 
-        chunk_texts = await chunker.chunk(processed_document)
+        chunks_with_metadata = await chunker.chunk(processed_document)
 
+        chunk_texts = [c.text for c in chunks_with_metadata]
         embeddings = await self.embedder.embed(chunk_texts)
 
         # Prepare all chunk records for batch insertion
         chunk_records = []
         created_chunks = []
 
-        for order, (chunk_text, embedding) in enumerate(zip(chunk_texts, embeddings)):
+        for order, (chunk_with_meta, embedding) in enumerate(
+            zip(chunks_with_metadata, embeddings)
+        ):
             chunk_id = str(uuid4())
 
+            # Convert ChunkMetadata to dict for storage
+            metadata_dict = chunk_with_meta.metadata.model_dump()
+
             chunk_record = self.store.ChunkRecord(
                 id=chunk_id,
                 document_id=document_id,
-                content=chunk_text,
+                content=chunk_with_meta.text,
-                metadata=json.dumps({}),
+                metadata=json.dumps(metadata_dict),
                 order=order,
                 vector=embedding,
             )
@@ -249,8 +255,8 @@ class ChunkRepository:
             chunk = Chunk(
                 id=chunk_id,
                 document_id=document_id,
-                content=chunk_text,
+                content=chunk_with_meta.text,
-                metadata={},
+                metadata=metadata_dict,
                 order=order,
             )
             created_chunks.append(chunk)

tests/test_chunker.py

@@ -32,7 +32,7 @@ async def test_local_chunker(qa_corpus: Dataset):
     # Ensure that chunks are reasonably sized (allowing more flexibility for structure-aware chunking)
     total_tokens = 0
     for chunk in chunks:
-        encoded_tokens = tokenizer.encode(chunk, add_special_tokens=False)
+        encoded_tokens = tokenizer.encode(chunk.text, add_special_tokens=False)
         token_count = len(encoded_tokens)
         total_tokens += token_count
 
@@ -95,7 +95,7 @@ async def test_local_chunker_hierarchical(qa_corpus: Dataset):
     assert len(chunks) > 0
     # Each chunk should be non-empty
     for chunk in chunks:
-        assert len(chunk.strip()) > 0
+        assert len(chunk.text.strip()) > 0
 
 
 def test_local_chunker_invalid_type():
@@ -127,8 +127,8 @@ async def test_local_chunker_markdown_tables():
     chunks_md = await chunker_md.chunk(doc)
 
     # Should contain markdown table format
-    assert any("|" in chunk for chunk in chunks_md)
+    assert any("|" in chunk.text for chunk in chunks_md)
-    assert any("Column 1" in chunk for chunk in chunks_md)
+    assert any("Column 1" in chunk.text for chunk in chunks_md)
 
     # Test with markdown tables disabled (narrative format)
     config_narrative = AppConfig()
@@ -137,11 +137,59 @@ async def test_local_chunker_markdown_tables():
     chunks_narrative = await chunker_narrative.chunk(doc)
 
     # Should contain narrative format (no pipe characters in table)
-    table_content = [chunk for chunk in chunks_narrative if "Value" in chunk][0]
+    table_content = [chunk.text for chunk in chunks_narrative if "Value" in chunk.text][
+        0
+    ]
     # Narrative format uses commas, not pipes for table structure
     assert "," in table_content and "|" not in table_content
 
 
+@pytest.mark.asyncio
+async def test_local_chunker_metadata_extraction():
+    """Test that DoclingLocalChunker extracts metadata correctly."""
+    sample_md = """# Chapter 1: Introduction
+
+This is the first paragraph of the introduction.
+
+## Section 1.1: Background
+
+Here is some background information.
+
+| Header 1 | Header 2 |
+|----------|----------|
+| Value 1  | Value 2  |
+"""
+    converter = get_converter(Config)
+    doc = await converter.convert_text(sample_md, name="test.md")
+
+    chunker = DoclingLocalChunker()
+    chunks = await chunker.chunk(doc)
+
+    assert len(chunks) > 0
+
+    # Check that at least one chunk has doc_item_refs
+    all_refs = []
+    all_labels = []
+    all_headings = []
+    for chunk in chunks:
+        all_refs.extend(chunk.metadata.doc_item_refs)
+        all_labels.extend(chunk.metadata.labels)
+        if chunk.metadata.headings:
+            all_headings.extend(chunk.metadata.headings)
+
+    # Should have JSON pointer refs like #/texts/0, #/tables/0
+    assert len(all_refs) > 0
+    assert any(ref.startswith("#/") for ref in all_refs)
+
+    # Should have labels
+    assert len(all_labels) > 0
+    assert "text" in all_labels or "table" in all_labels
+
+    # Should have headings
+    assert len(all_headings) > 0
+    assert any("Chapter" in h or "Section" in h for h in all_headings)
+
+
 def test_get_chunker_docling_serve():
     """Test factory returns DoclingServeChunker for docling-serve."""
     config = AppConfig()
@@ -193,8 +241,8 @@ class TestDoclingServeChunker:
 
         chunks = await chunker.chunk(doc)
         assert len(chunks) == 2
-        assert chunks[0] == "Chunk 1"
+        assert chunks[0].text == "Chunk 1"
-        assert chunks[1] == "Chunk 2"
+        assert chunks[1].text == "Chunk 2"
         mock_client.post.assert_called_once()
 
     @pytest.mark.asyncio
@@ -332,3 +380,156 @@ class TestDoclingServeChunker:
 
         with pytest.raises(ValueError, match="Authentication failed"):
             await chunker.chunk(doc)
+
+    @pytest.mark.asyncio
+    @patch("haiku.rag.chunkers.docling_serve.httpx.AsyncClient")
+    async def test_chunk_metadata_extraction(self, mock_client_class, chunker):
+        """Test that metadata is correctly extracted from API response.
+
+        Labels are resolved from the DoclingDocument using the refs, so we need
+        to create a document with matching structure for the mocked API response.
+        """
+        mock_response = Mock()
+        mock_response.status_code = 200
+        # docling-serve returns doc_items as list of ref strings
+        # We'll reference texts[0], texts[1], and tables[0]
+        mock_response.json.return_value = {
+            "chunks": [
+                {
+                    "text": "Chapter 1\nThis is content.",
+                    "doc_items": ["#/texts/0", "#/texts/1"],
+                    "headings": ["Chapter 1"],
+                    "page_numbers": [1],
+                },
+                {
+                    "text": "Table content here.",
+                    "doc_items": ["#/tables/0"],
+                    "headings": ["Chapter 1", "Section 1.1"],
+                    "page_numbers": [1, 2],
+                },
+            ]
+        }
+        mock_response.raise_for_status = Mock()
+
+        mock_client = AsyncMock()
+        mock_client.post.return_value = mock_response
+        mock_client_class.return_value.__aenter__.return_value = mock_client
+
+        # Create a document with texts and tables that match the mocked refs
+        converter = get_converter(Config)
+        doc = await converter.convert_text(
+            """# Chapter 1
+
+This is content.
+
+| Col1 | Col2 |
+|------|------|
+| A    | B    |
+""",
+            name="test.md",
+        )
+
+        chunks = await chunker.chunk(doc)
+
+        assert len(chunks) == 2
+
+        # First chunk - labels resolved from document
+        assert chunks[0].text == "Chapter 1\nThis is content."
+        assert chunks[0].metadata.doc_item_refs == ["#/texts/0", "#/texts/1"]
+        # texts[0] is title (# heading), texts[1] is text (paragraph)
+        assert chunks[0].metadata.labels == ["title", "text"]
+        assert chunks[0].metadata.headings == ["Chapter 1"]
+        assert chunks[0].metadata.page_numbers == [1]
+
+        # Second chunk - label resolved from document
+        assert chunks[1].text == "Table content here."
+        assert chunks[1].metadata.doc_item_refs == ["#/tables/0"]
+        assert chunks[1].metadata.labels == ["table"]
+        assert chunks[1].metadata.headings == ["Chapter 1", "Section 1.1"]
+        assert chunks[1].metadata.page_numbers == [1, 2]
+
+
+def is_docling_serve_available(base_url: str = "http://localhost:5001") -> bool:
+    """Check if docling-serve is running and accessible."""
+    import requests
+
+    try:
+        response = requests.get(f"{base_url}/health", timeout=2)
+        return response.status_code == 200
+    except Exception:
+        return False
+
+
+@pytest.mark.skipif(
+    not is_docling_serve_available(),
+    reason="docling-serve not available at http://localhost:5001",
+)
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_local_and_serve_chunkers_produce_same_output():
+    """Test that local and serve chunkers produce identical output for the same document.
+
+    This integration test requires docling-serve to be running on localhost:5001.
+
+    Note: Labels are resolved from the DoclingDocument since docling-serve API
+    only returns ref strings, not labels. See:
+    https://github.com/docling-project/docling-serve/issues/448
+    """
+    from pathlib import Path
+
+    from haiku.rag.chunkers.docling_local import DoclingLocalChunker
+    from haiku.rag.chunkers.docling_serve import DoclingServeChunker
+    from haiku.rag.converters.docling_serve import DoclingServeConverter
+
+    # Use docling-serve to convert the PDF (ensures same conversion for both chunkers)
+    converter = DoclingServeConverter(Config)
+    pdf_path = Path("tests/data/doclaynet.pdf")
+    doc = await converter.convert_file(pdf_path)
+
+    # Create both chunkers with same config
+    config = AppConfig()
+    config.processing.chunk_size = 256
+    config.processing.chunker_type = "hybrid"
+    config.processing.chunking_merge_peers = True
+    config.processing.chunking_use_markdown_tables = True
+
+    local_chunker = DoclingLocalChunker(config)
+    serve_chunker = DoclingServeChunker(config)
+
+    # Chunk with both
+    local_chunks = await local_chunker.chunk(doc)
+    serve_chunks = await serve_chunker.chunk(doc)
+
+    # Same number of chunks
+    assert len(local_chunks) == len(serve_chunks), (
+        f"Chunk count mismatch: local={len(local_chunks)}, serve={len(serve_chunks)}"
+    )
+
+    # Compare each chunk
+    for i, (local, serve) in enumerate(zip(local_chunks, serve_chunks)):
+        # Text should match
+        assert local.text == serve.text, f"Chunk {i} text mismatch"
+
+        # doc_item_refs should match
+        assert local.metadata.doc_item_refs == serve.metadata.doc_item_refs, (
+            f"Chunk {i} doc_item_refs mismatch: "
+            f"local={local.metadata.doc_item_refs}, serve={serve.metadata.doc_item_refs}"
+        )
+
+        # Labels should match (now that serve resolves from document)
+        assert local.metadata.labels == serve.metadata.labels, (
+            f"Chunk {i} labels mismatch: "
+            f"local={local.metadata.labels}, serve={serve.metadata.labels}"
+        )
+
+        # Headings should match
+        assert local.metadata.headings == serve.metadata.headings, (
+            f"Chunk {i} headings mismatch: "
+            f"local={local.metadata.headings}, serve={serve.metadata.headings}"
+        )
+
+        # Page numbers should match
+        assert local.metadata.page_numbers == serve.metadata.page_numbers, (
+            f"Chunk {i} page_numbers mismatch: "
+            f"local={local.metadata.page_numbers}, serve={serve.metadata.page_numbers}"
+        )