From ad4e79ab5d37c060a22ce2a6dfd20e6223c23e2f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 28 Nov 2025 10:49:17 +0200 Subject: [PATCH] Store DoclingDocument json --- haiku_rag_slim/haiku/rag/client.py | 49 ++++++++-- tests/test_client.py | 146 +++++++++++++++++++++++++++++ 2 files changed, 187 insertions(+), 8 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index d3c15513..46b1f9fa 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -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 ) diff --git a/tests/test_client.py b/tests/test_client.py index 8d648f70..1d258aec 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -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