diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 686d7414..8c6c5ba8 100644 --- a/src/haiku/rag/client.py +++ b/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() diff --git a/src/haiku/rag/monitor.py b/src/haiku/rag/monitor.py index 713c5834..3b269b38 100644 --- a/src/haiku/rag/monitor.py +++ b/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: diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index c8fac655..bfd2e624 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/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 diff --git a/src/haiku/rag/store/repositories/document.py b/src/haiku/rag/store/repositories/document.py index ef951d6c..70643023 100644 --- a/src/haiku/rag/store/repositories/document.py +++ b/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"