From ab57c55eaa2e6fdc5ded6d13d606cb8c04b508eb Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 5 Dec 2025 10:31:05 +0200 Subject: [PATCH] Drop ChunkWithMetadata, use good old Chunk --- CHANGELOG.md | 4 +- haiku_rag_slim/haiku/rag/chunkers/base.py | 14 +--- .../haiku/rag/chunkers/docling_local.py | 14 ++-- .../haiku/rag/chunkers/docling_serve.py | 14 ++-- .../haiku/rag/store/repositories/chunk.py | 32 ++++----- tests/test_chunker.py | 70 ++++++++++--------- 6 files changed, 69 insertions(+), 79 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1bb35f38..87f3c16a 100644 --- a/CHANGELOG.md +++ b/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]` - - `ChunkWithMetadata` combines chunk text with `ChunkMetadata` (refs, labels, headings, page_numbers) +- **BREAKING: Chunker Interface**: `DocumentChunker.chunk()` now returns `list[Chunk]` instead of `list[str]` + - 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 diff --git a/haiku_rag_slim/haiku/rag/chunkers/base.py b/haiku_rag_slim/haiku/rag/chunkers/base.py index 783a2cb7..e1daf93b 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/base.py +++ b/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 - -class ChunkWithMetadata: - """A chunk with its text content and structured metadata.""" - - def __init__(self, text: str, metadata: ChunkMetadata): - self.text = text - self.metadata = metadata + from haiku.rag.store.models.chunk import Chunk 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: diff --git a/haiku_rag_slim/haiku/rag/chunkers/docling_local.py b/haiku_rag_slim/haiku/rag/chunkers/docling_local.py index 7e40e9a3..74f5c19f 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/docling_local.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py index 4c4bdf9d..efb2323f 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 1e79f3bb..0c562628 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/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: - if c.metadata.headings: - embedding_text = "\n".join(c.metadata.headings) + "\n" + c.text + for chunk in chunks: + chunk_meta = chunk.get_chunk_metadata() + 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( - zip(chunks_with_metadata, embeddings) - ): + for order, (chunk, embedding) in enumerate(zip(chunks, 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, - metadata=json.dumps(metadata_dict), + content=chunk.content, + metadata=json.dumps(chunk.metadata), order=order, vector=embedding, ) chunk_records.append(chunk_record) - chunk = Chunk( - id=chunk_id, - document_id=document_id, - content=chunk_with_meta.text, - metadata=metadata_dict, - order=order, - ) + chunk.id = chunk_id + chunk.document_id = document_id + chunk.order = order created_chunks.append(chunk) # Batch insert all chunks at once diff --git a/tests/test_chunker.py b/tests/test_chunker.py index f75b0a5a..104382c8 100644 --- a/tests/test_chunker.py +++ b/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("Column 1" in chunk.text for chunk in chunks_md) + assert any("|" in chunk.content 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][ - 0 - ] + table_content = [ + 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) - all_labels.extend(chunk.metadata.labels) - if chunk.metadata.headings: - all_headings.extend(chunk.metadata.headings) + meta = chunk.get_chunk_metadata() + all_refs.extend(meta.doc_item_refs) + all_labels.extend(meta.labels) + 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[1].text == "Chunk 2" + assert chunks[0].content == "Chunk 1" + 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].metadata.doc_item_refs == ["#/texts/0", "#/texts/1"] + assert chunks[0].content == "Chapter 1\nThis is content." + 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 chunks[0].metadata.headings == ["Chapter 1"] - assert chunks[0].metadata.page_numbers == [1] + assert meta0.labels == ["title", "text"] + assert meta0.headings == ["Chapter 1"] + assert meta0.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] + assert chunks[1].content == "Table content here." + meta1 = chunks[1].get_chunk_metadata() + assert meta1.doc_item_refs == ["#/tables/0"] + assert meta1.labels == ["table"] + 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}" )