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("https://...")` - download and convert URL
|
||||||
- `client.convert("text content")` - convert plain text
|
- `client.convert("text content")` - convert plain text
|
||||||
- Supports `file://` URIs
|
- Supports `file://` URIs
|
||||||
|
- **New `chunk()` Method**: Chunk a DoclingDocument into Chunk objects
|
||||||
|
- `client.chunk(docling_doc)` - returns `list[Chunk]` without embeddings
|
||||||
|
|
||||||
### Changed
|
### Changed
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -172,6 +172,20 @@ class HaikuRAG:
|
||||||
# Treat as text content
|
# Treat as text content
|
||||||
return await converter.convert_text(source)
|
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(
|
async def _create_document_with_docling(
|
||||||
self,
|
self,
|
||||||
docling_document,
|
docling_document,
|
||||||
|
|
|
||||||
|
|
@ -1639,3 +1639,72 @@ async def test_client_convert_file_uri(temp_db_path):
|
||||||
assert isinstance(docling_doc, DoclingDocument)
|
assert isinstance(docling_doc, DoclingDocument)
|
||||||
markdown = docling_doc.export_to_markdown()
|
markdown = docling_doc.export_to_markdown()
|
||||||
assert "URI file content" in 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