Drop ChunkWithMetadata, use good old Chunk

CHANGELOG.md

@@ -44,8 +44,8 @@
 - **BREAKING: `create_document()` API**: Removed `chunks` parameter
   - `create_document()` now always processes content (converts, chunks, embeds)
   - Use new `import_document()` for pre-processed documents with custom chunks
-- **BREAKING: Chunker Interface**: `DocumentChunker.chunk()` now returns `list[ChunkWithMetadata]` instead of `list[str]`
+- **BREAKING: Chunker Interface**: `DocumentChunker.chunk()` now returns `list[Chunk]` instead of `list[str]`
-  - `ChunkWithMetadata` combines chunk text with `ChunkMetadata` (refs, labels, headings, page_numbers)
+  - Chunks include structured metadata (doc_item_refs, labels, headings, page_numbers) in the `metadata` dict
   - All chunker implementations updated: `DoclingLocalChunker`, `DoclingServeChunker`
 - **Page Image Generation**: `generate_page_images=True` is now always enabled for local docling converter
   - Required for visual grounding features

haiku_rag_slim/haiku/rag/chunkers/base.py

@@ -1,18 +1,10 @@
 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
 
-
+    from haiku.rag.store.models.chunk import Chunk
-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):
@@ -23,14 +15,14 @@ class DocumentChunker(ABC):
     """
 
     @abstractmethod
-    async def chunk(self, document: "DoclingDocument") -> list[ChunkWithMetadata]:
+    async def chunk(self, document: "DoclingDocument") -> list["Chunk"]:
         """Split a document into chunks with metadata.
 
         Args:
             document: The DoclingDocument to chunk.
 
         Returns:
-            List of ChunkWithMetadata containing text and structured metadata
+            List of Chunk with content and structured metadata in the metadata dict
             (doc_item_refs, headings, labels, page_numbers).
 
         Raises:

haiku_rag_slim/haiku/rag/chunkers/docling_local.py

@@ -1,8 +1,8 @@
 from typing import TYPE_CHECKING, cast
 
-from haiku.rag.chunkers.base import ChunkWithMetadata, DocumentChunker
+from haiku.rag.chunkers.base import DocumentChunker
 from haiku.rag.config import AppConfig, Config
-from haiku.rag.store.models.chunk import ChunkMetadata
+from haiku.rag.store.models.chunk import Chunk, ChunkMetadata
 
 if TYPE_CHECKING:
     from docling_core.transforms.chunker.doc_chunk import DocMeta
@@ -95,7 +95,7 @@ class DoclingLocalChunker(DocumentChunker):
                 "Must be 'hybrid' or 'hierarchical'."
             )
 
-    async def chunk(self, document: "DoclingDocument") -> list[ChunkWithMetadata]:
+    async def chunk(self, document: "DoclingDocument") -> list[Chunk]:
         """Split the document into chunks with metadata.
 
         Extracts structured metadata from each DocChunk including:
@@ -108,13 +108,13 @@ class DoclingLocalChunker(DocumentChunker):
             document: The DoclingDocument to be split into chunks.
 
         Returns:
-            List of ChunkWithMetadata containing text and structured metadata.
+            List of Chunk containing content and structured metadata.
         """
         if document is None:
             return []
 
         raw_chunks = list(self.chunker.chunk(document))
-        result: list[ChunkWithMetadata] = []
+        result: list[Chunk] = []
 
         for chunk in raw_chunks:
             # Use raw chunk text - headings are stored separately in metadata
@@ -149,12 +149,12 @@ class DoclingLocalChunker(DocumentChunker):
             if meta and meta.headings:
                 headings = list(meta.headings)
 
-            metadata = ChunkMetadata(
+            chunk_metadata = ChunkMetadata(
                 doc_item_refs=doc_item_refs,
                 headings=headings,
                 labels=labels,
                 page_numbers=sorted(page_numbers),
             )
-            result.append(ChunkWithMetadata(text=text, metadata=metadata))
+            result.append(Chunk(content=text, metadata=chunk_metadata.model_dump()))
 
         return result

haiku_rag_slim/haiku/rag/chunkers/docling_serve.py

@@ -4,9 +4,9 @@ from typing import TYPE_CHECKING
 
 import httpx
 
-from haiku.rag.chunkers.base import ChunkWithMetadata, DocumentChunker
+from haiku.rag.chunkers.base import DocumentChunker
 from haiku.rag.config import AppConfig, Config
-from haiku.rag.store.models.chunk import ChunkMetadata
+from haiku.rag.store.models.chunk import Chunk, ChunkMetadata
 
 if TYPE_CHECKING:
     from docling_core.types.doc.document import DoclingDocument
@@ -133,7 +133,7 @@ class DoclingServeChunker(DocumentChunker):
         except Exception as e:
             raise ValueError(f"Failed to chunk via docling-serve: {e}")
 
-    async def chunk(self, document: "DoclingDocument") -> list[ChunkWithMetadata]:
+    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:
@@ -146,7 +146,7 @@ class DoclingServeChunker(DocumentChunker):
             document: The DoclingDocument to be split into chunks.
 
         Returns:
-            List of ChunkWithMetadata containing text and structured metadata.
+            List of Chunk containing content and structured metadata.
 
         Raises:
             ValueError: If chunking fails or service is unavailable.
@@ -155,7 +155,7 @@ class DoclingServeChunker(DocumentChunker):
             return []
 
         raw_chunks = await self._call_chunk_api(document)
-        result: list[ChunkWithMetadata] = []
+        result: list[Chunk] = []
 
         for chunk in raw_chunks:
             text = chunk.get("text", "")
@@ -186,12 +186,12 @@ class DoclingServeChunker(DocumentChunker):
             # Get page numbers directly from chunk
             page_numbers = chunk.get("page_numbers", [])
 
-            metadata = ChunkMetadata(
+            chunk_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))
+            result.append(Chunk(content=text, metadata=chunk_metadata.model_dump()))
 
         return result

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

@@ -225,16 +225,17 @@ class ChunkRepository:
                 )
                 raise e
 
-        chunks_with_metadata = await chunker.chunk(processed_document)
+        chunks = await chunker.chunk(processed_document)
 
         # Build embedding texts with headings prepended for better semantic search
         # The stored content stays raw, but embeddings capture section context
         embedding_texts = []
-        for c in chunks_with_metadata:
+        for chunk in chunks:
-            if c.metadata.headings:
+            chunk_meta = chunk.get_chunk_metadata()
-                embedding_text = "\n".join(c.metadata.headings) + "\n" + c.text
+            if chunk_meta.headings:
+                embedding_text = "\n".join(chunk_meta.headings) + "\n" + chunk.content
             else:
-                embedding_text = c.text
+                embedding_text = chunk.content
             embedding_texts.append(embedding_text)
         embeddings = await self.embedder.embed(embedding_texts)
 
@@ -242,31 +243,22 @@ class ChunkRepository:
         chunk_records = []
         created_chunks = []
 
-        for order, (chunk_with_meta, embedding) in enumerate(
+        for order, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
-            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_with_meta.text,
+                content=chunk.content,
-                metadata=json.dumps(metadata_dict),
+                metadata=json.dumps(chunk.metadata),
                 order=order,
                 vector=embedding,
             )
             chunk_records.append(chunk_record)
 
-            chunk = Chunk(
+            chunk.id = chunk_id
-                id=chunk_id,
+            chunk.document_id = document_id
-                document_id=document_id,
+            chunk.order = order
-                content=chunk_with_meta.text,
-                metadata=metadata_dict,
-                order=order,
-            )
             created_chunks.append(chunk)
 
         # Batch insert all chunks at once

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.text, add_special_tokens=False)
+        encoded_tokens = tokenizer.encode(chunk.content, 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.text.strip()) > 0
+        assert len(chunk.content.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.text for chunk in chunks_md)
+    assert any("|" in chunk.content for chunk in chunks_md)
-    assert any("Column 1" in chunk.text for chunk in chunks_md)
+    assert any("Column 1" in chunk.content for chunk in chunks_md)
 
     # Test with markdown tables disabled (narrative format)
     config_narrative = AppConfig()
@@ -137,9 +137,9 @@ 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.text for chunk in chunks_narrative if "Value" in chunk.text][
+    table_content = [
-        0
+        chunk.content for chunk in chunks_narrative if "Value" in chunk.content
-    ]
+    ][0]
     # Narrative format uses commas, not pipes for table structure
     assert "," in table_content and "|" not in table_content
 
@@ -172,10 +172,11 @@ Here is some background information.
     all_labels = []
     all_headings = []
     for chunk in chunks:
-        all_refs.extend(chunk.metadata.doc_item_refs)
+        meta = chunk.get_chunk_metadata()
-        all_labels.extend(chunk.metadata.labels)
+        all_refs.extend(meta.doc_item_refs)
-        if chunk.metadata.headings:
+        all_labels.extend(meta.labels)
-            all_headings.extend(chunk.metadata.headings)
+        if meta.headings:
+            all_headings.extend(meta.headings)
 
     # Should have JSON pointer refs like #/texts/0, #/tables/0
     assert len(all_refs) > 0
@@ -241,8 +242,8 @@ class TestDoclingServeChunker:
 
         chunks = await chunker.chunk(doc)
         assert len(chunks) == 2
-        assert chunks[0].text == "Chunk 1"
+        assert chunks[0].content == "Chunk 1"
-        assert chunks[1].text == "Chunk 2"
+        assert chunks[1].content == "Chunk 2"
         mock_client.post.assert_called_once()
 
     @pytest.mark.asyncio
@@ -434,19 +435,21 @@ This is content.
         assert len(chunks) == 2
 
         # First chunk - labels resolved from document
-        assert chunks[0].text == "Chapter 1\nThis is content."
+        assert chunks[0].content == "Chapter 1\nThis is content."
-        assert chunks[0].metadata.doc_item_refs == ["#/texts/0", "#/texts/1"]
+        meta0 = chunks[0].get_chunk_metadata()
+        assert meta0.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 meta0.labels == ["title", "text"]
-        assert chunks[0].metadata.headings == ["Chapter 1"]
+        assert meta0.headings == ["Chapter 1"]
-        assert chunks[0].metadata.page_numbers == [1]
+        assert meta0.page_numbers == [1]
 
         # Second chunk - label resolved from document
-        assert chunks[1].text == "Table content here."
+        assert chunks[1].content == "Table content here."
-        assert chunks[1].metadata.doc_item_refs == ["#/tables/0"]
+        meta1 = chunks[1].get_chunk_metadata()
-        assert chunks[1].metadata.labels == ["table"]
+        assert meta1.doc_item_refs == ["#/tables/0"]
-        assert chunks[1].metadata.headings == ["Chapter 1", "Section 1.1"]
+        assert meta1.labels == ["table"]
-        assert chunks[1].metadata.page_numbers == [1, 2]
+        assert meta1.headings == ["Chapter 1", "Section 1.1"]
+        assert meta1.page_numbers == [1, 2]
 
 
 def is_docling_serve_available(base_url: str = "http://localhost:5001") -> bool:
@@ -508,28 +511,31 @@ async def test_local_and_serve_chunkers_produce_same_output():
     # 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"
+        assert local.content == serve.content, f"Chunk {i} content mismatch"
+
+        local_meta = local.get_chunk_metadata()
+        serve_meta = serve.get_chunk_metadata()
 
         # doc_item_refs should match
-        assert local.metadata.doc_item_refs == serve.metadata.doc_item_refs, (
+        assert local_meta.doc_item_refs == serve_meta.doc_item_refs, (
             f"Chunk {i} doc_item_refs mismatch: "
-            f"local={local.metadata.doc_item_refs}, serve={serve.metadata.doc_item_refs}"
+            f"local={local_meta.doc_item_refs}, serve={serve_meta.doc_item_refs}"
         )
 
         # Labels should match (now that serve resolves from document)
-        assert local.metadata.labels == serve.metadata.labels, (
+        assert local_meta.labels == serve_meta.labels, (
             f"Chunk {i} labels mismatch: "
-            f"local={local.metadata.labels}, serve={serve.metadata.labels}"
+            f"local={local_meta.labels}, serve={serve_meta.labels}"
         )
 
         # Headings should match
-        assert local.metadata.headings == serve.metadata.headings, (
+        assert local_meta.headings == serve_meta.headings, (
             f"Chunk {i} headings mismatch: "
-            f"local={local.metadata.headings}, serve={serve.metadata.headings}"
+            f"local={local_meta.headings}, serve={serve_meta.headings}"
         )
 
         # Page numbers should match
-        assert local.metadata.page_numbers == serve.metadata.page_numbers, (
+        assert local_meta.page_numbers == serve_meta.page_numbers, (
             f"Chunk {i} page_numbers mismatch: "
-            f"local={local.metadata.page_numbers}, serve={serve.metadata.page_numbers}"
+            f"local={local_meta.page_numbers}, serve={serve_meta.page_numbers}"
         )