Several cleanups, addressing review comments

Signed-off-by: Christoph Auer <cau@zurich.ibm.com>

Author: cau-git

docling_jobkit/config/target_config.py

@@ -4,7 +4,13 @@
 
 
 class S3PresignedConfig(BaseModel):
+    """Server-managed presigned target policy injected by docling-serve.
+
+    ``docling-serve`` builds this internal config from its artifact storage settings
+    before constructing an orchestrator. The storage prefix lives on
+    ``s3_coords.key_prefix`` so jobkit does not carry a redundant top-level field.
+    """
+
     s3_coords: S3Coordinates
-    key_prefix: str = "converted/"
     date_partition_format: str = "%Y%m%d"
     url_expiration: int = Field(default=3600, ge=60, le=604800)

docling_jobkit/connectors/google_drive_source_processor.py

@@ -3,6 +3,8 @@
 from io import BytesIO
 from typing import TYPE_CHECKING, Iterator
 
+from typing_extensions import override
+
 from docling.datamodel.base_models import DocumentStream
 
 if TYPE_CHECKING:
@@ -84,6 +86,7 @@ def _fetch_document_by_id(self, info: GoogleDriveFileIdentifier) -> DocumentStre
             stream=buffer,
         )
 
+    @override
     def _make_document_ref(
         self, info: GoogleDriveFileIdentifier, source_index: int
     ) -> SourceDocumentRef[GoogleDriveFileIdentifier]:
@@ -93,5 +96,4 @@ def _make_document_ref(
             source_index=source_index,
             source_uri=source_uri,
             filename=info["name"],
-            metadata={"mimeType": info["mimeType"]},
         )

docling_jobkit/connectors/http_source_processor.py

@@ -1,5 +1,7 @@
 from typing import Iterator, TypedDict
 
+from typing_extensions import override
+
 from docling.datamodel.service.sources import FileSource, HttpSource
 from docling_core.types.io import DocumentStream
 
@@ -43,6 +45,7 @@ def _fetch_document_by_id(self, identifier: HttpFileIdentifier) -> DocumentStrea
         else:
             raise ValueError(f"Unsupported source type: {type(source)}")
 
+    @override
     def _make_document_ref(
         self, identifier: HttpFileIdentifier, source_index: int
     ) -> SourceDocumentRef[HttpFileIdentifier]:
@@ -60,14 +63,16 @@ def _make_document_ref(
             filename=filename,
         )
 
+    @override
     def fetch_converter_source_by_ref(
         self, ref: SourceDocumentRef[HttpFileIdentifier]
     ) -> ConverterSource:
         source = ref.id["source"]
         if isinstance(source, HttpSource):
             return str(source.url)
-        return self.fetch_document_by_ref(ref)
+        return self._fetch_document_by_id(ref.id)
 
+    @override
     def headers_for_ref(
         self, ref: SourceDocumentRef[HttpFileIdentifier]
     ) -> dict[str, object] | None:

docling_jobkit/connectors/local_path_source_processor.py

@@ -1,6 +1,8 @@
 from pathlib import Path
 from typing import Iterator, TypedDict
 
+from typing_extensions import override
+
 from docling_core.types.io import DocumentStream
 
 from docling_jobkit.connectors.source_processor import (
@@ -125,6 +127,7 @@ def _fetch_document_by_id(
 
         return DocumentStream(name=str(file_path), stream=buffer)
 
+    @override
     def _make_document_ref(
         self, identifier: LocalPathFileIdentifier, source_index: int
     ) -> SourceDocumentRef[LocalPathFileIdentifier]:
@@ -134,10 +137,6 @@ def _make_document_ref(
             source_index=source_index,
             source_uri=str(file_path),
             filename=str(file_path),
-            metadata={
-                "size": identifier["size"],
-                "last_modified": identifier["last_modified"],
-            },
         )
 
     def _fetch_documents(self) -> Iterator[DocumentStream]:

docling_jobkit/connectors/s3_source_processor.py

@@ -1,6 +1,8 @@
 from io import BytesIO
 from typing import Iterator, TypedDict
 
+from typing_extensions import override
+
 from docling.datamodel.service.sources import S3Coordinates
 from docling_core.types.io import DocumentStream
 
@@ -52,6 +54,7 @@ def _count_documents(self) -> int:
             total += len(page.get("Contents", []))
         return total
 
+    @override
     def _make_document_ref(
         self, identifier: S3FileIdentifier, source_index: int
     ) -> SourceDocumentRef[S3FileIdentifier]:
@@ -61,10 +64,6 @@ def _make_document_ref(
             source_index=source_index,
             source_uri=f"s3://{self._coords.bucket}/{key}",
             filename=key,
-            metadata={
-                "size": identifier["size"],
-                "last_modified": identifier["last_modified"],
-            },
         )
 
     # ----------------- Document fetch -----------------

docling_jobkit/connectors/s3_upload_support.py

@@ -66,12 +66,14 @@ def build_task_scoped_s3_key(
     artifact_filename: str,
 ) -> str:
     # PresignedUrlTarget writes into operator-managed storage, so the full key
-    # includes the managed prefix/tenant/date/task structure before the per-source hash.
+    # includes the managed prefix/tenant/date/task structure before the per-source
+    # hash. The storage prefix is injected by docling-serve onto
+    # config.s3_coords.key_prefix when the orchestrator is constructed.
     source_key = hash_path_component(source_uri)
     date_partition = datetime.now(timezone.utc).strftime(config.date_partition_format)
 
     path_parts: list[str] = []
-    key_prefix = config.key_prefix.strip("/")
+    key_prefix = config.s3_coords.key_prefix.strip("/")
     if key_prefix:
         path_parts.append(key_prefix)
 

docling_jobkit/connectors/source_processor.py

@@ -1,9 +1,9 @@
 from abc import ABC, abstractmethod
 from contextlib import AbstractContextManager
 from itertools import islice
-from typing import Any, Callable, Generic, Iterator, Sequence, TypeVar
+from typing import Callable, Generic, Iterator, Sequence, TypeVar
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict
 
 from docling.datamodel.base_models import DocumentStream
 
@@ -21,11 +21,16 @@ class SourceDocumentRef(BaseModel, Generic[FileIdentifierT]):
     source_index: int
     source_uri: str
     filename: str
-    metadata: dict[str, Any] = Field(default_factory=dict)
 
 
 class DocumentChunk(BaseModel, Generic[SourceT, FileIdentifierT]):
-    """A data-only source chunk plus a local fetcher convenience."""
+    """A serializable source chunk plus an optional local fetcher convenience.
+
+    Local/CLI callers may attach a fetcher so ``iter_documents()`` can materialize
+    streams lazily from the refs. Cross-process callers such as Ray must strip that
+    fetcher because it may capture initialized connector state that is not safe to
+    serialize.
+    """
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
@@ -53,6 +58,7 @@ def index(self) -> int:
         return self.chunk_index
 
     def iter_documents(self) -> Iterator[DocumentStream]:
+        """Materialize documents for local callers when a fetcher is attached."""
         if self._fetcher is None:
             raise RuntimeError("DocumentChunk does not have an attached fetcher.")
         for ref in self.refs:
@@ -118,19 +124,21 @@ def source(self) -> SourceT:
     def _count_documents(self) -> int | None:
         return None
 
-    def fetch_document_by_ref(
-        self, ref: SourceDocumentRef[FileIdentifierT]
-    ) -> DocumentStream:
-        return self._fetch_document_by_id(ref.id)
-
     def fetch_converter_source_by_ref(
         self, ref: SourceDocumentRef[FileIdentifierT]
     ) -> ConverterSource:
-        return self.fetch_document_by_ref(ref)
+        """Resolve a ref into the converter input expected by the backend.
+
+        Most connectors materialize a ``DocumentStream`` from the ref's identifier.
+        Connectors with remote-fetch semantics may override this to return a lighter
+        representation such as a source URL.
+        """
+        return self._fetch_document_by_id(ref.id)
 
     def headers_for_ref(
         self, ref: SourceDocumentRef[FileIdentifierT]
-    ) -> dict[str, Any] | None:
+    ) -> dict[str, object] | None:
+        """Return per-ref request headers when the converter should fetch remotely."""
         del ref
         return None