Merge pull request #458 from mcdonc/thread-chunker-chunk

Thread chunker.chunk() off the asyncio event loop

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:

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]:
+    def _chunk_sync(self, document: "DoclingDocument") -> list[Chunk]:
-        """Split the document into chunks with metadata.
+        """Synchronous chunking helper (CPU-bound, no I/O).
 
-        Extracts structured metadata from each DocChunk including:
+        Runs the underlying HybridChunker/HierarchicalChunker and extracts
-        - doc_item_refs: JSON pointer references to DocItems (e.g., "#/texts/5")
+        structured metadata from each DocChunk.
-        - 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 []
-
         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)

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:

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."""