danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Drop ChunkWithMetadata, use good old Chunk

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-12-05 10:31:05 +02:00
parent 33e6a36290
commit ab57c55eaa
No known key found for this signature in database
6 changed files with 69 additions and 79 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
CHANGELOG.md
View file

@ -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

14
haiku_rag_slim/haiku/rag/chunkers/base.py
View file

@ -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:

14
haiku_rag_slim/haiku/rag/chunkers/docling_local.py
View file

@ -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

14
haiku_rag_slim/haiku/rag/chunkers/docling_serve.py
View file

@ -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

32
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -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

70
tests/test_chunker.py
View file

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