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

Store DoclingDocument json

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-11-28 10:49:17 +02:00
parent 343272d6fb
commit ad4e79ab5d
No known key found for this signature in database
2 changed files with 187 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -101,6 +101,8 @@ class HaikuRAG:
uri=uri,
title=title,
metadata=metadata or {},
docling_document_json=docling_document.model_dump_json(),
docling_version=docling_document.version,
)
return await self.document_repository._create_and_chunk(
document, docling_document, chunks
@ -125,21 +127,27 @@ class HaikuRAG:
Returns:
The created Document instance.
"""
document = Document(
content=content,
uri=uri,
title=title,
metadata=metadata or {},
)
# Only create docling_document if we need to generate chunks
if chunks is None:
# Use converter to convert text
converter = get_converter(self._config)
docling_document = await converter.convert_text(content)
docling_json = docling_document.model_dump_json()
docling_version = docling_document.version
else:
# Chunks already provided, no conversion needed
docling_document = None
docling_json = None
docling_version = None
document = Document(
content=content,
uri=uri,
title=title,
metadata=metadata or {},
docling_document_json=docling_json,
docling_version=docling_version,
)
return await self.document_repository._create_and_chunk(
document, docling_document, chunks
@ -269,6 +277,8 @@ class HaikuRAG:
# Update existing document
existing_doc.content = docling_document.export_to_markdown()
existing_doc.metadata = metadata
existing_doc.docling_document_json = docling_document.model_dump_json()
existing_doc.docling_version = docling_document.version
if title is not None:
existing_doc.title = title
return await self.document_repository._update_and_rechunk(
@ -363,6 +373,8 @@ class HaikuRAG:
if existing_doc:
existing_doc.content = docling_document.export_to_markdown()
existing_doc.metadata = metadata
existing_doc.docling_document_json = docling_document.model_dump_json()
existing_doc.docling_version = docling_document.version
if title is not None:
existing_doc.title = title
return await self.document_repository._update_and_rechunk(
@ -435,6 +447,10 @@ class HaikuRAG:
converter = get_converter(self._config)
docling_document = await converter.convert_text(document.content)
# Store DoclingDocument JSON
document.docling_document_json = docling_document.model_dump_json()
document.docling_version = docling_document.version
return await self.document_repository._update_and_rechunk(
document, docling_document
)
@ -476,7 +492,7 @@ class HaikuRAG:
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
# Use custom chunks - no docling document to store
# Delete existing chunks
await self.chunk_repository.delete_by_document_id(document_id)
@ -495,6 +511,11 @@ class HaikuRAG:
# 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
)
@ -843,12 +864,24 @@ class HaikuRAG:
"Source missing for %s, re-embedding from content", doc.uri
)
docling_document = await converter.convert_text(doc.content)
# Update document with docling JSON
doc.docling_document_json = docling_document.model_dump_json()
doc.docling_version = docling_document.version
await self.document_repository.update(doc)
await self.chunk_repository.create_chunks_for_document(
doc.id, docling_document
)
yield doc.id
else:
docling_document = await converter.convert_text(doc.content)
# Update document with docling JSON
doc.docling_document_json = docling_document.model_dump_json()
doc.docling_version = docling_document.version
await self.document_repository.update(doc)
await self.chunk_repository.create_chunks_for_document(
doc.id, docling_document
)

146
tests/test_client.py
View file

@ -1007,3 +1007,149 @@ async def test_client_expand_context_keeps_separate_non_overlapping(temp_db_path
assert "Chunk 7" in chunk5_expanded.content
assert "Chunk 0" not in chunk5_expanded.content
assert score2 == 0.7
@pytest.mark.asyncio
async def test_client_create_document_stores_docling_json(temp_db_path):
"""Test that create_document stores DoclingDocument JSON."""
async with HaikuRAG(temp_db_path) as client:
doc = await client.create_document(
content="Test content for docling storage",
uri="test://docling",
metadata={"test": "docling_storage"},
)
assert doc.id is not None
assert doc.docling_document_json is not None
assert doc.docling_version is not None
# Verify JSON is valid and can be parsed
import json
parsed = json.loads(doc.docling_document_json)
assert "version" in parsed
assert parsed["version"] == doc.docling_version
@pytest.mark.asyncio
async def test_client_create_document_with_custom_chunks_no_docling_json(temp_db_path):
"""Test that create_document with custom chunks does not store docling JSON."""
async with HaikuRAG(temp_db_path) as client:
custom_chunks = [Chunk(content="Custom chunk", order=0)]
doc = await client.create_document(content="Test content", chunks=custom_chunks)
assert doc.id is not None
# When custom chunks are provided, no conversion happens
assert doc.docling_document_json is None
assert doc.docling_version is None
@pytest.mark.asyncio
async def test_client_create_document_from_file_stores_docling_json(temp_db_path):
"""Test that create_document_from_source stores DoclingDocument JSON for files."""
async with HaikuRAG(temp_db_path) as client:
with tempfile.TemporaryDirectory() as temp_dir:
temp_path = Path(temp_dir) / "test.txt"
temp_path.write_text("Test file content")
doc = await client.create_document_from_source(temp_path)
assert isinstance(doc, Document)
assert doc.id is not None
assert doc.docling_document_json is not None
assert doc.docling_version is not None
# Verify the stored document also has the JSON
retrieved = await client.get_document_by_id(doc.id)
assert retrieved is not None
assert retrieved.docling_document_json == doc.docling_document_json
assert retrieved.docling_version == doc.docling_version
@pytest.mark.asyncio
async def test_client_update_document_stores_docling_json(temp_db_path):
"""Test that update_document stores DoclingDocument JSON."""
async with HaikuRAG(temp_db_path) as client:
# Create initial document
doc = await client.create_document(content="Initial content")
assert doc.id is not None
original_json = doc.docling_document_json
# Update the document
doc.content = "Updated content"
updated_doc = await client.update_document(doc)
assert updated_doc.docling_document_json is not None
assert updated_doc.docling_version is not None
# JSON should be different because content changed
assert updated_doc.docling_document_json != original_json
@pytest.mark.asyncio
async def test_client_update_document_fields_stores_docling_json(temp_db_path):
"""Test that update_document_fields stores DoclingDocument JSON when content changes."""
async with HaikuRAG(temp_db_path) as client:
# Create initial document
doc = await client.create_document(content="Initial content")
assert doc.id is not None
original_json = doc.docling_document_json
# Update content via update_document_fields
updated_doc = await client.update_document_fields(
document_id=doc.id, content="New content via fields update"
)
assert updated_doc.docling_document_json is not None
assert updated_doc.docling_version is not None
# JSON should be different because content changed
assert updated_doc.docling_document_json != original_json
@pytest.mark.asyncio
async def test_client_update_document_fields_with_custom_chunks_no_docling_json(
temp_db_path,
):
"""Test that update_document_fields with custom chunks does not update docling JSON."""
async with HaikuRAG(temp_db_path) as client:
# Create initial document
doc = await client.create_document(content="Initial content")
assert doc.id is not None
original_json = doc.docling_document_json
# Update with custom chunks
custom_chunks = [Chunk(content="Custom chunk", order=0)]
updated_doc = await client.update_document_fields(
document_id=doc.id, content="New content", chunks=custom_chunks
)
# Docling JSON should remain unchanged (no conversion when custom chunks provided)
assert updated_doc.docling_document_json == original_json
@pytest.mark.asyncio
async def test_client_file_update_stores_docling_json(temp_db_path):
"""Test that updating a file re-stores DoclingDocument JSON."""
async with HaikuRAG(temp_db_path) as client:
with tempfile.TemporaryDirectory() as temp_dir:
temp_path = Path(temp_dir) / "test.txt"
temp_path.write_text("Original content")
# Create initial document
doc1 = await client.create_document_from_source(temp_path)
assert isinstance(doc1, Document)
original_json = doc1.docling_document_json
original_version = doc1.docling_version
# Modify file
temp_path.write_text("Modified content")
# Update document from source
doc2 = await client.create_document_from_source(temp_path)
assert isinstance(doc2, Document)
assert doc2.id == doc1.id # Same document
# Docling JSON should be updated
assert doc2.docling_document_json is not None
assert doc2.docling_document_json != original_json
assert doc2.docling_version == original_version # Version stays same