retrieve picture image bytes via referenced+zip path, fix docling-serve not returning images even when set to do so. Remove xfail from relevant test

haiku_rag_slim/haiku/rag/converters/docling_serve.py

@@ -84,10 +84,38 @@ class DoclingServeConverter(DocumentConverter):
 
         raise ValueError(f"Unsupported VLM provider: {model.provider}")
 
+    def _picture_images_enabled(self) -> bool:
+        """Whether the conversion should produce embedded picture images.
+
+        True if the user explicitly enabled ``generate_picture_images`` or if
+        ``picture_description.enabled`` is on (the VLM needs the picture bytes).
+        """
+        opts = self.config.processing.conversion_options
+        return opts.generate_picture_images or opts.picture_description.enabled
+
     def _build_conversion_data(self) -> dict[str, str | list[str]]:
-        """Build form data for conversion request."""
+        """Build form data for conversion request.
+
+        When picture images are requested, switch to
+        ``image_export_mode="referenced"`` + ``target_type="zip"``. Per
+        docling-jobkit, ``generate_picture_images=True`` is only set when
+        ``image_export_mode == "referenced"``; with ``"embedded"`` the server
+        emits page images but leaves ``PictureItem.image=None`` (upstream
+        issue docling-project/docling-serve#576). The zip target then ships
+        the picture (and page) bytes as files under ``artifacts/`` which the
+        caller rehydrates back into ``data:`` URIs for parity with the local
+        converter.
+        """
         opts = self.config.processing.conversion_options
         pic_desc = opts.picture_description
+        picture_images_enabled = self._picture_images_enabled()
+
+        if picture_images_enabled:
+            image_export_mode = "referenced"
+        else:
+            image_export_mode = (
+                "embedded" if opts.generate_page_images else "placeholder"
+            )
 
         data: dict[str, str | list[str]] = {
             "to_formats": "json",
@@ -98,15 +126,14 @@ class DoclingServeConverter(DocumentConverter):
             "table_mode": opts.table_mode,
             "table_cell_matching": str(opts.table_cell_matching).lower(),
             "images_scale": str(opts.images_scale),
-            "image_export_mode": "embedded"
-            if opts.generate_page_images
-            else "placeholder",
-            "include_images": str(
-                opts.generate_picture_images or pic_desc.enabled
-            ).lower(),
+            "image_export_mode": image_export_mode,
+            "include_images": str(picture_images_enabled).lower(),
             "do_picture_description": str(pic_desc.enabled).lower(),
         }
 
+        if picture_images_enabled:
+            data["target_type"] = "zip"
+
         if opts.ocr_lang:
             data["ocr_lang"] = opts.ocr_lang
 
@@ -125,6 +152,63 @@ class DoclingServeConverter(DocumentConverter):
 
         return data
 
+    def _parse_zip_to_docling(self, zip_bytes: bytes, name: str) -> "DoclingDocument":
+        """Parse a docling-serve ``target_type=zip`` response.
+
+        The archive contains the document JSON at the root plus an
+        ``artifacts/`` directory holding referenced image files (PNG, one per
+        picture and per page). ``ImageRef.uri`` in the JSON points at the
+        bare filename inside ``artifacts/``; this helper inlines those bytes
+        as ``data:<mime>;base64,...`` URIs so the resulting
+        :class:`DoclingDocument` is shape-equivalent to what the local
+        converter produces.
+        """
+        import base64
+        import io
+        import zipfile
+
+        from docling_core.types.doc.document import DoclingDocument
+
+        with zipfile.ZipFile(io.BytesIO(zip_bytes)) as zf:
+            top_level_jsons = [
+                n for n in zf.namelist() if n.endswith(".json") and "/" not in n
+            ]
+            if not top_level_jsons:
+                raise ValueError(
+                    f"docling-serve zip for {name} has no top-level JSON document"
+                )
+            with zf.open(top_level_jsons[0]) as f:
+                doc_json = json.loads(f.read().decode("utf-8"))
+
+            artifact_cache: dict[str, bytes] = {}
+
+            def _inline(image_field: dict | None) -> None:
+                if not image_field:
+                    return
+                uri = image_field.get("uri")
+                # In target_type=zip mode docling-serve writes the URI as the
+                # in-archive path, e.g. "artifacts/image_000000_<sha>.png".
+                # Anything else (data: URI, missing field, weird shape) is left
+                # alone.
+                if not isinstance(uri, str) or not uri.startswith("artifacts/"):
+                    return
+                if uri not in artifact_cache:
+                    try:
+                        artifact_cache[uri] = zf.read(uri)
+                    except KeyError:
+                        return
+                mime = image_field.get("mimetype") or "image/png"
+                b64 = base64.b64encode(artifact_cache[uri]).decode("ascii")
+                image_field["uri"] = f"data:{mime};base64,{b64}"
+
+            for picture in doc_json.get("pictures") or []:
+                _inline(picture.get("image"))
+            for page in (doc_json.get("pages") or {}).values():
+                if isinstance(page, dict):
+                    _inline(page.get("image"))
+
+        return DoclingDocument.model_validate(doc_json)
+
     async def _make_request(self, files: dict, name: str) -> "DoclingDocument":
         """Make an async request to docling-serve and poll for results.
 
@@ -141,6 +225,16 @@ class DoclingServeConverter(DocumentConverter):
         from docling_core.types.doc.document import DoclingDocument
 
         data = self._build_conversion_data()
+
+        if self._picture_images_enabled():
+            zip_bytes = await self.client.submit_and_poll_zip(
+                endpoint="/v1/convert/file/async",
+                files=files,
+                data=data,
+                name=name,
+            )
+            return self._parse_zip_to_docling(zip_bytes, name)
+
         result = await self.client.submit_and_poll(
             endpoint="/v1/convert/file/async",
             files=files,

haiku_rag_slim/haiku/rag/providers/docling_serve.py

@@ -24,6 +24,48 @@ class DoclingServeClient:
             headers["X-Api-Key"] = self.api_key
         return headers
 
+    async def _submit_and_wait(
+        self,
+        client: httpx.AsyncClient,
+        endpoint: str,
+        files: dict[str, Any],
+        data: dict[str, Any],
+        headers: dict[str, str],
+        name: str,
+    ) -> str:
+        """Submit a task and poll until success. Returns the task_id.
+
+        Shared by submit_and_poll (JSON results) and submit_and_poll_zip
+        (binary zip results) — only the result-fetching step differs.
+        """
+        submit_url = f"{self.base_url}{endpoint}"
+        response = await client.post(
+            submit_url,
+            files=files,
+            data=data,
+            headers=headers,
+        )
+        response.raise_for_status()
+        submit_result = response.json()
+        task_id = submit_result.get("task_id")
+
+        if not task_id:
+            raise ValueError("docling-serve did not return a task_id")
+
+        poll_url = f"{self.base_url}/v1/status/poll/{task_id}"
+        while True:
+            poll_response = await client.get(poll_url, headers=headers)
+            poll_response.raise_for_status()
+            poll_result = poll_response.json()
+            status = poll_result.get("task_status")
+
+            if status == "success":
+                return task_id
+            elif status in ("failure", "error"):
+                raise ValueError(f"docling-serve task failed for {name}: {poll_result}")
+
+            await asyncio.sleep(1)
+
     async def submit_and_poll(
         self,
         endpoint: str,
@@ -31,7 +73,7 @@ class DoclingServeClient:
         data: dict[str, Any],
         name: str = "document",
     ) -> dict[str, Any]:
-        """Submit a task and poll until completion.
+        """Submit a task and poll until completion; fetch result as JSON.
 
         Args:
             endpoint: The async endpoint path (e.g., "/v1/convert/file/async")
@@ -49,39 +91,9 @@ class DoclingServeClient:
 
         try:
             async with httpx.AsyncClient(timeout=self.timeout) as client:
-                # Submit async task
-                submit_url = f"{self.base_url}{endpoint}"
-                response = await client.post(
-                    submit_url,
-                    files=files,
-                    data=data,
-                    headers=headers,
+                task_id = await self._submit_and_wait(
+                    client, endpoint, files, data, headers, name
                 )
-                response.raise_for_status()
-                submit_result = response.json()
-                task_id = submit_result.get("task_id")
-
-                if not task_id:
-                    raise ValueError("docling-serve did not return a task_id")
-
-                # Poll for completion
-                poll_url = f"{self.base_url}/v1/status/poll/{task_id}"
-                while True:
-                    poll_response = await client.get(poll_url, headers=headers)
-                    poll_response.raise_for_status()
-                    poll_result = poll_response.json()
-                    status = poll_result.get("task_status")
-
-                    if status == "success":
-                        break
-                    elif status in ("failure", "error"):
-                        raise ValueError(
-                            f"docling-serve task failed for {name}: {poll_result}"
-                        )
-
-                    await asyncio.sleep(1)
-
-                # Fetch result
                 result_url = f"{self.base_url}/v1/result/{task_id}"
                 result_response = await client.get(result_url, headers=headers)
                 result_response.raise_for_status()
@@ -106,3 +118,49 @@ class DoclingServeClient:
             raise
         except Exception as e:
             raise ValueError(f"Failed to process via docling-serve: {e}")
+
+    async def submit_and_poll_zip(
+        self,
+        endpoint: str,
+        files: dict[str, Any],
+        data: dict[str, Any],
+        name: str = "document",
+    ) -> bytes:
+        """Submit a task and poll until completion; fetch result as raw bytes.
+
+        Used when the caller requested ``target_type=zip`` (e.g. to retrieve
+        picture image bytes that docling-serve only emits as referenced files
+        bundled into a zip archive). The submit/poll flow is identical to
+        ``submit_and_poll``; only the result-fetching step differs.
+        """
+        headers = self._get_headers()
+
+        try:
+            async with httpx.AsyncClient(timeout=self.timeout) as client:
+                task_id = await self._submit_and_wait(
+                    client, endpoint, files, data, headers, name
+                )
+                result_url = f"{self.base_url}/v1/result/{task_id}"
+                result_response = await client.get(result_url, headers=headers)
+                result_response.raise_for_status()
+                return result_response.content
+
+        except httpx.ConnectError 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 httpx.TimeoutException as e:
+            raise ValueError(
+                f"Request to docling-serve timed out after {self.timeout}s. Error: {e}"
+            )
+        except httpx.HTTPStatusError 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 ValueError:
+            raise
+        except Exception as e:
+            raise ValueError(f"Failed to process via docling-serve: {e}")

tests/test_converters.py

@@ -79,6 +79,46 @@ def create_async_workflow_mocks(
     return submit_response, poll_response, result_response
 
 
+def create_async_workflow_zip_mocks(
+    doc_json: dict,
+    artifacts: dict[str, bytes] | None = None,
+    task_id: str = "test-task-zip",
+) -> tuple[Mock, Mock, Mock]:
+    """Mock responses for the ``target_type=zip`` async workflow.
+
+    Builds a real zip in memory containing the document JSON at the archive
+    root and each ``artifacts[name] = bytes`` entry under ``artifacts/<name>``.
+    The result response exposes the zip via ``.content`` (raw bytes) so the
+    converter's zip-parsing path is exercised end-to-end.
+    """
+    import io
+    import json as _json
+    import zipfile
+
+    submit_response = Mock()
+    submit_response.status_code = 200
+    submit_response.json.return_value = {"task_id": task_id, "task_status": "pending"}
+    submit_response.raise_for_status = Mock()
+
+    poll_response = Mock()
+    poll_response.status_code = 200
+    poll_response.json.return_value = {"task_id": task_id, "task_status": "success"}
+    poll_response.raise_for_status = Mock()
+
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, mode="w") as zf:
+        zf.writestr(f"{doc_json.get('name', 'document')}.json", _json.dumps(doc_json))
+        for filename, blob in (artifacts or {}).items():
+            zf.writestr(f"artifacts/{filename}", blob)
+
+    result_response = Mock()
+    result_response.status_code = 200
+    result_response.content = buf.getvalue()
+    result_response.raise_for_status = Mock()
+
+    return submit_response, poll_response, result_response
+
+
 class TestTextFileHandler:
     """Tests for TextFileHandler utility class."""
 
@@ -712,6 +752,122 @@ class TestDoclingServeConverter:
             data = call_kwargs["data"]
             assert data["ocr_engine"] == "rapidocr"
 
+    @pytest.mark.asyncio
+    async def test_picture_images_request_uses_referenced_zip(self, config):
+        """When generate_picture_images is on, the request flips to
+        image_export_mode=referenced + target_type=zip and consumes a zip
+        response — mirrors the upstream docling-serve#576 workaround.
+        """
+        config.processing.conversion_options.generate_picture_images = True
+        converter = DoclingServeConverter(config)
+
+        doc_json = create_mock_docling_document("test")
+        submit_resp, poll_resp, result_resp = create_async_workflow_zip_mocks(doc_json)
+
+        with patch("httpx.AsyncClient") as mock_client_class:
+            mock_client = AsyncMock()
+            mock_client.post = AsyncMock(return_value=submit_resp)
+            mock_client.get = AsyncMock(side_effect=[poll_resp, result_resp])
+            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+            mock_client.__aexit__ = AsyncMock(return_value=None)
+            mock_client_class.return_value = mock_client
+
+            await converter.convert_text("# Test")
+
+            data = mock_client.post.call_args.kwargs["data"]
+            assert data["image_export_mode"] == "referenced"
+            assert data["target_type"] == "zip"
+            assert data["include_images"] == "true"
+
+    def test_parse_zip_rehydrates_picture_uri(self, config):
+        """Zip path inlines artifact bytes as data: URIs on PictureItem.image.
+
+        target_type=zip mode emits the URI as ``artifacts/<filename>`` — the
+        same string we use to read the entry out of the archive.
+        """
+        import base64
+        import io
+        import json as _json
+        import zipfile
+
+        converter = DoclingServeConverter(config)
+
+        doc_json = create_mock_docling_document("test")
+        doc_json["pictures"] = [
+            {
+                "self_ref": "#/pictures/0",
+                "parent": {"cref": "#/body"},
+                "children": [],
+                "content_layer": "body",
+                "label": "picture",
+                "prov": [],
+                "captions": [],
+                "references": [],
+                "footnotes": [],
+                "annotations": [],
+                "image": {
+                    "mimetype": "image/png",
+                    "dpi": 144,
+                    "size": {"width": 1.0, "height": 1.0},
+                    "uri": "artifacts/image_000000_test.png",
+                },
+            }
+        ]
+        fake_png = b"\x89PNG\r\n\x1a\nfake-bytes-for-test"
+
+        buf = io.BytesIO()
+        with zipfile.ZipFile(buf, mode="w") as zf:
+            zf.writestr("test.json", _json.dumps(doc_json))
+            zf.writestr("artifacts/image_000000_test.png", fake_png)
+
+        doc = converter._parse_zip_to_docling(buf.getvalue(), "test")
+
+        assert len(doc.pictures) == 1
+        image = doc.pictures[0].image
+        assert image is not None
+        uri = str(image.uri)
+        assert uri.startswith("data:image/png;base64,")
+        decoded = base64.b64decode(uri.split(",", 1)[1])
+        assert decoded == fake_png
+
+    def test_parse_zip_leaves_unknown_artifact_uri_unchanged(self, config):
+        """If a PictureItem references an artifact that is not in the zip,
+        the URI is left as-is (no crash, no truncation)."""
+        import io
+        import json as _json
+        import zipfile
+
+        converter = DoclingServeConverter(config)
+        doc_json = create_mock_docling_document("test")
+        doc_json["pictures"] = [
+            {
+                "self_ref": "#/pictures/0",
+                "parent": {"cref": "#/body"},
+                "children": [],
+                "content_layer": "body",
+                "label": "picture",
+                "prov": [],
+                "captions": [],
+                "references": [],
+                "footnotes": [],
+                "annotations": [],
+                "image": {
+                    "mimetype": "image/png",
+                    "dpi": 144,
+                    "size": {"width": 1.0, "height": 1.0},
+                    "uri": "artifacts/missing.png",
+                },
+            }
+        ]
+
+        buf = io.BytesIO()
+        with zipfile.ZipFile(buf, mode="w") as zf:
+            zf.writestr("test.json", _json.dumps(doc_json))
+
+        doc = converter._parse_zip_to_docling(buf.getvalue(), "test")
+        assert doc.pictures[0].image is not None
+        assert str(doc.pictures[0].image.uri) == "artifacts/missing.png"
+
     @pytest.mark.asyncio
     async def test_convert_text_connection_error(self, converter):
         """Test handling of connection errors."""
@@ -880,7 +1036,12 @@ class TestDoclingServeConverterPictureDescription:
 
     @pytest.mark.asyncio
     async def test_picture_description_options_passed_to_api(self, config):
-        """Test that picture description options are passed to docling-serve API."""
+        """Test that picture description options are passed to docling-serve API.
+
+        Picture descriptions force ``generate_picture_images=True`` upstream,
+        which routes the request through the ``target_type=zip`` path so the
+        VLM can see the actual figures. The test mocks the zip workflow.
+        """
         import json
 
         config.processing.conversion_options.picture_description.enabled = True
@@ -896,7 +1057,7 @@ class TestDoclingServeConverterPictureDescription:
         converter = DoclingServeConverter(config)
 
         doc_json = create_mock_docling_document("test")
-        submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json)
+        submit_resp, poll_resp, result_resp = create_async_workflow_zip_mocks(doc_json)
 
         with patch("httpx.AsyncClient") as mock_client_class:
             mock_client = AsyncMock()
@@ -914,6 +1075,8 @@ class TestDoclingServeConverterPictureDescription:
 
             assert data["do_picture_description"] == "true"
             assert data["include_images"] == "true"
+            assert data["image_export_mode"] == "referenced"
+            assert data["target_type"] == "zip"
             assert "picture_description_api" in data
 
             api_config = json.loads(data["picture_description_api"])
@@ -1090,15 +1253,18 @@ class TestDoclingServeConverterIntegration:
         assert len(doc.pages) > 0
         assert len(doc.export_to_markdown().strip()) > 100
 
-    @pytest.mark.xfail(
-        reason="docling-serve does not return picture image data in JSON response "
-        "even with include_images=true. Page images work, but extracted picture/figure "
-        "images are not included. This is a docling-serve limitation."
-    )
     @pytest.mark.vcr()
     @pytest.mark.asyncio
     async def test_convert_pdf_with_picture_images(self, config):
-        """Test PDF conversion includes picture images when enabled."""
+        """Test PDF conversion includes picture images when enabled.
+
+        docling-serve only emits picture image bytes when
+        ``image_export_mode="referenced"`` (upstream issue
+        docling-project/docling-serve#576). The converter switches to
+        ``referenced`` + ``target_type="zip"`` when picture images are
+        wanted, then rehydrates the bundled artifact files into ``data:``
+        URIs so the result is shape-equivalent to the local converter.
+        """
         pdf_path = Path("tests/data/doclaynet.pdf")
         config.processing.conversion_options.generate_picture_images = True
         converter = DoclingServeConverter(config)
@@ -1106,9 +1272,13 @@ class TestDoclingServeConverterIntegration:
         doc = await converter.convert_file(pdf_path)
         assert isinstance(doc, DoclingDocument)
 
-        # Check that at least some pictures have image data
         pictures_with_images = [p for p in doc.pictures if p.image is not None]
-        if doc.pictures:
-            assert len(pictures_with_images) > 0, (
-                "Pictures should have image data when generate_picture_images=True"
-            )
+        assert doc.pictures, "doclaynet.pdf is expected to contain at least one picture"
+        assert len(pictures_with_images) > 0, (
+            "Pictures should have image data when generate_picture_images=True"
+        )
+        sample = pictures_with_images[0]
+        assert sample.image is not None
+        assert str(sample.image.uri).startswith("data:image/"), (
+            "Rehydrated picture URI should be a data: URI, not a bare artifact filename"
+        )