From 60ab864fc14036f6cbe631d89cbe8ce1788e7bd6 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 4 Dec 2025 12:42:10 +0200 Subject: [PATCH] Add docling_document_json to update_document_fields() and --- CHANGELOG.md | 4 + docs/python.md | 21 ++++- haiku_rag_slim/haiku/rag/client.py | 121 ++++++++++++++++++++--------- tests/test_client.py | 105 +++++++++++++++++++++++++ 4 files changed, 213 insertions(+), 38 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index abe6a9bc..24a339a3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -34,6 +34,10 @@ - Accepts `content`, `chunks` (required), and optional `docling_document_json`/`docling_version` - Validates that if docling parameters are provided, both must be present - Validates docling JSON parses correctly if provided +- **`update_document_fields()` DoclingDocument Support**: Added `docling_document_json` and `docling_version` parameters + - When `docling_document_json` is provided without `chunks`, content is extracted and document is rechunked + - When `docling_document_json` is provided with `chunks`, both are stored (chunks used as-is) + - `content` and `docling_document_json` are mutually exclusive to avoid ambiguity ### Changed diff --git a/docs/python.md b/docs/python.md index 59006257..44d920b9 100644 --- a/docs/python.md +++ b/docs/python.md @@ -174,9 +174,28 @@ await client.update_document_fields( document_id=doc.id, chunks=custom_chunks ) + +# Update with DoclingDocument (extracts content and rechunks) +await client.update_document_fields( + document_id=doc.id, + docling_document_json=docling_doc.model_dump_json(), + docling_version=docling_doc.version, +) + +# Update with DoclingDocument and custom chunks (stores both, uses chunks as-is) +await client.update_document_fields( + document_id=doc.id, + chunks=custom_chunks, + docling_document_json=docling_doc.model_dump_json(), + docling_version=docling_doc.version, +) ``` -**Performance Note:** Updates to only `metadata` or `title` skip re-chunking for efficiency. Updates to `content` or `chunks` will regenerate or replace the document's chunks. +**Notes:** + +- Updates to only `metadata` or `title` skip re-chunking for efficiency +- Updates to `content`, `chunks`, or `docling_document_json` will regenerate or replace chunks +- `content` and `docling_document_json` are mutually exclusive - provide one or the other ### Deleting Documents diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 3e9046f4..5ebaa2e8 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -514,67 +514,114 @@ class HaikuRAG: metadata: dict | None = None, chunks: list[Chunk] | None = None, title: str | None = None, + docling_document_json: str | None = None, + docling_version: str | None = None, ) -> Document: """Update specific fields of a document by ID. Args: document_id: The ID of the document to update - content: New content for the document + content: New content for the document (mutually exclusive with docling_document_json) metadata: New metadata for the document chunks: Custom chunks to use instead of auto-generating title: New title for the document + docling_document_json: Serialized DoclingDocument JSON (mutually exclusive with content) + docling_version: DoclingDocument schema version (required with docling_document_json) Returns: The updated Document instance. + + Raises: + ValueError: If both content and docling_document_json are provided, + if docling_document_json is provided without docling_version, + or if the JSON is invalid. """ + from docling_core.types.doc.document import DoclingDocument + + # Validate: content and docling_document_json are mutually exclusive + if content is not None and docling_document_json is not None: + raise ValueError( + "content and docling_document_json are mutually exclusive. " + "Provide one or the other, not both." + ) + + # Validate docling parameters must be provided together + if (docling_document_json is None) != (docling_version is None): + raise ValueError( + "docling_document_json and docling_version must both be provided or both be None" + ) + + # Parse and validate docling JSON if provided + docling_document: DoclingDocument | None = None + if docling_document_json is not None: + try: + docling_document = DoclingDocument.model_validate_json( + docling_document_json + ) + except Exception as e: + raise ValueError(f"Invalid docling_document_json: {e}") from e + # 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 only the provided fields - if content is not None: - existing_doc.content = content + # Update metadata/title fields if title is not None: existing_doc.title = title if metadata is not None: existing_doc.metadata = metadata - # Determine if we need to rechunk - if content is not None or chunks is not None: - # Content changed or custom chunks provided - need to rechunk - if chunks is not None: - # Use custom chunks - no docling document to store - # Delete existing chunks - await self.chunk_repository.delete_by_document_id(document_id) - - # Update document metadata - await self.document_repository.update(existing_doc) - - # Set document_id and order for all chunks - for order, chunk in enumerate(chunks): - chunk.document_id = document_id - chunk.order = order - # Batch create all chunks in a single operation - await self.chunk_repository.create(chunks) - - return existing_doc - else: - # Auto-generate chunks from content - converter = get_converter(self._config) - docling_document = await converter.convert_text(existing_doc.content) - - # Store DoclingDocument JSON - existing_doc.docling_document_json = docling_document.model_dump_json() - existing_doc.docling_version = docling_document.version - - return await self.document_repository._update_and_rechunk( - existing_doc, docling_document - ) - else: - # Only metadata/title changed - no rechunking needed + # 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: + # Update content field if provided + if content is not None: + existing_doc.content = content + + # Store docling data if provided + if docling_document is not None: + existing_doc.docling_document_json = docling_document_json + existing_doc.docling_version = docling_version + # Extract content from docling if not explicitly provided + if content is None: + existing_doc.content = docling_document.export_to_markdown() + + # Delete existing chunks and use custom ones + await self.chunk_repository.delete_by_document_id(document_id) + await self.document_repository.update(existing_doc) + + for order, chunk in enumerate(chunks): + chunk.document_id = document_id + chunk.order = order + await self.chunk_repository.create(chunks) + + return existing_doc + + # DoclingDocument provided without chunks - extract content and rechunk + if docling_document is not None: + existing_doc.content = docling_document.export_to_markdown() + existing_doc.docling_document_json = docling_document_json + existing_doc.docling_version = docling_version + + return await self.document_repository._update_and_rechunk( + existing_doc, docling_document + ) + + # Content provided without chunks - convert and rechunk + existing_doc.content = content # type: ignore[assignment] + converter = get_converter(self._config) + converted_docling = await converter.convert_text(existing_doc.content) + existing_doc.docling_document_json = converted_docling.model_dump_json() + existing_doc.docling_version = converted_docling.version + + return await self.document_repository._update_and_rechunk( + existing_doc, 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) diff --git a/tests/test_client.py b/tests/test_client.py index eb825173..7b7df9c4 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1289,6 +1289,111 @@ async def test_client_update_document_fields_with_custom_chunks_no_docling_json( assert updated_doc.docling_document_json == original_json +@pytest.mark.asyncio +async def test_client_update_document_fields_content_docling_mutually_exclusive( + temp_db_path, +): + """Test that content and docling_document_json cannot both be provided.""" + from docling_core.types.doc.document import DoclingDocument + + async with HaikuRAG(temp_db_path, create=True) as client: + doc = await client.create_document(content="Initial content") + assert doc.id is not None + + # Create a docling document + from docling_core.types.doc.labels import DocItemLabel + + docling_doc = DoclingDocument(name="test") + docling_doc.add_text(label=DocItemLabel.TEXT, text="Some text") + + with pytest.raises(ValueError, match="mutually exclusive"): + await client.update_document_fields( + document_id=doc.id, + content="New content", + docling_document_json=docling_doc.model_dump_json(), + docling_version=docling_doc.version, + ) + + +@pytest.mark.asyncio +async def test_client_update_document_fields_with_docling_rechunks(temp_db_path): + """Test that providing docling_document_json without chunks triggers rechunk.""" + from docling_core.types.doc.document import DoclingDocument + + async with HaikuRAG(temp_db_path, create=True) as client: + # Create initial document + doc = await client.create_document(content="Initial content") + assert doc.id is not None + original_chunks = await client.chunk_repository.get_by_document_id(doc.id) + + # Create a new docling document with different content + from docling_core.types.doc.labels import DocItemLabel + + docling_doc = DoclingDocument(name="updated") + docling_doc.add_text( + label=DocItemLabel.TEXT, + text="Completely different text from docling document", + ) + + # Update with docling document only - should rechunk from it + updated_doc = await client.update_document_fields( + document_id=doc.id, + docling_document_json=docling_doc.model_dump_json(), + docling_version=docling_doc.version, + ) + + # Content should be extracted from docling document + assert "Completely different text" in updated_doc.content + assert updated_doc.docling_document_json == docling_doc.model_dump_json() + assert updated_doc.docling_version == docling_doc.version + + # Chunks should be regenerated + new_chunks = await client.chunk_repository.get_by_document_id(doc.id) + assert len(new_chunks) > 0 + # Content should differ from original + assert new_chunks[0].content != original_chunks[0].content + + +@pytest.mark.asyncio +async def test_client_update_document_fields_docling_with_chunks(temp_db_path): + """Test that providing both docling_document_json and chunks stores both.""" + from docling_core.types.doc.document import DoclingDocument + + async with HaikuRAG(temp_db_path, create=True) as client: + # Create initial document + doc = await client.create_document(content="Initial content") + assert doc.id is not None + + # Create a docling document + from docling_core.types.doc.labels import DocItemLabel + + docling_doc = DoclingDocument(name="custom") + docling_doc.add_text(label=DocItemLabel.TEXT, text="Text from docling") + + # Provide both docling and custom chunks + custom_chunks = [ + Chunk(content="Custom chunk 1", order=0), + Chunk(content="Custom chunk 2", order=1), + ] + + updated_doc = await client.update_document_fields( + document_id=doc.id, + chunks=custom_chunks, + docling_document_json=docling_doc.model_dump_json(), + docling_version=docling_doc.version, + ) + + # Content should be extracted from docling (since content wasn't provided) + assert "Text from docling" in updated_doc.content + assert updated_doc.docling_document_json == docling_doc.model_dump_json() + + # Custom chunks should be used (not rechunked from docling) + chunks = await client.chunk_repository.get_by_document_id(doc.id) + assert len(chunks) == 2 + assert chunks[0].content == "Custom chunk 1" + assert chunks[1].content == "Custom chunk 2" + + @pytest.mark.asyncio async def test_client_file_update_stores_docling_json(temp_db_path): """Test that updating a file re-stores DoclingDocument JSON."""