danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add docling_document_json to update_document_fields() and

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-12-04 12:42:10 +02:00
parent 8e30f67266
commit 60ab864fc1
No known key found for this signature in database
4 changed files with 213 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
CHANGELOG.md
View file

@ -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

21
docs/python.md
View file

@ -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

121
haiku_rag_slim/haiku/rag/client.py
View file

@ -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)

105
tests/test_client.py
View file

@ -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."""