fix(file_processors): Preserve Docling chunk structure metadata

Signed-off-by: Peter Gustafsson <peter.gustafsson6@gmail.com>

Author: Peter Gustafsson

src/ogx/providers/inline/file_processor/docling/_metadata.py

@@ -0,0 +1,29 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Any
+
+
+def extract_structural_metadata(doc_chunk: Any) -> dict[str, Any]:
+    chunk_meta = getattr(doc_chunk, "meta", None)
+    if chunk_meta is None:
+        return {}
+
+    metadata: dict[str, Any] = {}
+    headings = getattr(chunk_meta, "headings", None)
+    if headings:
+        metadata["headings"] = headings
+
+    page_numbers = {
+        page_number
+        for doc_item in getattr(chunk_meta, "doc_items", [])
+        for provenance in (getattr(doc_item, "prov", None) or [])
+        if (page_number := getattr(provenance, "page_no", None)) is not None
+    }
+    if page_numbers:
+        metadata["page_numbers"] = sorted(page_numbers)
+
+    return metadata

src/ogx/providers/inline/file_processor/docling/docling.py

@@ -34,6 +34,7 @@
     VectorStoreChunkingStrategy,
 )
 
+from ._metadata import extract_structural_metadata
 from .config import DoclingFileProcessorConfig
 
 log = get_logger(name=__name__, category="providers::file_processors")
@@ -277,7 +278,6 @@ def _create_chunks(
             if not text or not text.strip():
                 continue
 
-            headings = getattr(doc_chunk, "headings", None)
             chunk_window = f"{i}"
 
             chunk_id = generate_chunk_id(document_id, text, chunk_window)
@@ -286,8 +286,7 @@ def _create_chunks(
                 "document_id": document_id,
                 **document_metadata,
             }
-            if headings:
-                meta["headings"] = headings
+            meta.update(extract_structural_metadata(doc_chunk))
 
             chunks.append(
                 Chunk(

src/ogx/providers/remote/file_processor/docling_serve/docling_serve.py

@@ -343,10 +343,15 @@ async def _convert_and_chunk(
                 **document_metadata,
             }
 
-            headings = raw_chunk.get("meta", {}).get("headings", None)
+            legacy_meta = raw_chunk.get("meta") or {}
+            headings = raw_chunk.get("headings") or legacy_meta.get("headings")
             if headings:
                 meta["headings"] = headings
 
+            page_numbers = raw_chunk.get("page_numbers") or legacy_meta.get("page_numbers")
+            if page_numbers:
+                meta["page_numbers"] = page_numbers
+
             chunks.append(
                 Chunk(
                     content=text,
@@ -426,14 +431,15 @@ async def _convert_and_chunk_async(
                 **document_metadata,
             }
 
-            # Extract headings from meta object
-            headings = None
-            if hasattr(raw_chunk, "meta") and hasattr(raw_chunk.meta, "headings"):
-                headings = raw_chunk.meta.headings
-
+            legacy_meta = getattr(raw_chunk, "meta", None)
+            headings = getattr(raw_chunk, "headings", None) or getattr(legacy_meta, "headings", None)
             if headings:
                 meta["headings"] = headings
 
+            page_numbers = getattr(raw_chunk, "page_numbers", None) or getattr(legacy_meta, "page_numbers", None)
+            if page_numbers:
+                meta["page_numbers"] = page_numbers
+
             chunks.append(
                 Chunk(
                     content=text,

tests/unit/providers/file_processor/test_docling_metadata.py

@@ -0,0 +1,34 @@
+# Copyright (c) The OGX Contributors.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from types import SimpleNamespace
+
+from ogx.providers.inline.file_processor.docling._metadata import extract_structural_metadata
+
+
+def test_extract_structural_metadata_from_native_docling_chunk():
+    doc_chunk = SimpleNamespace(
+        headings=["wrong location"],
+        meta=SimpleNamespace(
+            headings=["Introduction", "Architecture"],
+            doc_items=[
+                SimpleNamespace(prov=[SimpleNamespace(page_no=2), SimpleNamespace(page_no=1)]),
+                SimpleNamespace(prov=[SimpleNamespace(page_no=2)]),
+                SimpleNamespace(prov=None),
+            ],
+        ),
+    )
+
+    assert extract_structural_metadata(doc_chunk) == {
+        "headings": ["Introduction", "Architecture"],
+        "page_numbers": [1, 2],
+    }
+
+
+def test_extract_structural_metadata_omits_empty_values():
+    doc_chunk = SimpleNamespace(meta=SimpleNamespace(headings=None, doc_items=[]))
+
+    assert extract_structural_metadata(doc_chunk) == {}

tests/unit/providers/file_processor/test_docling_serve.py

@@ -41,9 +41,30 @@ def _make_httpx_response(json_body: dict, status_code: int = 200) -> httpx.Respo
 
 CHUNK_RESPONSE = {
     "chunks": [
-        {"text": "First chunk of text.", "meta": {"headings": ["Introduction"]}},
-        {"text": "Second chunk of text.", "meta": {}},
-        {"text": "Third chunk of text.", "meta": {"headings": ["Conclusion"]}},
+        {
+            "filename": "test.pdf",
+            "chunk_index": 0,
+            "text": "First chunk of text.",
+            "headings": ["Introduction"],
+            "doc_items": ["#/texts/0"],
+            "page_numbers": [1],
+        },
+        {
+            "filename": "test.pdf",
+            "chunk_index": 1,
+            "text": "Second chunk of text.",
+            "headings": None,
+            "doc_items": ["#/texts/1"],
+            "page_numbers": None,
+        },
+        {
+            "filename": "test.pdf",
+            "chunk_index": 2,
+            "text": "Third chunk of text.",
+            "headings": ["Conclusion"],
+            "doc_items": ["#/texts/2"],
+            "page_numbers": [2, 3],
+        },
     ],
 }
 
@@ -205,15 +226,18 @@ async def test_chunk_id_uniqueness(self, processor: DoclingServeFileProcessor, u
         ids = [c.chunk_id for c in response.chunks]
         assert len(ids) == len(set(ids))
 
-    async def test_headings_propagated(self, processor: DoclingServeFileProcessor, upload_file: UploadFile):
+    async def test_structural_metadata_propagated(self, processor: DoclingServeFileProcessor, upload_file: UploadFile):
         request = ProcessFileRequest(chunking_strategy=VectorStoreChunkingStrategyAuto())
 
         with patch("httpx.AsyncClient.post", return_value=_make_httpx_response(CHUNK_RESPONSE)):
             response = await processor.process_file(request, file=upload_file)
 
         assert response.chunks[0].metadata["headings"] == ["Introduction"]
+        assert response.chunks[0].metadata["page_numbers"] == [1]
         assert "headings" not in response.chunks[1].metadata
+        assert "page_numbers" not in response.chunks[1].metadata
         assert response.chunks[2].metadata["headings"] == ["Conclusion"]
+        assert response.chunks[2].metadata["page_numbers"] == [2, 3]
 
     async def test_chunk_window_set(self, processor: DoclingServeFileProcessor, upload_file: UploadFile):
         request = ProcessFileRequest(chunking_strategy=VectorStoreChunkingStrategyAuto())
@@ -473,7 +497,14 @@ async def test_local_docker_allows_chunking(self, upload_file: UploadFile):
 
             # Mock submit_chunk() returning a successful job
             mock_job = AsyncMock()
-            mock_chunk = SimpleNamespace(text="Chunk content", meta=SimpleNamespace(headings=None))
+            mock_chunk = SimpleNamespace(
+                filename="test.pdf",
+                chunk_index=0,
+                text="Chunk content",
+                headings=["Introduction"],
+                doc_items=["#/texts/0"],
+                page_numbers=[1, 2],
+            )
             mock_response = SimpleNamespace(chunks=[mock_chunk])
             mock_job.result.return_value = mock_response
             mock_instance.submit_chunk.return_value = mock_job
@@ -483,4 +514,6 @@ async def test_local_docker_allows_chunking(self, upload_file: UploadFile):
 
         assert result.chunks is not None
         assert len(result.chunks) > 0
+        assert result.chunks[0].metadata["headings"] == ["Introduction"]
+        assert result.chunks[0].metadata["page_numbers"] == [1, 2]
         assert result.metadata["conversion_method"] == "async"