HaikuRag.chunk() utility method
This commit is contained in:
parent
ace915473c
commit
4fbb4609b6
3 changed files with 85 additions and 0 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue