loader.py

#
# Copyright IBM Corp. 2025 - 2025
# SPDX-License-Identifier: MIT
#

"""Docling LangChain loader module."""

from abc import ABC, abstractmethod
from collections.abc import Callable, Iterable, Iterator
from enum import Enum
from typing import Any, Protocol

from docling.chunking import BaseChunk, BaseChunker, HybridChunker
from docling.datamodel.document import ConversionResult, DoclingDocument
from docling.document_converter import DocumentConverter
from langchain_core.document_loaders import BaseLoader
from langchain_core.documents import Document


class ExportType(str, Enum):
    """Enumeration of available export types."""

    MARKDOWN = "markdown"
    DOC_CHUNKS = "doc_chunks"


class ConversionBackend(Protocol):
    """Structural interface for local and service-backed Docling conversion."""

    @property
    def convert(self) -> Callable[..., ConversionResult]:
        """Return the backend-specific conversion callable."""
        ...


class BaseMetaExtractor(ABC):
    """BaseMetaExtractor."""

    @abstractmethod
    def extract_chunk_meta(self, file_path: str, chunk: BaseChunk) -> dict[str, Any]:
        """Extract chunk meta."""
        raise NotImplementedError()

    @abstractmethod
    def extract_dl_doc_meta(
        self, file_path: str, dl_doc: DoclingDocument
    ) -> dict[str, Any]:
        """Extract Docling document meta."""
        raise NotImplementedError()


class MetaExtractor(BaseMetaExtractor):
    """MetaExtractor."""

    def extract_chunk_meta(self, file_path: str, chunk: BaseChunk) -> dict[str, Any]:
        """Extract chunk meta."""
        return {
            "source": file_path,
            "dl_meta": chunk.meta.export_json_dict(),
        }

    def extract_dl_doc_meta(
        self, file_path: str, dl_doc: DoclingDocument
    ) -> dict[str, Any]:
        """Extract Docling document meta."""
        return {"source": file_path}


class DoclingLoader(BaseLoader):
    """Docling Loader."""

    def __init__(
        self,
        file_path: str | Iterable[str],
        *,
        converter: ConversionBackend | None = None,
        convert_kwargs: dict[str, Any] | None = None,
        export_type: ExportType = ExportType.DOC_CHUNKS,
        md_export_kwargs: dict[str, Any] | None = None,
        chunker: BaseChunker | None = None,
        meta_extractor: BaseMetaExtractor | None = None,
    ):
        """Initialize with a file path.

        Args:
            file_path: File source as single str (URL or local file) or Iterable
                thereof.
            converter: A local `DocumentConverter`, remote `DoclingServiceClient`, or
                compatible conversion backend. Defaults to `None` (i.e. a local
                `DocumentConverter` defined internally).
            convert_kwargs: Any backend-specific kwargs to pass to the conversion
                invocation. For example, pass `options` when using a
                `DoclingServiceClient`. Defaults to `None` (i.e. behavior defined
                internally).
            export_type: The type to export to: either `ExportType.MARKDOWN` (outputs
                Markdown of whole input file) or `ExportType.DOC_CHUNKS` (outputs chunks
                based on chunker).
            md_export_kwargs: Any specific kwargs to pass to Markdown export (in case of
                `ExportType.MARKDOWN`). Defaults to `None` (i.e. behavior defined
                internally).
            chunker: Any specific `BaseChunker` to use (in case of
                `ExportType.DOC_CHUNKS`). Defaults to `None` (i.e. chunker defined
                internally).
            meta_extractor: The extractor instance to use for populating the output
                document metadata; if not set, a system default is used.
        """
        self._file_paths = (
            file_path
            if isinstance(file_path, Iterable) and not isinstance(file_path, str)
            else [file_path]
        )

        self._converter: ConversionBackend = converter or DocumentConverter()
        self._convert_kwargs = convert_kwargs if convert_kwargs is not None else {}
        self._export_type = export_type
        self._md_export_kwargs = (
            md_export_kwargs
            if md_export_kwargs is not None
            else {"image_placeholder": ""}
        )
        if self._export_type == ExportType.DOC_CHUNKS:
            self._chunker: BaseChunker = chunker or HybridChunker()
        self._meta_extractor = meta_extractor or MetaExtractor()

    def lazy_load(
        self,
    ) -> Iterator[Document]:
        """Lazy load documents."""
        for file_path in self._file_paths:
            conv_res = self._converter.convert(
                source=file_path,
                **self._convert_kwargs,
            )
            dl_doc = conv_res.document
            if self._export_type == ExportType.MARKDOWN:
                yield Document(
                    page_content=dl_doc.export_to_markdown(**self._md_export_kwargs),
                    metadata=self._meta_extractor.extract_dl_doc_meta(
                        file_path=file_path,
                        dl_doc=dl_doc,
                    ),
                )
            elif self._export_type == ExportType.DOC_CHUNKS:
                chunk_iter = self._chunker.chunk(dl_doc)
                for chunk in chunk_iter:
                    yield Document(
                        page_content=self._chunker.contextualize(chunk=chunk),
                        metadata=self._meta_extractor.extract_chunk_meta(
                            file_path=file_path,
                            chunk=chunk,
                        ),
                    )

            else:
                raise ValueError(f"Unexpected export type: {self._export_type}")