feat(docling_serve): support async conversion jobs

Author: 2830500285

integrations/docling_serve/src/haystack_integrations/components/converters/docling_serve/__init__.py

@@ -2,6 +2,11 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from haystack_integrations.components.converters.docling_serve.converter import DoclingServeConverter, ExportType
+from haystack_integrations.components.converters.docling_serve.converter import (
+    ConversionMode,
+    DoclingServeConversionError,
+    DoclingServeConverter,
+    ExportType,
+)
 
-__all__ = ["DoclingServeConverter", "ExportType"]
+__all__ = ["ConversionMode", "DoclingServeConversionError", "DoclingServeConverter", "ExportType"]

integrations/docling_serve/src/haystack_integrations/components/converters/docling_serve/converter.py

@@ -2,8 +2,10 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import asyncio
 import json
 import mimetypes
+import time
 from enum import Enum
 from pathlib import Path
 from typing import Any
@@ -19,6 +21,10 @@
 
 _FILE_CONVERT_PATH = "/v1/convert/file"
 _SOURCE_CONVERT_PATH = "/v1/convert/source"
+_STATUS_POLL_PATH = "/v1/status/poll"
+_RESULT_PATH = "/v1/result"
+_TERMINAL_TASK_STATUSES = {"success", "failure"}
+_FAILED_CONVERSION_STATUSES = {"failure", "skipped"}
 
 
 class ExportType(str, Enum):
@@ -35,6 +41,22 @@ class ExportType(str, Enum):
     JSON = "json"
 
 
+class ConversionMode(str, Enum):
+    """
+    Execution mode for DoclingServe conversions.
+
+    - `SYNC`: Uses DoclingServe's synchronous conversion endpoints.
+    - `ASYNC`: Uses DoclingServe's async job endpoints and polls for completion.
+    """
+
+    SYNC = "sync"
+    ASYNC = "async"
+
+
+class DoclingServeConversionError(Exception):
+    """Raised when DoclingServe reports an async task or conversion failure."""
+
+
 def _is_url(source: str) -> bool:
     parsed = urlparse(source)
     return parsed.scheme in ("http", "https")
@@ -88,6 +110,9 @@ def __init__(
         convert_options: dict[str, Any] | None = None,
         timeout: float = 120.0,
         api_key: Secret | None = Secret.from_env_var("DOCLING_SERVE_API_KEY", strict=False),
+        mode: ConversionMode | str = ConversionMode.SYNC,
+        poll_interval: float = 2.0,
+        job_timeout: float = 600.0,
     ) -> None:
         """
         Initializes the DoclingServeConverter.
@@ -108,12 +133,29 @@ def __init__(
             API key for authenticating with a secured DoclingServe instance. Reads from the
             `DOCLING_SERVE_API_KEY` environment variable by default. Set to `None` to disable
             authentication.
+        :param mode:
+            Conversion mode. `sync` uses DoclingServe's synchronous endpoints. `async` submits
+            conversion jobs to DoclingServe's async endpoints and polls until completion.
+        :param poll_interval:
+            Maximum server-side long-poll wait in seconds when `mode="async"`.
+        :param job_timeout:
+            Maximum time in seconds to wait for each async conversion job.
         """
+        if poll_interval <= 0:
+            msg = "poll_interval must be greater than 0."
+            raise ValueError(msg)
+        if job_timeout <= 0:
+            msg = "job_timeout must be greater than 0."
+            raise ValueError(msg)
+
         self.base_url = base_url.rstrip("/")
         self.export_type = ExportType(export_type)
         self.convert_options = dict(convert_options) if convert_options else {}
         self.timeout = timeout
         self.api_key = api_key
+        self.mode = ConversionMode(mode)
+        self.poll_interval = poll_interval
+        self.job_timeout = job_timeout
 
     def to_dict(self) -> dict[str, Any]:
         """
@@ -129,6 +171,9 @@ def to_dict(self) -> dict[str, Any]:
             convert_options=self.convert_options,
             timeout=self.timeout,
             api_key=self.api_key.to_dict() if self.api_key else None,
+            mode=self.mode.value,
+            poll_interval=self.poll_interval,
+            job_timeout=self.job_timeout,
         )
 
     @classmethod
@@ -166,6 +211,41 @@ def _extract_content(self, data: dict[str, Any]) -> str | None:
             return json.dumps(content) if content is not None else None
         return None
 
+    def _raise_for_failed_conversion(self, data: dict[str, Any]) -> None:
+        status = data.get("status")
+        if status not in _FAILED_CONVERSION_STATUSES:
+            return
+
+        errors = data.get("errors") or []
+        details = "; ".join(
+            str(error.get("error_message", error)) if isinstance(error, dict) else str(error) for error in errors
+        )
+        msg = f"DoclingServe conversion finished with status '{status}'"
+        if details:
+            msg = f"{msg}: {details}"
+        raise DoclingServeConversionError(msg)
+
+    def _extract_task_id(self, data: dict[str, Any]) -> str:
+        task_id = data.get("task_id")
+        if not isinstance(task_id, str) or not task_id:
+            msg = "DoclingServe async task response did not include a task_id."
+            raise DoclingServeConversionError(msg)
+        return task_id
+
+    def _raise_for_failed_task(self, data: dict[str, Any]) -> None:
+        if data.get("task_status") != "failure":
+            return
+
+        error_message = data.get("error_message") or "DoclingServe async task failed."
+        raise DoclingServeConversionError(str(error_message))
+
+    def _async_source_payload(self, url: str) -> dict[str, Any]:
+        return {
+            "options": {**self.convert_options, "to_formats": [self._to_format()]},
+            "sources": [{"kind": "http", "url": url}],
+            "target": {"kind": "inbody"},
+        }
+
     def _post_file(self, client: httpx.Client, source: str | Path | ByteStream) -> dict[str, Any]:
         filename = _resolve_filename(source)
         file_bytes = source.data if isinstance(source, ByteStream) else Path(source).read_bytes()
@@ -184,6 +264,24 @@ def _post_file(self, client: httpx.Client, source: str | Path | ByteStream) -> d
         response.raise_for_status()
         return response.json()
 
+    def _submit_file_job(self, client: httpx.Client, source: str | Path | ByteStream) -> str:
+        filename = _resolve_filename(source)
+        file_bytes = source.data if isinstance(source, ByteStream) else Path(source).read_bytes()
+        mime_type = (
+            (source.mime_type or _guess_mime_type(filename))
+            if isinstance(source, ByteStream)
+            else _guess_mime_type(filename)
+        )
+        options = {**self.convert_options, "to_formats": self._to_format(), "target_type": "inbody"}
+        response = client.post(
+            f"{self.base_url}{_FILE_CONVERT_PATH}/async",
+            files={"files": (filename, file_bytes, mime_type)},
+            data=options,
+            headers=self._headers(),
+        )
+        response.raise_for_status()
+        return self._extract_task_id(response.json())
+
     def _post_url(self, client: httpx.Client, url: str) -> dict[str, Any]:
         payload: dict[str, Any] = {
             "options": {**self.convert_options, "to_formats": [self._to_format()]},
@@ -197,6 +295,64 @@ def _post_url(self, client: httpx.Client, url: str) -> dict[str, Any]:
         response.raise_for_status()
         return response.json()
 
+    def _submit_url_job(self, client: httpx.Client, url: str) -> str:
+        response = client.post(
+            f"{self.base_url}{_SOURCE_CONVERT_PATH}/async",
+            json=self._async_source_payload(url),
+            headers=self._headers(),
+        )
+        response.raise_for_status()
+        return self._extract_task_id(response.json())
+
+    def _poll_job_status(self, client: httpx.Client, task_id: str, wait: float) -> dict[str, Any]:
+        response = client.get(
+            f"{self.base_url}{_STATUS_POLL_PATH}/{task_id}",
+            params={"wait": wait},
+            headers=self._headers(),
+        )
+        response.raise_for_status()
+        return response.json()
+
+    def _wait_for_job(self, client: httpx.Client, task_id: str) -> None:
+        deadline = time.monotonic() + self.job_timeout
+        while True:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                msg = f"Timed out waiting for DoclingServe task {task_id} after {self.job_timeout:.2f}s."
+                raise DoclingServeConversionError(msg)
+
+            wait = min(self.poll_interval, remaining)
+            poll_started = time.monotonic()
+            status = self._poll_job_status(client, task_id, wait)
+            task_status = status.get("task_status")
+            if task_status in _TERMINAL_TASK_STATUSES:
+                self._raise_for_failed_task(status)
+                return
+
+            sleep_for = min(self.poll_interval, remaining) - (time.monotonic() - poll_started)
+            if sleep_for > 0:
+                time.sleep(sleep_for)
+
+    def _fetch_job_result(self, client: httpx.Client, task_id: str) -> dict[str, Any]:
+        response = client.get(
+            f"{self.base_url}{_RESULT_PATH}/{task_id}",
+            headers=self._headers(),
+        )
+        response.raise_for_status()
+        data = response.json()
+        self._raise_for_failed_conversion(data)
+        return data
+
+    def _post_file_job(self, client: httpx.Client, source: str | Path | ByteStream) -> dict[str, Any]:
+        task_id = self._submit_file_job(client, source)
+        self._wait_for_job(client, task_id)
+        return self._fetch_job_result(client, task_id)
+
+    def _post_url_job(self, client: httpx.Client, url: str) -> dict[str, Any]:
+        task_id = self._submit_url_job(client, url)
+        self._wait_for_job(client, task_id)
+        return self._fetch_job_result(client, task_id)
+
     async def _post_file_async(self, client: httpx.AsyncClient, source: str | Path | ByteStream) -> dict[str, Any]:
         filename = _resolve_filename(source)
         file_bytes = source.data if isinstance(source, ByteStream) else Path(source).read_bytes()
@@ -215,6 +371,24 @@ async def _post_file_async(self, client: httpx.AsyncClient, source: str | Path |
         response.raise_for_status()
         return response.json()
 
+    async def _submit_file_job_async(self, client: httpx.AsyncClient, source: str | Path | ByteStream) -> str:
+        filename = _resolve_filename(source)
+        file_bytes = source.data if isinstance(source, ByteStream) else Path(source).read_bytes()
+        mime_type = (
+            (source.mime_type or _guess_mime_type(filename))
+            if isinstance(source, ByteStream)
+            else _guess_mime_type(filename)
+        )
+        options = {**self.convert_options, "to_formats": self._to_format(), "target_type": "inbody"}
+        response = await client.post(
+            f"{self.base_url}{_FILE_CONVERT_PATH}/async",
+            files={"files": (filename, file_bytes, mime_type)},
+            data=options,
+            headers=self._headers(),
+        )
+        response.raise_for_status()
+        return self._extract_task_id(response.json())
+
     async def _post_url_async(self, client: httpx.AsyncClient, url: str) -> dict[str, Any]:
         payload: dict[str, Any] = {
             "options": {**self.convert_options, "to_formats": [self._to_format()]},
@@ -228,6 +402,64 @@ async def _post_url_async(self, client: httpx.AsyncClient, url: str) -> dict[str
         response.raise_for_status()
         return response.json()
 
+    async def _submit_url_job_async(self, client: httpx.AsyncClient, url: str) -> str:
+        response = await client.post(
+            f"{self.base_url}{_SOURCE_CONVERT_PATH}/async",
+            json=self._async_source_payload(url),
+            headers=self._headers(),
+        )
+        response.raise_for_status()
+        return self._extract_task_id(response.json())
+
+    async def _poll_job_status_async(self, client: httpx.AsyncClient, task_id: str, wait: float) -> dict[str, Any]:
+        response = await client.get(
+            f"{self.base_url}{_STATUS_POLL_PATH}/{task_id}",
+            params={"wait": wait},
+            headers=self._headers(),
+        )
+        response.raise_for_status()
+        return response.json()
+
+    async def _wait_for_job_async(self, client: httpx.AsyncClient, task_id: str) -> None:
+        deadline = time.monotonic() + self.job_timeout
+        while True:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                msg = f"Timed out waiting for DoclingServe task {task_id} after {self.job_timeout:.2f}s."
+                raise DoclingServeConversionError(msg)
+
+            wait = min(self.poll_interval, remaining)
+            poll_started = time.monotonic()
+            status = await self._poll_job_status_async(client, task_id, wait)
+            task_status = status.get("task_status")
+            if task_status in _TERMINAL_TASK_STATUSES:
+                self._raise_for_failed_task(status)
+                return
+
+            sleep_for = min(self.poll_interval, remaining) - (time.monotonic() - poll_started)
+            if sleep_for > 0:
+                await asyncio.sleep(sleep_for)
+
+    async def _fetch_job_result_async(self, client: httpx.AsyncClient, task_id: str) -> dict[str, Any]:
+        response = await client.get(
+            f"{self.base_url}{_RESULT_PATH}/{task_id}",
+            headers=self._headers(),
+        )
+        response.raise_for_status()
+        data = response.json()
+        self._raise_for_failed_conversion(data)
+        return data
+
+    async def _post_file_job_async(self, client: httpx.AsyncClient, source: str | Path | ByteStream) -> dict[str, Any]:
+        task_id = await self._submit_file_job_async(client, source)
+        await self._wait_for_job_async(client, task_id)
+        return await self._fetch_job_result_async(client, task_id)
+
+    async def _post_url_job_async(self, client: httpx.AsyncClient, url: str) -> dict[str, Any]:
+        task_id = await self._submit_url_job_async(client, url)
+        await self._wait_for_job_async(client, task_id)
+        return await self._fetch_job_result_async(client, task_id)
+
     @component.output_types(documents=list[Document])
     def run(
         self,
@@ -256,9 +488,17 @@ def run(
                 merged_meta = {**bytestream_meta, **source_meta}
                 try:
                     if isinstance(source, str) and _is_url(source):
-                        result = self._post_url(client, source)
+                        result = (
+                            self._post_url_job(client, source)
+                            if self.mode == ConversionMode.ASYNC
+                            else self._post_url(client, source)
+                        )
                     else:
-                        result = self._post_file(client, source)
+                        result = (
+                            self._post_file_job(client, source)
+                            if self.mode == ConversionMode.ASYNC
+                            else self._post_file(client, source)
+                        )
                     content = self._extract_content(result)
                     if content is not None:
                         documents.append(Document(content=content, meta=merged_meta))
@@ -277,6 +517,12 @@ def run(
                         source=source,
                         error=e,
                     )
+                except DoclingServeConversionError as e:
+                    logger.warning(
+                        "DoclingServe conversion failed for {source}. Skipping it. Error: {error}",
+                        source=source,
+                        error=e,
+                    )
 
         return {"documents": documents}
 
@@ -310,9 +556,17 @@ async def run_async(
                 merged_meta = {**bytestream_meta, **source_meta}
                 try:
                     if isinstance(source, str) and _is_url(source):
-                        result = await self._post_url_async(client, source)
+                        result = (
+                            await self._post_url_job_async(client, source)
+                            if self.mode == ConversionMode.ASYNC
+                            else await self._post_url_async(client, source)
+                        )
                     else:
-                        result = await self._post_file_async(client, source)
+                        result = (
+                            await self._post_file_job_async(client, source)
+                            if self.mode == ConversionMode.ASYNC
+                            else await self._post_file_async(client, source)
+                        )
                     content = self._extract_content(result)
                     if content is not None:
                         documents.append(Document(content=content, meta=merged_meta))
@@ -331,5 +585,11 @@ async def run_async(
                         source=source,
                         error=e,
                     )
+                except DoclingServeConversionError as e:
+                    logger.warning(
+                        "DoclingServe conversion failed for {source}. Skipping it. Error: {error}",
+                        source=source,
+                        error=e,
+                    )
 
         return {"documents": documents}