From 3224d1f20cf4a362cfecf72d6fb14d9690db72b6 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 24 Apr 2026 12:30:15 +0300 Subject: [PATCH] Extract document orchestration into client/documents.py --- haiku_rag_slim/haiku/rag/client/__init__.py | 590 +------------------ haiku_rag_slim/haiku/rag/client/documents.py | 491 +++++++++++++++ haiku_rag_slim/haiku/rag/client/rebuild.py | 3 +- tests/store/test_document_items.py | 12 +- tests/test_client.py | 64 +- tests/test_context.py | 5 +- tests/test_context_enhancement.py | 3 +- 7 files changed, 561 insertions(+), 607 deletions(-) create mode 100644 haiku_rag_slim/haiku/rag/client/documents.py diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index d33cdb8c..5dd47305 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -183,120 +183,6 @@ class HaikuRAG: return await generate_title(self._config, document) - async def _store_document_with_chunks( - self, - document: Document, - chunks: list[Chunk], - docling_document: "DoclingDocument", - ) -> Document: - """Store a document with chunks, embedding any that lack embeddings. - - Handles versioning/rollback on failure. - - Args: - document: The document to store (will be created). - chunks: Chunks to store (will be embedded if lacking embeddings). - docling_document: The DoclingDocument to extract items from. - - Returns: - The created Document instance with ID set. - """ - # Ensure all chunks have embeddings before storing - chunks = await self._ensure_chunks_embedded(chunks) - - # Snapshot table versions for versioned rollback (if supported) - versions = await self.store.current_table_versions() - - # Create the document - created_doc = await self.document_repository.create(document) - - try: - assert created_doc.id is not None, ( - "Document ID should not be None after creation" - ) - # Set document_id and order for all chunks - for order, chunk in enumerate(chunks): - chunk.document_id = created_doc.id - chunk.order = order - - # Batch create all chunks in a single operation - await self.chunk_repository.create(chunks) - - # Extract and store document items for context expansion - items = extract_items(created_doc.id, docling_document) - await self.document_item_repository.create_items(created_doc.id, items) - - # Vacuum old versions in background (non-blocking) if auto_vacuum enabled - if self._config.storage.auto_vacuum: - self._schedule_vacuum() - - return created_doc - except Exception: - # Roll back to the captured versions and re-raise - await self.store.restore_table_versions(versions) - raise - - async def _update_document_with_chunks( - self, - document: Document, - chunks: list[Chunk], - docling_document: "DoclingDocument | None" = None, - ) -> Document: - """Update a document and replace its chunks, embedding any that lack embeddings. - - Handles versioning/rollback on failure. - - Args: - document: The document to update (must have ID set). - chunks: Chunks to replace existing (will be embedded if lacking embeddings). - docling_document: The DoclingDocument to extract items from. - When None, existing items are preserved. - - Returns: - The updated Document instance. - """ - assert document.id is not None, "Document ID is required for update" - - # Ensure all chunks have embeddings before storing - chunks = await self._ensure_chunks_embedded(chunks) - - # Snapshot table versions for versioned rollback - versions = await self.store.current_table_versions() - - # Delete existing chunks before writing new ones - await self.chunk_repository.delete_by_document_id(document.id) - - try: - # Update the document - updated_doc = await self.document_repository.update(document) - - # Set document_id and order for all chunks - assert updated_doc.id is not None - for order, chunk in enumerate(chunks): - chunk.document_id = updated_doc.id - chunk.order = order - - # Batch create all chunks in a single operation - await self.chunk_repository.create(chunks) - - # Replace document items when a new DoclingDocument is provided - if docling_document is not None: - await self.document_item_repository.delete_by_document_id( - updated_doc.id - ) - items = extract_items(updated_doc.id, docling_document) - await self.document_item_repository.create_items(updated_doc.id, items) - - # Vacuum old versions in background (non-blocking) if auto_vacuum enabled - if self._config.storage.auto_vacuum: - self._schedule_vacuum() - - return updated_doc - except Exception: - # Roll back to the captured versions and re-raise - await self.store.restore_table_versions(versions) - raise - async def create_document( self, content: str, @@ -305,49 +191,9 @@ class HaikuRAG: metadata: dict | None = None, format: str = "md", ) -> Document: - """Create a new document from text content. + from haiku.rag.client.documents import create_document - Converts the content, chunks it, and generates embeddings. - - Args: - content: The text content of the document. - uri: Optional URI identifier for the document. - title: Optional title for the document. - metadata: Optional metadata dictionary. - format: The format of the content ("md", "html", or "plain"). - Defaults to "md". Use "plain" for plain text without parsing. - - Returns: - The created Document instance. - """ - from haiku.rag.embeddings import embed_chunks - - # Convert → Chunk → Embed using primitives - converter = get_converter(self._config) - docling_document = await converter.convert_text(content, format=format) - chunks = await self.chunk(docling_document) - embedded_chunks = await embed_chunks(chunks, self._config) - - # Store markdown export as content for better display/readability - # The original content is preserved in docling_document - stored_content = docling_document.export_to_markdown() - - if title is None: - title = await self._resolve_title(docling_document, stored_content) - - # Create document model - document = Document( - content=stored_content, - uri=uri, - title=title, - metadata=metadata or {}, - ) - document.set_docling(docling_document) - - # Store document and chunks - return await self._store_document_with_chunks( - document, embedded_chunks, docling_document - ) + return await create_document(self, content, uri, title, metadata, format) async def import_document( self, @@ -357,316 +203,43 @@ class HaikuRAG: title: str | None = None, metadata: dict | None = None, ) -> Document: - """Import a pre-processed document with chunks. + from haiku.rag.client.documents import import_document - Use this when document conversion, chunking, and embedding were done - externally and you want to store the results in haiku.rag. - - Args: - docling_document: The DoclingDocument to import. - chunks: Pre-created chunks. Chunks without embeddings will be - automatically embedded. - uri: Optional URI identifier for the document. - title: Optional title for the document. - metadata: Optional metadata dictionary. - - Returns: - The created Document instance. - """ - content = docling_document.export_to_markdown() - if title is None: - title = await self._resolve_title(docling_document, content) - - document = Document( - content=content, - uri=uri, - title=title, - metadata=metadata or {}, - ) - document.set_docling(docling_document) - - return await self._store_document_with_chunks( - document, chunks, docling_document + return await import_document( + self, docling_document, chunks, uri, title, metadata ) async def create_document_from_source( - self, source: str | Path, title: str | None = None, metadata: dict | None = None + self, + source: str | Path, + title: str | None = None, + metadata: dict | None = None, ) -> Document | list[Document]: - """Create or update document(s) from a file path, directory, or URL. + from haiku.rag.client.documents import create_document_from_source - Checks if a document with the same URI already exists: - - If MD5 is unchanged, returns existing document - - If MD5 changed, updates the document - - If no document exists, creates a new one + return await create_document_from_source(self, source, title, metadata) - Args: - source: File path, directory (as string or Path), or URL to parse - title: Optional title (only used for single files, not directories) - metadata: Optional metadata dictionary + async def update_document( + self, + document_id: str, + content: str | None = None, + metadata: dict | None = None, + chunks: list[Chunk] | None = None, + title: str | None = None, + docling_document: "DoclingDocument | None" = None, + ) -> Document: + from haiku.rag.client.documents import update_document - Returns: - Document instance (created, updated, or existing) for single files/URLs - List of Document instances for directories - - Raises: - ValueError: If the file/URL cannot be parsed or doesn't exist - httpx.RequestError: If URL request fails - """ - # Normalize metadata - metadata = metadata or {} - - # Check if it's a URL - source_str = str(source) - parsed_url = urlparse(source_str) - if parsed_url.scheme in ("http", "https"): - return await self._create_or_update_document_from_url( - source_str, title=title, metadata=metadata - ) - elif parsed_url.scheme == "file": - # Handle file:// URI by converting to path - source_path = Path(parsed_url.path) - else: - # Handle as regular file path - source_path = Path(source) if isinstance(source, str) else source - - # Handle directories - if source_path.is_dir(): - from haiku.rag.monitor import FileFilter - - documents = [] - filter = FileFilter( - ignore_patterns=self._config.monitor.ignore_patterns or None, - include_patterns=self._config.monitor.include_patterns or None, - ) - for path in source_path.rglob("*"): - if path.is_file() and filter.include_file(str(path)): - doc = await self._create_document_from_file( - path, title=None, metadata=metadata - ) - documents.append(doc) - return documents - - # Handle single file - return await self._create_document_from_file( - source_path, title=title, metadata=metadata + return await update_document( + self, + document_id, + content, + metadata, + chunks, + title, + docling_document, ) - async def _create_document_from_file( - self, source_path: Path, title: str | None = None, metadata: dict | None = None - ) -> Document: - """Create or update a document from a single file path. - - Args: - source_path: Path to the file - title: Optional title - metadata: Optional metadata dictionary - - Returns: - Document instance (created, updated, or existing) - - Raises: - ValueError: If the file cannot be parsed or doesn't exist - """ - from haiku.rag.embeddings import embed_chunks - - metadata = metadata or {} - - converter = get_converter(self._config) - if source_path.suffix.lower() not in converter.supported_extensions: - raise ValueError(f"Unsupported file extension: {source_path.suffix}") - - if not source_path.exists(): - raise ValueError(f"File does not exist: {source_path}") - - uri = source_path.absolute().as_uri() - md5_hash = hashlib.md5( - source_path.read_bytes(), usedforsecurity=False - ).hexdigest() - - # Get content type from file extension (do before early return) - content_type, _ = mimetypes.guess_type(str(source_path)) - if not content_type: - content_type = "application/octet-stream" - # Merge metadata with contentType and md5 - metadata.update({"contentType": content_type, "md5": md5_hash}) - - # Check if document already exists - existing_doc = await self.get_document_by_uri(uri) - if existing_doc and existing_doc.metadata.get("md5") == md5_hash: - # MD5 unchanged; update title/metadata if provided - updated = False - if title is not None and title != existing_doc.title: - existing_doc.title = title - updated = True - - # Check if metadata actually changed (beyond contentType and md5) - merged_metadata = {**(existing_doc.metadata or {}), **metadata} - if merged_metadata != existing_doc.metadata: - existing_doc.metadata = merged_metadata - updated = True - - if updated: - return await self.document_repository.update(existing_doc) - return existing_doc - - # Convert → Chunk → Embed using primitives - docling_document = await self.convert(source_path) - chunks = await self.chunk(docling_document) - embedded_chunks = await embed_chunks(chunks, self._config) - - stored_content = docling_document.export_to_markdown() - - if existing_doc: - # Update existing document and rechunk - existing_doc.content = stored_content - existing_doc.metadata = metadata - existing_doc.set_docling(docling_document) - if title is not None: - existing_doc.title = title - elif existing_doc.title is None: - existing_doc.title = await self._resolve_title( - docling_document, stored_content - ) - return await self._update_document_with_chunks( - existing_doc, embedded_chunks, docling_document - ) - else: - # Create new document - if title is None: - title = await self._resolve_title(docling_document, stored_content) - document = Document( - content=stored_content, - uri=uri, - title=title, - metadata=metadata, - ) - document.set_docling(docling_document) - return await self._store_document_with_chunks( - document, embedded_chunks, docling_document - ) - - async def _create_or_update_document_from_url( - self, url: str, title: str | None = None, metadata: dict | None = None - ) -> Document: - """Create or update a document from a URL by downloading and parsing the content. - - Checks if a document with the same URI already exists: - - If MD5 is unchanged, returns existing document - - If MD5 changed, updates the document - - If no document exists, creates a new one - - Args: - url: URL to download and parse - metadata: Optional metadata dictionary - - Returns: - Document instance (created, updated, or existing) - - Raises: - ValueError: If the content cannot be parsed - httpx.RequestError: If URL request fails - """ - from haiku.rag.embeddings import embed_chunks - - metadata = metadata or {} - - converter = get_converter(self._config) - supported_extensions = converter.supported_extensions - - async with httpx.AsyncClient() as client: - response = await client.get(url) - response.raise_for_status() - - md5_hash = hashlib.md5(response.content).hexdigest() - - # Get content type early (used for potential no-op update) - content_type = response.headers.get("content-type", "").lower() - - # Check if document already exists - existing_doc = await self.get_document_by_uri(url) - if existing_doc and existing_doc.metadata.get("md5") == md5_hash: - # MD5 unchanged; update title/metadata if provided - updated = False - if title is not None and title != existing_doc.title: - existing_doc.title = title - updated = True - - metadata.update({"contentType": content_type, "md5": md5_hash}) - # Check if metadata actually changed (beyond contentType and md5) - merged_metadata = {**(existing_doc.metadata or {}), **metadata} - if merged_metadata != existing_doc.metadata: - existing_doc.metadata = merged_metadata - updated = True - - if updated: - return await self.document_repository.update(existing_doc) - return existing_doc - file_extension = self._get_extension_from_content_type_or_url( - url, content_type - ) - - if file_extension not in supported_extensions: - raise ValueError( - f"Unsupported content type/extension: {content_type}/{file_extension}" - ) - - # Create a temporary file with the appropriate extension - with tempfile.NamedTemporaryFile( - mode="wb", suffix=file_extension, delete=False - ) as temp_file: - temp_file.write(response.content) - temp_file.flush() - temp_path = Path(temp_file.name) - - try: - # Convert → Chunk → Embed using primitives - docling_document = await self.convert(temp_path) - chunks = await self.chunk(docling_document) - embedded_chunks = await embed_chunks(chunks, self._config) - finally: - temp_path.unlink(missing_ok=True) - - # Merge metadata with contentType and md5 - metadata.update({"contentType": content_type, "md5": md5_hash}) - - stored_content = docling_document.export_to_markdown() - - if existing_doc: - # Update existing document and rechunk - existing_doc.content = stored_content - existing_doc.metadata = metadata - existing_doc.set_docling(docling_document) - if title is not None: - existing_doc.title = title - elif existing_doc.title is None: - existing_doc.title = await self._resolve_title( - docling_document, stored_content - ) - return await self._update_document_with_chunks( - existing_doc, embedded_chunks, docling_document - ) - else: - # Create new document - if title is None: - title = await self._resolve_title(docling_document, stored_content) - document = Document( - content=stored_content, - uri=url, - title=title, - metadata=metadata, - ) - document.set_docling(docling_document) - return await self._store_document_with_chunks( - document, embedded_chunks, docling_document - ) - - def _get_extension_from_content_type_or_url( - self, url: str, content_type: str - ) -> str: - from haiku.rag.client.processing import get_extension_from_content_type_or_url - - return get_extension_from_content_type_or_url(url, content_type) - async def get_document_by_id(self, document_id: str) -> Document | None: """Get a document by its ID. @@ -724,99 +297,6 @@ class HaikuRAG: return None - async def update_document( - self, - document_id: str, - content: str | None = None, - metadata: dict | None = None, - chunks: list[Chunk] | None = None, - title: str | None = None, - docling_document: "DoclingDocument | None" = None, - ) -> Document: - """Update a document by ID. - - Updates specified fields. When content or docling_document is provided, - the document is rechunked and re-embedded. Updates to only metadata or title - skip rechunking for efficiency. - - Args: - document_id: The ID of the document to update. - content: New content (mutually exclusive with docling_document). - metadata: New metadata dict. - chunks: Custom chunks (will be embedded if missing embeddings). - title: New title. - docling_document: DoclingDocument to replace content (mutually exclusive with content). - - Returns: - The updated Document instance. - - Raises: - ValueError: If document not found, or if both content and docling_document - are provided. - """ - from haiku.rag.embeddings import embed_chunks - - # Validate: content and docling_document are mutually exclusive - if content is not None and docling_document is not None: - raise ValueError( - "content and docling_document are mutually exclusive. " - "Provide one or the other, not both." - ) - - # Fetch the existing document - existing_doc = await self.get_document_by_id(document_id) - if existing_doc is None: - raise ValueError(f"Document with ID {document_id} not found") - - # Update metadata/title fields - if title is not None: - existing_doc.title = title - if metadata is not None: - existing_doc.metadata = metadata - - # Only metadata/title update - no rechunking needed - if content is None and chunks is None and docling_document is None: - return await self.document_repository.update(existing_doc) - - # Custom chunks provided - use them as-is - if chunks is not None: - # Store docling data if provided - if docling_document is not None: - existing_doc.content = docling_document.export_to_markdown() - existing_doc.set_docling(docling_document) - elif content is not None: - existing_doc.content = content - - return await self._update_document_with_chunks( - existing_doc, chunks, docling_document - ) - - # DoclingDocument provided without chunks - chunk and embed using primitives - if docling_document is not None: - existing_doc.content = docling_document.export_to_markdown() - existing_doc.set_docling(docling_document) - - new_chunks = await self.chunk(docling_document) - embedded_chunks = await embed_chunks(new_chunks, self._config) - return await self._update_document_with_chunks( - existing_doc, embedded_chunks, docling_document - ) - - # Content provided without chunks - convert, chunk, and embed using primitives - assert content is not None - existing_doc.content = content - converter = get_converter(self._config) - converted_docling = await converter.convert_text( - existing_doc.content, format="md" - ) - existing_doc.set_docling(converted_docling) - - new_chunks = await self.chunk(converted_docling) - embedded_chunks = await embed_chunks(new_chunks, self._config) - return await self._update_document_with_chunks( - existing_doc, embedded_chunks, converted_docling - ) - async def delete_document(self, document_id: str) -> bool: """Delete a document by its ID.""" return await self.document_repository.delete(document_id) @@ -920,18 +400,6 @@ class HaikuRAG: async for doc_id in rebuild_database(self, mode): yield doc_id - def _check_source_accessible(self, uri: str) -> bool: - """Check if a document's source URI is accessible.""" - parsed_url = urlparse(uri) - try: - if parsed_url.scheme == "file": - return Path(parsed_url.path).exists() - elif parsed_url.scheme in ("http", "https"): - return True - return False - except Exception: - return False - async def vacuum(self) -> None: """Optimize and clean up old versions across all tables.""" await self.store.vacuum() diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py new file mode 100644 index 00000000..3435f91f --- /dev/null +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -0,0 +1,491 @@ +import hashlib +import mimetypes +import tempfile +from pathlib import Path +from typing import TYPE_CHECKING +from urllib.parse import urlparse + +import httpx + +from haiku.rag.converters import get_converter +from haiku.rag.store.models.chunk import Chunk +from haiku.rag.store.models.document import Document +from haiku.rag.store.models.document_item import extract_items + +if TYPE_CHECKING: + from docling_core.types.doc.document import DoclingDocument + + from haiku.rag.client import HaikuRAG + + +async def _store_document_with_chunks( + client: "HaikuRAG", + document: Document, + chunks: list[Chunk], + docling_document: "DoclingDocument", +) -> Document: + """Store a document with chunks, embedding any that lack embeddings. + + Handles versioning/rollback on failure. + """ + # Ensure all chunks have embeddings before storing + chunks = await client._ensure_chunks_embedded(chunks) + + # Snapshot table versions for versioned rollback (if supported) + versions = await client.store.current_table_versions() + + # Create the document + created_doc = await client.document_repository.create(document) + + try: + assert created_doc.id is not None, ( + "Document ID should not be None after creation" + ) + # Set document_id and order for all chunks + for order, chunk in enumerate(chunks): + chunk.document_id = created_doc.id + chunk.order = order + + # Batch create all chunks in a single operation + await client.chunk_repository.create(chunks) + + # Extract and store document items for context expansion + items = extract_items(created_doc.id, docling_document) + await client.document_item_repository.create_items(created_doc.id, items) + + # Vacuum old versions in background (non-blocking) if auto_vacuum enabled + if client._config.storage.auto_vacuum: + client._schedule_vacuum() + + return created_doc + except Exception: + # Roll back to the captured versions and re-raise + await client.store.restore_table_versions(versions) + raise + + +async def _update_document_with_chunks( + client: "HaikuRAG", + document: Document, + chunks: list[Chunk], + docling_document: "DoclingDocument | None" = None, +) -> Document: + """Update a document and replace its chunks, embedding any that lack embeddings. + + Handles versioning/rollback on failure. When `docling_document` is None, + existing items are preserved. + """ + assert document.id is not None, "Document ID is required for update" + + chunks = await client._ensure_chunks_embedded(chunks) + + versions = await client.store.current_table_versions() + + # Delete existing chunks before writing new ones + await client.chunk_repository.delete_by_document_id(document.id) + + try: + updated_doc = await client.document_repository.update(document) + + assert updated_doc.id is not None + for order, chunk in enumerate(chunks): + chunk.document_id = updated_doc.id + chunk.order = order + + await client.chunk_repository.create(chunks) + + # Replace document items when a new DoclingDocument is provided + if docling_document is not None: + await client.document_item_repository.delete_by_document_id(updated_doc.id) + items = extract_items(updated_doc.id, docling_document) + await client.document_item_repository.create_items(updated_doc.id, items) + + if client._config.storage.auto_vacuum: + client._schedule_vacuum() + + return updated_doc + except Exception: + await client.store.restore_table_versions(versions) + raise + + +async def create_document( + client: "HaikuRAG", + content: str, + uri: str | None = None, + title: str | None = None, + metadata: dict | None = None, + format: str = "md", +) -> Document: + """Create a new document from text content. + + Converts the content, chunks it, and generates embeddings. + """ + from haiku.rag.embeddings import embed_chunks + + # Convert → Chunk → Embed using primitives + converter = get_converter(client._config) + docling_document = await converter.convert_text(content, format=format) + chunks = await client.chunk(docling_document) + embedded_chunks = await embed_chunks(chunks, client._config) + + # Store markdown export as content for better display/readability. + # The original is preserved in docling_document. + stored_content = docling_document.export_to_markdown() + + if title is None: + title = await client._resolve_title(docling_document, stored_content) + + document = Document( + content=stored_content, + uri=uri, + title=title, + metadata=metadata or {}, + ) + document.set_docling(docling_document) + + return await _store_document_with_chunks( + client, document, embedded_chunks, docling_document + ) + + +async def import_document( + client: "HaikuRAG", + docling_document: "DoclingDocument", + chunks: list[Chunk], + uri: str | None = None, + title: str | None = None, + metadata: dict | None = None, +) -> Document: + """Import a pre-processed document with chunks. + + Use this when conversion, chunking, and embedding were done externally. + Chunks without embeddings will be automatically embedded. + """ + content = docling_document.export_to_markdown() + if title is None: + title = await client._resolve_title(docling_document, content) + + document = Document( + content=content, + uri=uri, + title=title, + metadata=metadata or {}, + ) + document.set_docling(docling_document) + + return await _store_document_with_chunks(client, document, chunks, docling_document) + + +async def create_document_from_source( + client: "HaikuRAG", + source: str | Path, + title: str | None = None, + metadata: dict | None = None, +) -> Document | list[Document]: + """Create or update document(s) from a file path, directory, or URL. + + Checks if a document with the same URI already exists: + - If MD5 is unchanged, returns existing document + - If MD5 changed, updates the document + - If no document exists, creates a new one + + Returns a single Document for files/URLs, a list for directories. + """ + metadata = metadata or {} + + source_str = str(source) + parsed_url = urlparse(source_str) + if parsed_url.scheme in ("http", "https"): + return await _create_or_update_document_from_url( + client, source_str, title=title, metadata=metadata + ) + elif parsed_url.scheme == "file": + source_path = Path(parsed_url.path) + else: + source_path = Path(source) if isinstance(source, str) else source + + if source_path.is_dir(): + from haiku.rag.monitor import FileFilter + + documents = [] + filter = FileFilter( + ignore_patterns=client._config.monitor.ignore_patterns or None, + include_patterns=client._config.monitor.include_patterns or None, + ) + for path in source_path.rglob("*"): + if path.is_file() and filter.include_file(str(path)): + doc = await _create_document_from_file( + client, path, title=None, metadata=metadata + ) + documents.append(doc) + return documents + + return await _create_document_from_file( + client, source_path, title=title, metadata=metadata + ) + + +async def _create_document_from_file( + client: "HaikuRAG", + source_path: Path, + title: str | None = None, + metadata: dict | None = None, +) -> Document: + """Create or update a document from a single file path.""" + from haiku.rag.embeddings import embed_chunks + + metadata = metadata or {} + + converter = get_converter(client._config) + if source_path.suffix.lower() not in converter.supported_extensions: + raise ValueError(f"Unsupported file extension: {source_path.suffix}") + + if not source_path.exists(): + raise ValueError(f"File does not exist: {source_path}") + + uri = source_path.absolute().as_uri() + md5_hash = hashlib.md5(source_path.read_bytes(), usedforsecurity=False).hexdigest() + + content_type, _ = mimetypes.guess_type(str(source_path)) + if not content_type: + content_type = "application/octet-stream" + metadata.update({"contentType": content_type, "md5": md5_hash}) + + # Check if document already exists + existing_doc = await client.get_document_by_uri(uri) + if existing_doc and existing_doc.metadata.get("md5") == md5_hash: + # MD5 unchanged; update title/metadata if provided + updated = False + if title is not None and title != existing_doc.title: + existing_doc.title = title + updated = True + + merged_metadata = {**(existing_doc.metadata or {}), **metadata} + if merged_metadata != existing_doc.metadata: + existing_doc.metadata = merged_metadata + updated = True + + if updated: + return await client.document_repository.update(existing_doc) + return existing_doc + + # Convert → Chunk → Embed + docling_document = await client.convert(source_path) + chunks = await client.chunk(docling_document) + embedded_chunks = await embed_chunks(chunks, client._config) + + stored_content = docling_document.export_to_markdown() + + if existing_doc: + # Update existing document and rechunk + existing_doc.content = stored_content + existing_doc.metadata = metadata + existing_doc.set_docling(docling_document) + if title is not None: + existing_doc.title = title + elif existing_doc.title is None: + existing_doc.title = await client._resolve_title( + docling_document, stored_content + ) + return await _update_document_with_chunks( + client, existing_doc, embedded_chunks, docling_document + ) + else: + if title is None: + title = await client._resolve_title(docling_document, stored_content) + document = Document( + content=stored_content, + uri=uri, + title=title, + metadata=metadata, + ) + document.set_docling(docling_document) + return await _store_document_with_chunks( + client, document, embedded_chunks, docling_document + ) + + +async def _create_or_update_document_from_url( + client: "HaikuRAG", + url: str, + title: str | None = None, + metadata: dict | None = None, +) -> Document: + """Create or update a document from a URL by downloading and parsing the content.""" + from haiku.rag.client.processing import get_extension_from_content_type_or_url + from haiku.rag.embeddings import embed_chunks + + metadata = metadata or {} + + converter = get_converter(client._config) + supported_extensions = converter.supported_extensions + + async with httpx.AsyncClient() as http: + response = await http.get(url) + response.raise_for_status() + + md5_hash = hashlib.md5(response.content).hexdigest() + + content_type = response.headers.get("content-type", "").lower() + + # Check if document already exists + existing_doc = await client.get_document_by_uri(url) + if existing_doc and existing_doc.metadata.get("md5") == md5_hash: + updated = False + if title is not None and title != existing_doc.title: + existing_doc.title = title + updated = True + + metadata.update({"contentType": content_type, "md5": md5_hash}) + merged_metadata = {**(existing_doc.metadata or {}), **metadata} + if merged_metadata != existing_doc.metadata: + existing_doc.metadata = merged_metadata + updated = True + + if updated: + return await client.document_repository.update(existing_doc) + return existing_doc + + file_extension = get_extension_from_content_type_or_url(url, content_type) + + if file_extension not in supported_extensions: + raise ValueError( + f"Unsupported content type/extension: {content_type}/{file_extension}" + ) + + with tempfile.NamedTemporaryFile( + mode="wb", suffix=file_extension, delete=False + ) as temp_file: + temp_file.write(response.content) + temp_file.flush() + temp_path = Path(temp_file.name) + + try: + docling_document = await client.convert(temp_path) + chunks = await client.chunk(docling_document) + embedded_chunks = await embed_chunks(chunks, client._config) + finally: + temp_path.unlink(missing_ok=True) + + metadata.update({"contentType": content_type, "md5": md5_hash}) + + stored_content = docling_document.export_to_markdown() + + if existing_doc: + existing_doc.content = stored_content + existing_doc.metadata = metadata + existing_doc.set_docling(docling_document) + if title is not None: + existing_doc.title = title + elif existing_doc.title is None: + existing_doc.title = await client._resolve_title( + docling_document, stored_content + ) + return await _update_document_with_chunks( + client, existing_doc, embedded_chunks, docling_document + ) + else: + if title is None: + title = await client._resolve_title(docling_document, stored_content) + document = Document( + content=stored_content, + uri=url, + title=title, + metadata=metadata, + ) + document.set_docling(docling_document) + return await _store_document_with_chunks( + client, document, embedded_chunks, docling_document + ) + + +async def update_document( + client: "HaikuRAG", + document_id: str, + content: str | None = None, + metadata: dict | None = None, + chunks: list[Chunk] | None = None, + title: str | None = None, + docling_document: "DoclingDocument | None" = None, +) -> Document: + """Update a document by ID. + + Updates specified fields. When content or docling_document is provided, the + document is rechunked and re-embedded. Updates to only metadata or title + skip rechunking for efficiency. + + Raises: + ValueError: If document not found, or if both content and + docling_document are provided. + """ + from haiku.rag.embeddings import embed_chunks + + # Validate: content and docling_document are mutually exclusive + if content is not None and docling_document is not None: + raise ValueError( + "content and docling_document are mutually exclusive. " + "Provide one or the other, not both." + ) + + existing_doc = await client.get_document_by_id(document_id) + if existing_doc is None: + raise ValueError(f"Document with ID {document_id} not found") + + if title is not None: + existing_doc.title = title + if metadata is not None: + existing_doc.metadata = metadata + + # Only metadata/title update - no rechunking needed + if content is None and chunks is None and docling_document is None: + return await client.document_repository.update(existing_doc) + + # Custom chunks provided - use them as-is + if chunks is not None: + if docling_document is not None: + existing_doc.content = docling_document.export_to_markdown() + existing_doc.set_docling(docling_document) + elif content is not None: + existing_doc.content = content + + return await _update_document_with_chunks( + client, existing_doc, chunks, docling_document + ) + + # DoclingDocument provided without chunks - chunk and embed + if docling_document is not None: + existing_doc.content = docling_document.export_to_markdown() + existing_doc.set_docling(docling_document) + + new_chunks = await client.chunk(docling_document) + embedded_chunks = await embed_chunks(new_chunks, client._config) + return await _update_document_with_chunks( + client, existing_doc, embedded_chunks, docling_document + ) + + # Content provided without chunks - convert, chunk, and embed + assert content is not None + existing_doc.content = content + converter = get_converter(client._config) + converted_docling = await converter.convert_text(existing_doc.content, format="md") + existing_doc.set_docling(converted_docling) + + new_chunks = await client.chunk(converted_docling) + embedded_chunks = await embed_chunks(new_chunks, client._config) + return await _update_document_with_chunks( + client, existing_doc, embedded_chunks, converted_docling + ) + + +def check_source_accessible(uri: str) -> bool: + """Check if a document's source URI is accessible.""" + parsed_url = urlparse(uri) + try: + if parsed_url.scheme == "file": + return Path(parsed_url.path).exists() + elif parsed_url.scheme in ("http", "https"): + return True + return False + except Exception: + return False diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index f9b9e8f5..955c7bd7 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -4,6 +4,7 @@ from collections.abc import AsyncGenerator from datetime import datetime from typing import TYPE_CHECKING +from haiku.rag.client.documents import check_source_accessible from haiku.rag.converters import get_converter from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document @@ -261,7 +262,7 @@ async def _rebuild_full( assert doc.id is not None # Try to rebuild from source if available - if doc.uri and client._check_source_accessible(doc.uri): + if doc.uri and check_source_accessible(doc.uri): try: # Flush pending batch before source rebuild (creates new doc) if pending_docs: diff --git a/tests/store/test_document_items.py b/tests/store/test_document_items.py index 58dfede0..b495353a 100644 --- a/tests/store/test_document_items.py +++ b/tests/store/test_document_items.py @@ -1,6 +1,10 @@ import pytest from haiku.rag.client import HaikuRAG +from haiku.rag.client.documents import ( + _store_document_with_chunks, + _update_document_with_chunks, +) from haiku.rag.store.engine import Store from haiku.rag.store.models.document_item import ( DocumentItem, @@ -217,7 +221,7 @@ class TestDocumentItemPopulation: # Use _store_document_with_chunks directly with empty chunks # to avoid needing embeddings - created = await rag._store_document_with_chunks(document, [], docling_doc) + created = await _store_document_with_chunks(rag, document, [], docling_doc) assert created.id is not None count = await rag.document_item_repository.get_item_count(created.id) @@ -245,7 +249,7 @@ class TestDocumentItemPopulation: uri="test://doc", ) document.set_docling(docling_doc) - created = await rag._store_document_with_chunks(document, [], docling_doc) + created = await _store_document_with_chunks(rag, document, [], docling_doc) assert created.id is not None assert await rag.document_item_repository.get_item_count(created.id) == 6 @@ -254,7 +258,7 @@ class TestDocumentItemPopulation: new_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Only one item now.") created.set_docling(new_doc) - await rag._update_document_with_chunks(created, [], new_doc) + await _update_document_with_chunks(rag, created, [], new_doc) assert await rag.document_item_repository.get_item_count(created.id) == 1 async def test_delete_document_cascades_items(self, temp_db_path): @@ -269,7 +273,7 @@ class TestDocumentItemPopulation: uri="test://doc", ) document.set_docling(docling_doc) - created = await rag._store_document_with_chunks(document, [], docling_doc) + created = await _store_document_with_chunks(rag, document, [], docling_doc) assert created.id is not None assert await rag.document_item_repository.get_item_count(created.id) == 6 diff --git a/tests/test_client.py b/tests/test_client.py index dbb79983..4eafd9fc 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -481,49 +481,35 @@ async def test_client_create_document_from_url_http_error(temp_db_path): ) -@pytest.mark.vcr() -async def test_get_extension_from_content_type_or_url(temp_db_path): - """Test the helper method for determining file extensions.""" - async with HaikuRAG(temp_db_path, create=True) as client: - # Test content type mappings - assert ( - client._get_extension_from_content_type_or_url("", "text/html") == ".html" - ) - assert ( - client._get_extension_from_content_type_or_url("", "application/pdf") - == ".pdf" - ) - assert ( - client._get_extension_from_content_type_or_url("", "text/plain") == ".txt" - ) +def test_get_extension_from_content_type_or_url(): + """Test the helper function for determining file extensions.""" + from haiku.rag.client.processing import get_extension_from_content_type_or_url - # Test URL extension detection - assert ( - client._get_extension_from_content_type_or_url( - "https://example.com/doc.pdf", "" - ) - == ".pdf" - ) - assert ( - client._get_extension_from_content_type_or_url( - "https://example.com/data.json", "" - ) - == ".json" - ) + # Content type mappings + assert get_extension_from_content_type_or_url("", "text/html") == ".html" + assert get_extension_from_content_type_or_url("", "application/pdf") == ".pdf" + assert get_extension_from_content_type_or_url("", "text/plain") == ".txt" - # Test default fallback - assert ( - client._get_extension_from_content_type_or_url("https://example.com/", "") - == ".html" - ) + # URL extension detection + assert ( + get_extension_from_content_type_or_url("https://example.com/doc.pdf", "") + == ".pdf" + ) + assert ( + get_extension_from_content_type_or_url("https://example.com/data.json", "") + == ".json" + ) - # Test content type priority over URL extension - assert ( - client._get_extension_from_content_type_or_url( - "https://example.com/file.txt", "application/pdf" - ) - == ".pdf" + # Default fallback + assert get_extension_from_content_type_or_url("https://example.com/", "") == ".html" + + # Content type priority over URL extension + assert ( + get_extension_from_content_type_or_url( + "https://example.com/file.txt", "application/pdf" ) + == ".pdf" + ) @pytest.mark.vcr() diff --git a/tests/test_context.py b/tests/test_context.py index 7eb5f1f0..0ad6665a 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -1,5 +1,6 @@ import pytest +from haiku.rag.client.documents import _store_document_with_chunks from haiku.rag.context import ( _expand_outward, _find_expansion_range, @@ -249,7 +250,8 @@ class TestExpandWithItems: from haiku.rag.store.models.document import Document async with HaikuRAG(temp_db_path, create=True) as rag: - doc = await rag._store_document_with_chunks( + doc = await _store_document_with_chunks( + rag, Document(content="test"), [], __import__( @@ -262,6 +264,7 @@ class TestExpandWithItems: document_id=doc.id, doc_item_refs=["#/texts/999999"], ) + assert doc.id is not None expanded = await expand_with_items( rag.document_item_repository, doc.id, [result], 5000 ) diff --git a/tests/test_context_enhancement.py b/tests/test_context_enhancement.py index 6f59ec7c..0d21cd75 100644 --- a/tests/test_context_enhancement.py +++ b/tests/test_context_enhancement.py @@ -3,6 +3,7 @@ from docling_core.types.doc.document import DoclingDocument, TableData from docling_core.types.doc.labels import DocItemLabel from haiku.rag.client import HaikuRAG +from haiku.rag.client.documents import _store_document_with_chunks from haiku.rag.config.models import AppConfig from haiku.rag.store.models import SearchResult @@ -314,7 +315,7 @@ async def test_expand_context_single_item_document(temp_db_path): async with HaikuRAG(temp_db_path, create=True) as client: document = Document(content="Simple test content") document.set_docling(docling_doc) - doc = await client._store_document_with_chunks(document, [], docling_doc) + doc = await _store_document_with_chunks(client, document, [], docling_doc) assert doc.id is not None # Create a search result with a doc_item_ref pointing to the item