Make docling optional, depend on docling-core
This commit is contained in:
parent
2f9c907031
commit
b02514dc82
5 changed files with 50 additions and 10 deletions
|
|
@ -1,12 +1,18 @@
|
||||||
from typing import ClassVar
|
from typing import ClassVar
|
||||||
|
|
||||||
import tiktoken
|
import tiktoken
|
||||||
from docling.chunking import HybridChunker # type: ignore
|
|
||||||
from docling_core.transforms.chunker.tokenizer.openai import OpenAITokenizer
|
from docling_core.transforms.chunker.tokenizer.openai import OpenAITokenizer
|
||||||
from docling_core.types.doc.document import DoclingDocument
|
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
|
||||||
|
try:
|
||||||
|
import docling # noqa: F401
|
||||||
|
|
||||||
|
DOCLING_AVAILABLE = True
|
||||||
|
except ImportError:
|
||||||
|
DOCLING_AVAILABLE = False
|
||||||
|
|
||||||
|
|
||||||
class Chunker:
|
class Chunker:
|
||||||
"""A class that chunks text into smaller pieces for embedding and retrieval.
|
"""A class that chunks text into smaller pieces for embedding and retrieval.
|
||||||
|
|
@ -24,6 +30,13 @@ class Chunker:
|
||||||
self,
|
self,
|
||||||
chunk_size: int = Config.processing.chunk_size,
|
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
|
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
|
||||||
|
|
|
||||||
|
|
@ -1,11 +1,17 @@
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import ClassVar
|
from typing import ClassVar
|
||||||
|
|
||||||
from docling.document_converter import DocumentConverter
|
|
||||||
from docling_core.types.doc.document import DoclingDocument
|
from docling_core.types.doc.document import DoclingDocument
|
||||||
|
|
||||||
from haiku.rag.utils import text_to_docling_document
|
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:
|
class FileReader:
|
||||||
# Extensions supported by docling
|
# Extensions supported by docling
|
||||||
|
|
@ -95,6 +101,13 @@ class FileReader:
|
||||||
|
|
||||||
if file_extension in FileReader.docling_extensions:
|
if file_extension in FileReader.docling_extensions:
|
||||||
# Use docling for complex document formats
|
# 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()
|
converter = DocumentConverter()
|
||||||
result = converter.convert(path)
|
result = converter.convert(path)
|
||||||
return result.document
|
return result.document
|
||||||
|
|
|
||||||
|
|
@ -103,9 +103,16 @@ def text_to_docling_document(text: str, name: str = "content.md"):
|
||||||
Returns:
|
Returns:
|
||||||
A DoclingDocument created from the text content.
|
A DoclingDocument created from the text content.
|
||||||
"""
|
"""
|
||||||
# Lazy import docling deps to keep import-time light
|
try:
|
||||||
from docling.document_converter import DocumentConverter # type: ignore
|
import docling # noqa: F401
|
||||||
from docling_core.types.io import DocumentStream # type: ignore
|
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"))
|
bytes_io = BytesIO(text.encode("utf-8"))
|
||||||
doc_stream = DocumentStream(name=name, stream=bytes_io)
|
doc_stream = DocumentStream(name=name, stream=bytes_io)
|
||||||
|
|
@ -168,11 +175,15 @@ def load_callable(path: str):
|
||||||
def prefetch_models():
|
def prefetch_models():
|
||||||
"""Prefetch runtime models (Docling + Ollama as configured)."""
|
"""Prefetch runtime models (Docling + Ollama as configured)."""
|
||||||
import httpx
|
import httpx
|
||||||
from docling.utils.model_downloader import download_models
|
|
||||||
|
|
||||||
from haiku.rag.config import Config
|
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
|
# Collect Ollama models from config
|
||||||
required_models: set[str] = set()
|
required_models: set[str] = set()
|
||||||
|
|
|
||||||
|
|
@ -22,6 +22,7 @@ classifiers = [
|
||||||
]
|
]
|
||||||
|
|
||||||
dependencies = [
|
dependencies = [
|
||||||
|
"docling-core>=2.48.4",
|
||||||
"fastmcp>=2.13.0.2",
|
"fastmcp>=2.13.0.2",
|
||||||
"httpx>=0.28.1",
|
"httpx>=0.28.1",
|
||||||
"lancedb>=0.25.2",
|
"lancedb>=0.25.2",
|
||||||
|
|
|
||||||
2
uv.lock
2
uv.lock
|
|
@ -1222,6 +1222,7 @@ name = "haiku-rag-slim"
|
||||||
version = "0.13.3"
|
version = "0.13.3"
|
||||||
source = { editable = "haiku_rag_slim" }
|
source = { editable = "haiku_rag_slim" }
|
||||||
dependencies = [
|
dependencies = [
|
||||||
|
{ name = "docling-core" },
|
||||||
{ name = "fastmcp" },
|
{ name = "fastmcp" },
|
||||||
{ name = "httpx" },
|
{ name = "httpx" },
|
||||||
{ name = "lancedb" },
|
{ name = "lancedb" },
|
||||||
|
|
@ -1261,6 +1262,7 @@ zeroentropy = [
|
||||||
requires-dist = [
|
requires-dist = [
|
||||||
{ name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.0.0" },
|
{ name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.0.0" },
|
||||||
{ name = "docling", marker = "extra == 'docling'", specifier = ">=2.58.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 = "fasta2a", marker = "extra == 'a2a'", specifier = ">=0.1.0" },
|
||||||
{ name = "fastmcp", specifier = ">=2.13.0.2" },
|
{ name = "fastmcp", specifier = ">=2.13.0.2" },
|
||||||
{ name = "httpx", specifier = ">=0.28.1" },
|
{ name = "httpx", specifier = ">=0.28.1" },
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue