Lazy load docling & friends

This commit is contained in:
Yiorgis Gozadinos 2025-11-04 15:31:14 +02:00
parent 7ba1c2376e
commit eb2b65135a
No known key found for this signature in database
3 changed files with 24 additions and 59 deletions

View file

@ -1,18 +1,11 @@
from typing import ClassVar from typing import TYPE_CHECKING, ClassVar
import tiktoken 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 from haiku.rag.config import Config
# Check if docling is available if TYPE_CHECKING:
try: from docling_core.types.doc.document import DoclingDocument
import docling # noqa: F401
DOCLING_AVAILABLE = True
except ImportError:
DOCLING_AVAILABLE = False
class Chunker: class Chunker:
@ -31,21 +24,17 @@ class Chunker:
self, self,
chunk_size: int = Config.processing.chunk_size, chunk_size: int = Config.processing.chunk_size,
): ):
if not DOCLING_AVAILABLE: from docling_core.transforms.chunker.hybrid_chunker import HybridChunker
raise ImportError( from docling_core.transforms.chunker.tokenizer.openai import OpenAITokenizer
"Docling is required for chunking. "
"Install with: pip install haiku.rag-slim[docling]"
)
from docling.chunking import HybridChunker # type: ignore
self.chunk_size = chunk_size self.chunk_size = chunk_size
tokenizer = OpenAITokenizer( tokenizer = OpenAITokenizer(
tokenizer=tiktoken.encoding_for_model("gpt-4o"), max_tokens=chunk_size 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. """Split the document into chunks using docling's structure-aware chunking.
Args: Args:

View file

@ -2,8 +2,10 @@
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from pathlib import Path 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): class DocumentConverter(ABC):
@ -24,7 +26,7 @@ class DocumentConverter(ABC):
pass pass
@abstractmethod @abstractmethod
def convert_file(self, path: Path) -> DoclingDocument: def convert_file(self, path: Path) -> "DoclingDocument":
"""Convert a file to DoclingDocument format. """Convert a file to DoclingDocument format.
Args: Args:
@ -39,7 +41,7 @@ class DocumentConverter(ABC):
pass pass
@abstractmethod @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. """Convert text content to DoclingDocument format.
Args: Args:

View file

@ -2,26 +2,12 @@
from io import BytesIO from io import BytesIO
from pathlib import Path from pathlib import Path
from typing import ClassVar from typing import TYPE_CHECKING, 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 haiku.rag.converters.base import DocumentConverter 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 if TYPE_CHECKING:
try: from docling_core.types.doc.document import DoclingDocument
import docling # noqa: F401
DOCLING_AVAILABLE = True
except ImportError:
DOCLING_AVAILABLE = False
class DoclingLocalConverter(DocumentConverter): class DoclingLocalConverter(DocumentConverter):
@ -114,7 +100,7 @@ class DoclingLocalConverter(DocumentConverter):
"""Return list of file extensions supported by this converter.""" """Return list of file extensions supported by this converter."""
return self.docling_extensions + self.text_extensions 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. """Convert a file to DoclingDocument using local docling.
Args: Args:
@ -126,23 +112,14 @@ class DoclingLocalConverter(DocumentConverter):
Raises: Raises:
ValueError: If the file cannot be converted. ValueError: If the file cannot be converted.
""" """
from docling.document_converter import DocumentConverter as DoclingDocConverter
try: try:
file_extension = path.suffix.lower() file_extension = path.suffix.lower()
if file_extension in self.docling_extensions: if file_extension in self.docling_extensions:
# Use docling for complex document formats # 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() 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) result = converter.convert(path)
return result.document return result.document
elif file_extension in self.text_extensions: elif file_extension in self.text_extensions:
@ -159,17 +136,11 @@ class DoclingLocalConverter(DocumentConverter):
else: else:
# Fallback: try to read as text and convert to DoclingDocument # Fallback: try to read as text and convert to DoclingDocument
content = path.read_text(encoding="utf-8") 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") 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: except Exception:
raise ValueError(f"Failed to parse file: {path}") 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. """Convert text content to DoclingDocument using local docling.
Args: Args:
@ -182,6 +153,9 @@ class DoclingLocalConverter(DocumentConverter):
Raises: Raises:
ValueError: If the text cannot be converted. ValueError: If the text cannot be converted.
""" """
from docling.document_converter import DocumentConverter as DoclingDocConverter
from docling_core.types.io import DocumentStream
try: try:
bytes_io = BytesIO(text.encode("utf-8")) bytes_io = BytesIO(text.encode("utf-8"))
doc_stream = DocumentStream(name=name, stream=bytes_io) doc_stream = DocumentStream(name=name, stream=bytes_io)