From 33e6a36290a8d8b4e85b50c9e1ab7e1c2bfbfc7f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 4 Dec 2025 12:52:47 +0200 Subject: [PATCH] Make document import/update handling consistent for DoclingDocument --- CHANGELOG.md | 4 +-- docs/inspector.md | 2 +- docs/python.md | 17 +++++++++---- haiku_rag_slim/haiku/rag/client.py | 39 ++++++++++++++++++++---------- tests/test_client.py | 31 ++++++++++++++++++++++++ 5 files changed, 72 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 24a339a3..1bb35f38 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,8 +31,8 @@ - Requires terminal with image support (iTerm2, Kitty, etc.) - **New `import_document()` Method**: Import pre-processed documents with custom chunks - Use when document conversion, chunking, and embedding were done externally - - Accepts `content`, `chunks` (required), and optional `docling_document_json`/`docling_version` - - Validates that if docling parameters are provided, both must be present + - `chunks` required; `content` optional if `docling_document_json` is provided + - When `docling_document_json` is provided without `content`, content is extracted from the DoclingDocument - 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 diff --git a/docs/inspector.md b/docs/inspector.md index 389c1d2b..f5df2637 100644 --- a/docs/inspector.md +++ b/docs/inspector.md @@ -62,4 +62,4 @@ Press `v` while viewing a chunk to open the visual grounding modal: - Press `Esc` to close the modal !!! note - Visual grounding requires documents processed with docling that have page images stored. Text-only documents or documents imported without DoclingDocument won't have visual grounding available. + Visual grounding requires documents with a stored DoclingDocument that includes page images. Text-only documents or documents imported without DoclingDocument won't have visual grounding available. diff --git a/docs/python.md b/docs/python.md index 44d920b9..6a98f13f 100644 --- a/docs/python.md +++ b/docs/python.md @@ -82,20 +82,27 @@ doc = await client.import_document( ) ``` -If you also have a DoclingDocument from your processing pipeline, include it for visual grounding support: +If you also have a DoclingDocument from your processing pipeline, include it for rich metadata support (visual grounding, page numbers, section headings). When providing a DoclingDocument, you can omit `content` - it will be extracted automatically: ```python +# With explicit content +doc = await client.import_document( + chunks=chunks, + content="Full document content", + docling_document_json=docling_doc.model_dump_json(), + docling_version=docling_doc.version, +) + +# Or let content be extracted from DoclingDocument doc = await client.import_document( - content="Full document content", chunks=chunks, - uri="doc://custom", docling_document_json=docling_doc.model_dump_json(), docling_version=docling_doc.version, ) ``` !!! note - When providing `docling_document_json`, you must also provide `docling_version`. The JSON is validated to ensure it's a valid DoclingDocument. + Either `content` or `docling_document_json` must be provided. When `docling_document_json` is provided without `content`, the content is extracted from the DoclingDocument. ### Retrieving Documents @@ -175,7 +182,7 @@ await client.update_document_fields( chunks=custom_chunks ) -# Update with DoclingDocument (extracts content and rechunks) +# Update with DoclingDocument for rich metadata (extracts content and rechunks) await client.update_document_fields( document_id=doc.id, docling_document_json=docling_doc.model_dump_json(), diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 5ebaa2e8..8f22316e 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -150,8 +150,8 @@ class HaikuRAG: async def import_document( self, - content: str, chunks: list[Chunk], + content: str | None = None, uri: str | None = None, title: str | None = None, metadata: dict | None = None, @@ -164,38 +164,51 @@ class HaikuRAG: externally and you want to store the results in haiku.rag. Args: - content: The document content. - chunks: Pre-created chunks (must include embeddings). + chunks: Pre-created chunks. + content: The document content. Optional if docling_document_json is provided. uri: Optional URI identifier for the document. title: Optional title for the document. metadata: Optional metadata dictionary. - docling_document_json: Optional serialized DoclingDocument JSON. - docling_version: Optional DoclingDocument schema version. + docling_document_json: Serialized DoclingDocument JSON. If provided without + content, content is extracted from the DoclingDocument. + docling_version: DoclingDocument schema version (required with docling_document_json). Returns: The created Document instance. Raises: - ValueError: If docling_document_json is provided without docling_version - or vice versa, or if the JSON is invalid. + ValueError: If neither content nor docling_document_json is provided, + if docling_document_json is provided without docling_version, + or if the JSON is invalid. """ - # Validate docling parameters + from docling_core.types.doc.document import DoclingDocument + + # 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" ) - # Validate docling JSON parses if provided + # Validate that we have at least one content source + if content is None and docling_document_json is None: + raise ValueError("Either content or docling_document_json must be provided") + + # Parse and validate docling JSON if provided + docling_document: DoclingDocument | None = None if docling_document_json is not None: try: - from docling_core.types.doc.document import DoclingDocument - - DoclingDocument.model_validate_json(docling_document_json) + docling_document = DoclingDocument.model_validate_json( + docling_document_json + ) except Exception as e: raise ValueError(f"Invalid docling_document_json: {e}") from e + # Extract content from docling if not explicitly provided + if content is None and docling_document is not None: + content = docling_document.export_to_markdown() + document = Document( - content=content, + content=content, # type: ignore[arg-type] uri=uri, title=title, metadata=metadata or {}, diff --git a/tests/test_client.py b/tests/test_client.py index 7b7df9c4..ce78bff8 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1206,6 +1206,37 @@ async def test_client_import_document_validates_docling_params(temp_db_path): docling_version="1.0.0", ) + # Should fail if neither content nor docling_document_json provided + with pytest.raises(ValueError, match="Either content or docling_document_json"): + await client.import_document(chunks=custom_chunks) + + +@pytest.mark.asyncio +async def test_client_import_document_extracts_content_from_docling(temp_db_path): + """Test that import_document extracts content from DoclingDocument when not provided.""" + from docling_core.types.doc.document import DoclingDocument + from docling_core.types.doc.labels import DocItemLabel + + async with HaikuRAG(temp_db_path, create=True) as client: + # Create a docling document with some content + docling_doc = DoclingDocument(name="test") + docling_doc.add_text( + label=DocItemLabel.TEXT, text="Content from docling document" + ) + + custom_chunks = [Chunk(content="Chunk content", order=0)] + + # Import without content - should extract from docling + doc = await client.import_document( + chunks=custom_chunks, + docling_document_json=docling_doc.model_dump_json(), + docling_version=docling_doc.version, + ) + + assert doc.id is not None + assert "Content from docling document" in doc.content + assert doc.docling_document_json == docling_doc.model_dump_json() + @pytest.mark.asyncio async def test_client_create_document_from_file_stores_docling_json(temp_db_path):