HaikuRag.chunk() utility method

This commit is contained in:
Yiorgis Gozadinos 2025-12-05 11:49:57 +02:00
parent ace915473c
commit 4fbb4609b6
No known key found for this signature in database
3 changed files with 85 additions and 0 deletions

View file

@ -43,6 +43,8 @@
- `client.convert("https://...")` - download and convert URL
- `client.convert("text content")` - convert plain text
- Supports `file://` URIs
- **New `chunk()` Method**: Chunk a DoclingDocument into Chunk objects
- `client.chunk(docling_doc)` - returns `list[Chunk]` without embeddings
### Changed

View file

@ -172,6 +172,20 @@ class HaikuRAG:
# Treat as text content
return await converter.convert_text(source)
async def chunk(self, docling_document: "DoclingDocument") -> list[Chunk]:
"""Chunk a DoclingDocument into Chunks.
Args:
docling_document: The DoclingDocument to chunk.
Returns:
List of Chunk objects (without embeddings, without document_id).
"""
from haiku.rag.chunkers import get_chunker
chunker = get_chunker(self._config)
return await chunker.chunk(docling_document)
async def _create_document_with_docling(
self,
docling_document,

View file

@ -1639,3 +1639,72 @@ async def test_client_convert_file_uri(temp_db_path):
assert isinstance(docling_doc, DoclingDocument)
markdown = docling_doc.export_to_markdown()
assert "URI file content" in markdown
# =============================================================================
# chunk() method tests
# =============================================================================
@pytest.mark.asyncio
async def test_client_chunk_basic(temp_db_path):
"""Test chunk() produces Chunk objects from DoclingDocument."""
async with HaikuRAG(temp_db_path, create=True) as client:
# First convert some text
docling_doc = await client.convert("This is test content for chunking.")
# Then chunk it
chunks = await client.chunk(docling_doc)
assert isinstance(chunks, list)
assert len(chunks) > 0
assert all(isinstance(c, Chunk) for c in chunks)
# Chunks should have content but no embedding yet
assert all(c.content for c in chunks)
assert all(c.embedding is None for c in chunks)
# Chunks should not have document_id yet (not stored)
assert all(c.document_id is None for c in chunks)
@pytest.mark.asyncio
async def test_client_chunk_preserves_metadata(temp_db_path):
"""Test chunk() preserves structured metadata from DoclingDocument."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Convert structured markdown
markdown = """# Chapter 1
This is the first paragraph.
## Section 1.1
This is a subsection.
"""
docling_doc = await client.convert(markdown)
chunks = await client.chunk(docling_doc)
assert len(chunks) > 0
# Check that at least some chunks have metadata
has_metadata = False
for chunk in chunks:
meta = chunk.get_chunk_metadata()
if meta.doc_item_refs or meta.headings:
has_metadata = True
break
assert has_metadata, "Chunks should have structured metadata"
@pytest.mark.asyncio
async def test_client_chunk_empty_document(temp_db_path):
"""Test chunk() with empty DoclingDocument."""
from docling_core.types.doc.document import DoclingDocument
async with HaikuRAG(temp_db_path, create=True) as client:
# Create an empty DoclingDocument
empty_doc = DoclingDocument(name="empty")
chunks = await client.chunk(empty_doc)
assert isinstance(chunks, list)
assert len(chunks) == 0