From 91de82ec143a86991379b2cb157eac557914d6dc Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 14 Nov 2025 17:26:54 +0200 Subject: [PATCH] docling-serve chunker implementation. Lacks tests --- .../haiku/rag/chunkers/docling_serve.py | 92 +++++++++++++++++-- 1 file changed, 86 insertions(+), 6 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py index 9547cf73..10e6dd2b 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py +++ b/haiku_rag_slim/haiku/rag/chunkers/docling_serve.py @@ -1,7 +1,10 @@ +from io import BytesIO from typing import TYPE_CHECKING +import requests + from haiku.rag.chunkers.base import DocumentChunker -from haiku.rag.config import AppConfig +from haiku.rag.config import AppConfig, Config if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument @@ -10,15 +13,26 @@ if TYPE_CHECKING: class DoclingServeChunker(DocumentChunker): """Remote document chunker using docling-serve API. - Placeholder - will be implemented in a future commit. + Sends DoclingDocument JSON to docling-serve for chunking. Supports both hybrid + and hierarchical chunking strategies via remote API. + + Args: + config: Application configuration containing docling-serve settings. """ - def __init__(self, config: AppConfig): - raise NotImplementedError("DoclingServeChunker not yet implemented") + def __init__(self, config: AppConfig = Config): + self.config = config + self.base_url = config.providers.docling_serve.base_url.rstrip("/") + self.api_key = config.providers.docling_serve.api_key + self.timeout = config.providers.docling_serve.timeout + self.chunker_type = config.processing.chunker_type async def chunk(self, document: "DoclingDocument") -> list[str]: """Split the document into chunks via docling-serve. + Exports the DoclingDocument to JSON and sends it to docling-serve's chunking + endpoint. The API will chunk the document and return the text chunks. + Args: document: The DoclingDocument to be split into chunks. @@ -26,6 +40,72 @@ class DoclingServeChunker(DocumentChunker): A list of text chunks with semantic boundaries. Raises: - NotImplementedError: This chunker is not yet implemented. + ValueError: If chunking fails or service is unavailable. """ - raise NotImplementedError("DoclingServeChunker not yet implemented") + if document is None: + return [] + + try: + # Determine endpoint based on chunker_type + if self.chunker_type == "hierarchical": + url = f"{self.base_url}/v1/chunk/hierarchical/file" + else: + url = f"{self.base_url}/v1/chunk/hybrid/file" + + # Export document to JSON + doc_json = document.model_dump_json() + doc_bytes = doc_json.encode("utf-8") + + # Prepare multipart request with DoclingDocument JSON + files = {"files": ("document.json", BytesIO(doc_bytes), "application/json")} + + # Build form data with chunking parameters + data = { + "chunking_max_tokens": str(self.config.processing.chunk_size), + "chunking_tokenizer": self.config.processing.chunking_tokenizer, + "chunking_merge_peers": str( + self.config.processing.chunking_merge_peers + ).lower(), + "chunking_use_markdown_tables": str( + self.config.processing.chunking_use_markdown_tables + ).lower(), + } + + headers = {} + if self.api_key: + headers["X-Api-Key"] = self.api_key + + response = requests.post( + url, + files=files, + data=data, + headers=headers, + timeout=self.timeout, + ) + + response.raise_for_status() + + result = response.json() + + # Extract text from chunks + chunks = result.get("chunks", []) + return [chunk["text"] for chunk in chunks] + + except requests.exceptions.ConnectionError as e: + raise ValueError( + f"Could not connect to docling-serve at {self.base_url}. " + f"Ensure the service is running and accessible. Error: {e}" + ) + except requests.exceptions.Timeout as e: + raise ValueError( + f"Request to docling-serve timed out after {self.timeout}s. " + f"Consider increasing the timeout in configuration. Error: {e}" + ) + except requests.exceptions.HTTPError as e: + if e.response.status_code == 401: + raise ValueError( + "Authentication failed. Check your API key configuration." + ) + raise ValueError(f"HTTP error from docling-serve: {e}") + except Exception as e: + raise ValueError(f"Failed to chunk via docling-serve: {e}")