Make document import/update handling consistent for DoclingDocument

This commit is contained in:
Yiorgis Gozadinos 2025-12-04 12:52:47 +02:00
parent 60ab864fc1
commit 33e6a36290
No known key found for this signature in database
5 changed files with 72 additions and 21 deletions

View file

@ -31,8 +31,8 @@
- Requires terminal with image support (iTerm2, Kitty, etc.) - Requires terminal with image support (iTerm2, Kitty, etc.)
- **New `import_document()` Method**: Import pre-processed documents with custom chunks - **New `import_document()` Method**: Import pre-processed documents with custom chunks
- Use when document conversion, chunking, and embedding were done externally - Use when document conversion, chunking, and embedding were done externally
- Accepts `content`, `chunks` (required), and optional `docling_document_json`/`docling_version` - `chunks` required; `content` optional if `docling_document_json` is provided
- Validates that if docling parameters are provided, both must be present - When `docling_document_json` is provided without `content`, content is extracted from the DoclingDocument
- Validates docling JSON parses correctly if provided - Validates docling JSON parses correctly if provided
- **`update_document_fields()` DoclingDocument Support**: Added `docling_document_json` and `docling_version` parameters - **`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 without `chunks`, content is extracted and document is rechunked

View file

@ -62,4 +62,4 @@ Press `v` while viewing a chunk to open the visual grounding modal:
- Press `Esc` to close the modal - Press `Esc` to close the modal
!!! note !!! 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.

View file

@ -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 ```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( doc = await client.import_document(
content="Full document content",
chunks=chunks, chunks=chunks,
uri="doc://custom",
docling_document_json=docling_doc.model_dump_json(), docling_document_json=docling_doc.model_dump_json(),
docling_version=docling_doc.version, docling_version=docling_doc.version,
) )
``` ```
!!! note !!! 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 ### Retrieving Documents
@ -175,7 +182,7 @@ await client.update_document_fields(
chunks=custom_chunks 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( await client.update_document_fields(
document_id=doc.id, document_id=doc.id,
docling_document_json=docling_doc.model_dump_json(), docling_document_json=docling_doc.model_dump_json(),

View file

@ -150,8 +150,8 @@ class HaikuRAG:
async def import_document( async def import_document(
self, self,
content: str,
chunks: list[Chunk], chunks: list[Chunk],
content: str | None = None,
uri: str | None = None, uri: str | None = None,
title: str | None = None, title: str | None = None,
metadata: dict | None = None, metadata: dict | None = None,
@ -164,38 +164,51 @@ class HaikuRAG:
externally and you want to store the results in haiku.rag. externally and you want to store the results in haiku.rag.
Args: Args:
content: The document content. chunks: Pre-created chunks.
chunks: Pre-created chunks (must include embeddings). content: The document content. Optional if docling_document_json is provided.
uri: Optional URI identifier for the document. uri: Optional URI identifier for the document.
title: Optional title for the document. title: Optional title for the document.
metadata: Optional metadata dictionary. metadata: Optional metadata dictionary.
docling_document_json: Optional serialized DoclingDocument JSON. docling_document_json: Serialized DoclingDocument JSON. If provided without
docling_version: Optional DoclingDocument schema version. content, content is extracted from the DoclingDocument.
docling_version: DoclingDocument schema version (required with docling_document_json).
Returns: Returns:
The created Document instance. The created Document instance.
Raises: Raises:
ValueError: If docling_document_json is provided without docling_version ValueError: If neither content nor docling_document_json is provided,
or vice versa, or if the JSON is invalid. 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): if (docling_document_json is None) != (docling_version is None):
raise ValueError( raise ValueError(
"docling_document_json and docling_version must both be provided or both be None" "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: if docling_document_json is not None:
try: try:
from docling_core.types.doc.document import DoclingDocument docling_document = DoclingDocument.model_validate_json(
docling_document_json
DoclingDocument.model_validate_json(docling_document_json) )
except Exception as e: except Exception as e:
raise ValueError(f"Invalid docling_document_json: {e}") from 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( document = Document(
content=content, content=content, # type: ignore[arg-type]
uri=uri, uri=uri,
title=title, title=title,
metadata=metadata or {}, metadata=metadata or {},

View file

@ -1206,6 +1206,37 @@ async def test_client_import_document_validates_docling_params(temp_db_path):
docling_version="1.0.0", 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 @pytest.mark.asyncio
async def test_client_create_document_from_file_stores_docling_json(temp_db_path): async def test_client_create_document_from_file_stores_docling_json(temp_db_path):