Merge pull request #115 from ggozad/fix/docling-lazy-load

Lazy-load  docling-related functions, speed up client

src/haiku/rag/client.py

@@ -9,7 +9,6 @@ from urllib.parse import urlparse
 import httpx
 
 from haiku.rag.config import Config
-from haiku.rag.reader import FileReader
 from haiku.rag.reranking import get_reranker
 from haiku.rag.store.engine import Store
 from haiku.rag.store.models.chunk import Chunk
@@ -17,7 +16,6 @@ from haiku.rag.store.models.document import Document
 from haiku.rag.store.repositories.chunk import ChunkRepository
 from haiku.rag.store.repositories.document import DocumentRepository
 from haiku.rag.store.repositories.settings import SettingsRepository
-from haiku.rag.utils import text_to_docling_document
 
 logger = logging.getLogger(__name__)
 
@@ -91,6 +89,9 @@ class HaikuRAG:
         Returns:
             The created Document instance.
         """
+        # Lazy import to avoid loading docling
+        from haiku.rag.utils import text_to_docling_document
+
         # Convert content to DoclingDocument for processing
         docling_document = text_to_docling_document(content)
 
@@ -127,6 +128,8 @@ class HaikuRAG:
             ValueError: If the file/URL cannot be parsed or doesn't exist
             httpx.RequestError: If URL request fails
         """
+        # Lazy import to avoid loading docling
+        from haiku.rag.reader import FileReader
 
         # Normalize metadata
         metadata = metadata or {}
@@ -181,6 +184,9 @@ class HaikuRAG:
         Raises:
             ValueError: If the file cannot be parsed or doesn't exist
         """
+        # Lazy import to avoid loading docling
+        from haiku.rag.reader import FileReader
+
         metadata = metadata or {}
 
         if source_path.suffix.lower() not in FileReader.extensions:
@@ -256,6 +262,9 @@ class HaikuRAG:
             ValueError: If the content cannot be parsed
             httpx.RequestError: If URL request fails
         """
+        # Lazy import to avoid loading docling
+        from haiku.rag.reader import FileReader
+
         metadata = metadata or {}
 
         async with httpx.AsyncClient() as client:
@@ -379,6 +388,9 @@ class HaikuRAG:
 
     async def update_document(self, document: Document) -> Document:
         """Update an existing document."""
+        # Lazy import to avoid loading docling
+        from haiku.rag.utils import text_to_docling_document
+
         # Convert content to DoclingDocument
         docling_document = text_to_docling_document(document.content)
 
@@ -597,6 +609,9 @@ class HaikuRAG:
         Yields:
             int: The ID of the document currently being processed
         """
+        # Lazy import to avoid loading docling
+        from haiku.rag.utils import text_to_docling_document
+
         await self.chunk_repository.delete_all()
         self.store.recreate_embeddings_table()
 

src/haiku/rag/monitor.py

@@ -1,21 +1,27 @@
 import logging
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 from watchfiles import Change, DefaultFilter, awatch
 
 from haiku.rag.client import HaikuRAG
-from haiku.rag.reader import FileReader
 from haiku.rag.store.models.document import Document
 
+if TYPE_CHECKING:
+    from haiku.rag.reader import FileReader
+
 logger = logging.getLogger(__name__)
 
 
 class FileFilter(DefaultFilter):
     def __init__(self, *, ignore_paths: list[Path] | None = None) -> None:
+        # Lazy import to avoid loading docling
+        from haiku.rag.reader import FileReader
+
         self.extensions = tuple(FileReader.extensions)
         super().__init__(ignore_paths=ignore_paths)
 
-    def __call__(self, change: "Change", path: str) -> bool:
+    def __call__(self, change: Change, path: str) -> bool:
         return path.endswith(self.extensions) and super().__call__(change, path)
 
 
@@ -50,11 +56,15 @@ class FileWatcher:
             uri = file.as_uri()
             existing_doc = await self.client.get_document_by_uri(uri)
             if existing_doc:
-                doc = await self.client.create_document_from_source(str(file))
+                result = await self.client.create_document_from_source(str(file))
+                # Since we're passing a file (not directory), result should be a single Document
+                doc = result if isinstance(result, Document) else result[0]
                 logger.info(f"Updated document {existing_doc.id} from {file}")
                 return doc
             else:
-                doc = await self.client.create_document_from_source(str(file))
+                result = await self.client.create_document_from_source(str(file))
+                # Since we're passing a file (not directory), result should be a single Document
+                doc = result if isinstance(result, Document) else result[0]
                 logger.info(f"Created new document {doc.id} from {file}")
                 return doc
         except Exception as e:

src/haiku/rag/store/repositories/chunk.py

@@ -1,17 +1,19 @@
 import inspect
 import json
 import logging
+from typing import TYPE_CHECKING
 from uuid import uuid4
 
-from docling_core.types.doc.document import DoclingDocument
 from lancedb.rerankers import RRFReranker
 
-from haiku.rag.chunker import chunker
 from haiku.rag.config import Config
 from haiku.rag.embeddings import get_embedder
 from haiku.rag.store.engine import DocumentRecord, Store
 from haiku.rag.store.models.chunk import Chunk
-from haiku.rag.utils import load_callable, text_to_docling_document
+from haiku.rag.utils import load_callable
+
+if TYPE_CHECKING:
+    from docling_core.types.doc.document import DoclingDocument
 
 logger = logging.getLogger(__name__)
 
@@ -142,9 +144,13 @@ class ChunkRepository:
         return chunks
 
     async def create_chunks_for_document(
-        self, document_id: str, document: DoclingDocument
+        self, document_id: str, document: "DoclingDocument"
     ) -> list[Chunk]:
         """Create chunks and embeddings for a document from DoclingDocument."""
+        # Lazy imports to avoid loading docling during module import
+        from haiku.rag.chunker import chunker
+        from haiku.rag.utils import text_to_docling_document
+
         # Optionally preprocess markdown before chunking
         processed_document = document
         preprocessor_path = Config.MARKDOWN_PREPROCESSOR

src/haiku/rag/store/repositories/document.py

@@ -4,12 +4,12 @@ from datetime import datetime
 from typing import TYPE_CHECKING
 from uuid import uuid4
 
-from docling_core.types.doc.document import DoclingDocument
-
 from haiku.rag.store.engine import DocumentRecord, Store
 from haiku.rag.store.models.document import Document
 
 if TYPE_CHECKING:
+    from docling_core.types.doc.document import DoclingDocument
+
     from haiku.rag.store.models.chunk import Chunk
 
 
@@ -171,7 +171,7 @@ class DocumentRepository:
     async def _create_with_docling(
         self,
         entity: Document,
-        docling_document: DoclingDocument,
+        docling_document: "DoclingDocument",
         chunks: list["Chunk"] | None = None,
     ) -> Document:
         """Create a document with its chunks and embeddings."""
@@ -211,7 +211,7 @@ class DocumentRepository:
             raise
 
     async def _update_with_docling(
-        self, entity: Document, docling_document: DoclingDocument
+        self, entity: Document, docling_document: "DoclingDocument"
     ) -> Document:
         """Update a document and regenerate its chunks."""
         assert entity.id is not None, "Document ID is required for update"