diff --git a/haiku_rag_slim/haiku/rag/chunker.py b/haiku_rag_slim/haiku/rag/chunker.py index f478d64c..fce09df5 100644 --- a/haiku_rag_slim/haiku/rag/chunker.py +++ b/haiku_rag_slim/haiku/rag/chunker.py @@ -1,12 +1,18 @@ from typing import ClassVar import tiktoken -from docling.chunking import HybridChunker # type: ignore 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 + class Chunker: """A class that chunks text into smaller pieces for embedding and retrieval. @@ -24,6 +30,13 @@ 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 + self.chunk_size = chunk_size tokenizer = OpenAITokenizer( tokenizer=tiktoken.encoding_for_model("gpt-4o"), max_tokens=chunk_size diff --git a/haiku_rag_slim/haiku/rag/reader.py b/haiku_rag_slim/haiku/rag/reader.py index 8e7fba2c..c710be49 100644 --- a/haiku_rag_slim/haiku/rag/reader.py +++ b/haiku_rag_slim/haiku/rag/reader.py @@ -1,11 +1,17 @@ from pathlib import Path from typing import ClassVar -from docling.document_converter import DocumentConverter from docling_core.types.doc.document import DoclingDocument - from haiku.rag.utils import text_to_docling_document +# Check if docling is available +try: + import docling # noqa: F401 + + DOCLING_AVAILABLE = True +except ImportError: + DOCLING_AVAILABLE = False + class FileReader: # Extensions supported by docling @@ -95,6 +101,13 @@ class FileReader: if file_extension in FileReader.docling_extensions: # Use docling for complex document formats + 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() result = converter.convert(path) return result.document diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index 228aa7c2..950cfe87 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -103,9 +103,16 @@ def text_to_docling_document(text: str, name: str = "content.md"): Returns: A DoclingDocument created from the text content. """ - # Lazy import docling deps to keep import-time light - from docling.document_converter import DocumentConverter # type: ignore - from docling_core.types.io import DocumentStream # type: ignore + try: + import docling # noqa: F401 + except ImportError as e: + raise ImportError( + "Docling is required for document conversion. " + "Install with: pip install haiku.rag-slim[docling]" + ) from e + + from docling.document_converter import DocumentConverter + from docling_core.types.io import DocumentStream bytes_io = BytesIO(text.encode("utf-8")) doc_stream = DocumentStream(name=name, stream=bytes_io) @@ -168,11 +175,15 @@ def load_callable(path: str): def prefetch_models(): """Prefetch runtime models (Docling + Ollama as configured).""" import httpx - from docling.utils.model_downloader import download_models - from haiku.rag.config import Config - download_models() + try: + from docling.utils.model_downloader import download_models + + download_models() + except ImportError: + # Docling not installed, skip downloading docling models + pass # Collect Ollama models from config required_models: set[str] = set() diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 6d5ff663..1e867157 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -22,6 +22,7 @@ classifiers = [ ] dependencies = [ + "docling-core>=2.48.4", "fastmcp>=2.13.0.2", "httpx>=0.28.1", "lancedb>=0.25.2", diff --git a/uv.lock b/uv.lock index 8a5ed5e6..549468d5 100644 --- a/uv.lock +++ b/uv.lock @@ -1222,6 +1222,7 @@ name = "haiku-rag-slim" version = "0.13.3" source = { editable = "haiku_rag_slim" } dependencies = [ + { name = "docling-core" }, { name = "fastmcp" }, { name = "httpx" }, { name = "lancedb" }, @@ -1261,6 +1262,7 @@ zeroentropy = [ requires-dist = [ { name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.0.0" }, { name = "docling", marker = "extra == 'docling'", specifier = ">=2.58.0" }, + { name = "docling-core", specifier = ">=2.48.4" }, { name = "fasta2a", marker = "extra == 'a2a'", specifier = ">=0.1.0" }, { name = "fastmcp", specifier = ">=2.13.0.2" }, { name = "httpx", specifier = ">=0.28.1" },