diff --git a/haiku_rag_slim/haiku/rag/chunkers/base.py b/haiku_rag_slim/haiku/rag/chunkers/base.py index e1daf93b..96c9fc9e 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/base.py +++ b/haiku_rag_slim/haiku/rag/chunkers/base.py @@ -15,7 +15,7 @@ class DocumentChunker(ABC): """ @abstractmethod - async def chunk(self, document: "DoclingDocument") -> list["Chunk"]: + async def chunk(self, document: "DoclingDocument | None") -> list["Chunk"]: """Split a document into chunks with metadata. Args: diff --git a/haiku_rag_slim/haiku/rag/chunkers/docling_local.py b/haiku_rag_slim/haiku/rag/chunkers/docling_local.py index fc07bb22..85b7f690 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/docling_local.py +++ b/haiku_rag_slim/haiku/rag/chunkers/docling_local.py @@ -1,3 +1,4 @@ +import asyncio from functools import cache from typing import TYPE_CHECKING, cast @@ -105,24 +106,12 @@ class DoclingLocalChunker(DocumentChunker): "Must be 'hybrid' or 'hierarchical'." ) - async def chunk(self, document: "DoclingDocument") -> list[Chunk]: - """Split the document into chunks with metadata. + def _chunk_sync(self, document: "DoclingDocument") -> list[Chunk]: + """Synchronous chunking helper (CPU-bound, no I/O). - Extracts structured metadata from each DocChunk including: - - doc_item_refs: JSON pointer references to DocItems (e.g., "#/texts/5") - - headings: Section heading hierarchy - - labels: Semantic labels for each doc_item (e.g., "paragraph", "table") - - page_numbers: Page numbers where content appears - - Args: - document: The DoclingDocument to be split into chunks. - - Returns: - List of Chunk containing content and structured metadata. + Runs the underlying HybridChunker/HierarchicalChunker and extracts + structured metadata from each DocChunk. """ - if document is None: - return [] - raw_chunks = list(self.chunker.chunk(document)) result: list[Chunk] = [] @@ -172,3 +161,23 @@ class DoclingLocalChunker(DocumentChunker): ) return result + + async def chunk(self, document: "DoclingDocument | None") -> list[Chunk]: + """Split the document into chunks with metadata. + + Extracts structured metadata from each DocChunk including: + - doc_item_refs: JSON pointer references to DocItems (e.g., "#/texts/5") + - headings: Section heading hierarchy + - labels: Semantic labels for each doc_item (e.g., "paragraph", "table") + - page_numbers: Page numbers where content appears + + Args: + document: The DoclingDocument to be split into chunks. + + Returns: + List of Chunk containing content and structured metadata. + """ + if document is None: + return [] + + return await asyncio.to_thread(self._chunk_sync, document) diff --git a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py index c0380134..1bfb7648 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py @@ -129,7 +129,7 @@ class DoclingServeChunker(DocumentChunker): return result.get("chunks", []) - async def chunk(self, document: "DoclingDocument") -> list[Chunk]: + async def chunk(self, document: "DoclingDocument | None") -> list[Chunk]: """Split the document into chunks with metadata via docling-serve. Extracts structured metadata from the API response including: diff --git a/tests/test_chunker.py b/tests/test_chunker.py index e59681a0..c1f1a833 100644 --- a/tests/test_chunker.py +++ b/tests/test_chunker.py @@ -55,6 +55,42 @@ async def test_local_chunker(qa_corpus: list[dict[str, str]]): assert abs(total_tokens - original_tokens) <= original_tokens * 0.1 +@pytest.mark.asyncio +async def test_local_chunker_none_document(): + """Test DoclingLocalChunker returns empty list for None document.""" + chunker = DoclingLocalChunker() + assert await chunker.chunk(None) == [] + + +@pytest.mark.asyncio +async def test_local_chunker_runs_off_event_loop_thread(): + """Chunking is CPU-bound; verify it runs in a worker thread.""" + import threading + from unittest.mock import patch + + chunker = DoclingLocalChunker() + event_loop_thread = threading.current_thread() + called_from: list[threading.Thread] = [] + + original = chunker._chunk_sync + + def recording_chunk_sync(self, document): + called_from.append(threading.current_thread()) + return original(document) + + converter = get_converter(Config) + doc = await converter.convert_text("# Hello\n\nWorld", name="test.md") + + with patch.object(DoclingLocalChunker, "_chunk_sync", recording_chunk_sync): + await chunker.chunk(doc) + + assert called_from, "_chunk_sync was never called" + assert called_from[0] is not event_loop_thread, ( + "_chunk_sync ran on the event-loop thread; " + "it must be dispatched via asyncio.to_thread" + ) + + @pytest.mark.asyncio async def test_local_chunker_custom_config(): """Test DoclingLocalChunker with custom configuration."""