danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Lazy load docling & friends

Browse source
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
Show stats Download patch file Download diff file Expand all files Collapse all files

27
haiku_rag_slim/haiku/rag/chunker.py
View file

@ -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()

8
src/haiku/rag/converters/base.py
View file

@ -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:

48
src/haiku/rag/converters/docling_local.py
View file

@ -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}")