From eb2b65135af9ef6596855489bebd59d0d32fabfd Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 4 Nov 2025 15:31:14 +0200 Subject: [PATCH] Lazy load docling & friends --- haiku_rag_slim/haiku/rag/chunker.py | 27 ++++--------- src/haiku/rag/converters/base.py | 8 ++-- src/haiku/rag/converters/docling_local.py | 48 ++++++----------------- 3 files changed, 24 insertions(+), 59 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/chunker.py b/haiku_rag_slim/haiku/rag/chunker.py index 03ca0e19..e7941fa1 100644 --- a/haiku_rag_slim/haiku/rag/chunker.py +++ b/haiku_rag_slim/haiku/rag/chunker.py @@ -1,18 +1,11 @@ -from typing import ClassVar +from typing import TYPE_CHECKING, ClassVar import tiktoken -from docling_core.transforms.chunker.tokenizer.openai import OpenAITokenizer -from docling_core.types.doc.document import DoclingDocument from haiku.rag.config import Config -# Check if docling is available -try: - import docling # noqa: F401 - - DOCLING_AVAILABLE = True -except ImportError: - DOCLING_AVAILABLE = False +if TYPE_CHECKING: + from docling_core.types.doc.document import DoclingDocument class Chunker: @@ -31,21 +24,17 @@ class Chunker: self, chunk_size: int = Config.processing.chunk_size, ): - if not DOCLING_AVAILABLE: - raise ImportError( - "Docling is required for chunking. " - "Install with: pip install haiku.rag-slim[docling]" - ) - from docling.chunking import HybridChunker # type: ignore + from docling_core.transforms.chunker.hybrid_chunker import HybridChunker + from docling_core.transforms.chunker.tokenizer.openai import OpenAITokenizer self.chunk_size = chunk_size tokenizer = OpenAITokenizer( tokenizer=tiktoken.encoding_for_model("gpt-4o"), max_tokens=chunk_size ) - self.chunker = HybridChunker(tokenizer=tokenizer) # type: ignore + self.chunker = HybridChunker(tokenizer=tokenizer) - async def chunk(self, document: DoclingDocument) -> list[str]: + async def chunk(self, document: "DoclingDocument") -> list[str]: """Split the document into chunks using docling's structure-aware chunking. Args: @@ -62,4 +51,4 @@ class Chunker: return [self.chunker.contextualize(chunk) for chunk in chunks] -chunker = Chunker() +chunker = Chunker() \ No newline at end of file diff --git a/src/haiku/rag/converters/base.py b/src/haiku/rag/converters/base.py index ed1826cb..10baa87e 100644 --- a/src/haiku/rag/converters/base.py +++ b/src/haiku/rag/converters/base.py @@ -2,8 +2,10 @@ from abc import ABC, abstractmethod from pathlib import Path +from typing import TYPE_CHECKING -from docling_core.types.doc.document import DoclingDocument +if TYPE_CHECKING: + from docling_core.types.doc.document import DoclingDocument class DocumentConverter(ABC): @@ -24,7 +26,7 @@ class DocumentConverter(ABC): pass @abstractmethod - def convert_file(self, path: Path) -> DoclingDocument: + def convert_file(self, path: Path) -> "DoclingDocument": """Convert a file to DoclingDocument format. Args: @@ -39,7 +41,7 @@ class DocumentConverter(ABC): pass @abstractmethod - def convert_text(self, text: str, name: str = "content.md") -> DoclingDocument: + def convert_text(self, text: str, name: str = "content.md") -> "DoclingDocument": """Convert text content to DoclingDocument format. Args: diff --git a/src/haiku/rag/converters/docling_local.py b/src/haiku/rag/converters/docling_local.py index 97a094ed..d52f697d 100644 --- a/src/haiku/rag/converters/docling_local.py +++ b/src/haiku/rag/converters/docling_local.py @@ -2,26 +2,12 @@ from io import BytesIO from pathlib import Path -from typing import ClassVar - -<<<<<<<< HEAD:haiku_rag_slim/haiku/rag/reader.py -from docling_core.types.doc.document import DoclingDocument -from haiku.rag.utils import text_to_docling_document -======== -from docling.document_converter import DocumentConverter as DoclingDocConverter -from docling_core.types.doc.document import DoclingDocument -from docling_core.types.io import DocumentStream +from typing import TYPE_CHECKING, ClassVar from haiku.rag.converters.base import DocumentConverter ->>>>>>>> f7975a3d5946 (Introduce converters for supporting more than local docling document conversion. Transform existing FileReader and utils to "docling-local" converter):src/haiku/rag/converters/docling_local.py -# Check if docling is available -try: - import docling # noqa: F401 - - DOCLING_AVAILABLE = True -except ImportError: - DOCLING_AVAILABLE = False +if TYPE_CHECKING: + from docling_core.types.doc.document import DoclingDocument class DoclingLocalConverter(DocumentConverter): @@ -114,7 +100,7 @@ class DoclingLocalConverter(DocumentConverter): """Return list of file extensions supported by this converter.""" return self.docling_extensions + self.text_extensions - def convert_file(self, path: Path) -> DoclingDocument: + def convert_file(self, path: Path) -> "DoclingDocument": """Convert a file to DoclingDocument using local docling. Args: @@ -126,23 +112,14 @@ class DoclingLocalConverter(DocumentConverter): Raises: ValueError: If the file cannot be converted. """ + from docling.document_converter import DocumentConverter as DoclingDocConverter + try: file_extension = path.suffix.lower() if file_extension in self.docling_extensions: # Use docling for complex document formats -<<<<<<<< HEAD:haiku_rag_slim/haiku/rag/reader.py - if not DOCLING_AVAILABLE: - raise ImportError( - "Docling is required for processing this file type. " - "Install with: pip install haiku.rag-slim[docling]" - ) - from docling.document_converter import DocumentConverter - - converter = DocumentConverter() -======== converter = DoclingDocConverter() ->>>>>>>> f7975a3d5946 (Introduce converters for supporting more than local docling document conversion. Transform existing FileReader and utils to "docling-local" converter):src/haiku/rag/converters/docling_local.py result = converter.convert(path) return result.document elif file_extension in self.text_extensions: @@ -159,17 +136,11 @@ class DoclingLocalConverter(DocumentConverter): else: # Fallback: try to read as text and convert to DoclingDocument content = path.read_text(encoding="utf-8") -<<<<<<<< HEAD:haiku_rag_slim/haiku/rag/reader.py - return text_to_docling_document(content, name=f"{path.stem}.md") - except ImportError: - raise -======== return self.convert_text(content, name=f"{path.stem}.md") ->>>>>>>> f7975a3d5946 (Introduce converters for supporting more than local docling document conversion. Transform existing FileReader and utils to "docling-local" converter):src/haiku/rag/converters/docling_local.py except Exception: raise ValueError(f"Failed to parse file: {path}") - def convert_text(self, text: str, name: str = "content.md") -> DoclingDocument: + def convert_text(self, text: str, name: str = "content.md") -> "DoclingDocument": """Convert text content to DoclingDocument using local docling. Args: @@ -182,6 +153,9 @@ class DoclingLocalConverter(DocumentConverter): Raises: ValueError: If the text cannot be converted. """ + from docling.document_converter import DocumentConverter as DoclingDocConverter + from docling_core.types.io import DocumentStream + try: bytes_io = BytesIO(text.encode("utf-8")) doc_stream = DocumentStream(name=name, stream=bytes_io) @@ -189,4 +163,4 @@ class DoclingLocalConverter(DocumentConverter): result = converter.convert(doc_stream) return result.document except Exception as e: - raise ValueError(f"Failed to convert text to DoclingDocument: {e}") + raise ValueError(f"Failed to convert text to DoclingDocument: {e}") \ No newline at end of file