diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 46b1f9fa..523515b2 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -827,6 +827,12 @@ class HaikuRAG: for doc in documents: assert doc.id is not None 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_rebuild.py b/tests/test_rebuild.py index 2f749faf..fa410a70 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -10,6 +10,7 @@ async def test_rebuild_full(qa_corpus: Dataset, temp_db_path): async with HaikuRAG(temp_db_path, create=True) as client: doc = await client.create_document(content=qa_corpus["document_extracted"][0]) assert doc.id is not None + assert doc.docling_document_json is not None chunks_before = await client.chunk_repository.get_by_document_id(doc.id) assert len(chunks_before) > 0 @@ -19,6 +20,12 @@ async def test_rebuild_full(qa_corpus: Dataset, temp_db_path): assert doc.id in processed_ids + # Verify DoclingDocument JSON is preserved after rebuild + doc_after = await client.document_repository.get_by_id(doc.id) + assert doc_after is not None + assert doc_after.docling_document_json is not None + assert doc_after.docling_version is not None + chunks_after = await client.chunk_repository.get_by_document_id(doc.id) assert len(chunks_after) > 0 chunk_ids_after = {c.id for c in chunks_after} @@ -33,6 +40,7 @@ async def test_rebuild_embed_only(qa_corpus: Dataset, temp_db_path): async with HaikuRAG(temp_db_path, create=True) as client: doc = await client.create_document(content=qa_corpus["document_extracted"][0]) assert doc.id is not None + original_docling_json = doc.docling_document_json chunks_before = await client.chunk_repository.get_by_document_id(doc.id) assert len(chunks_before) > 0 @@ -46,6 +54,11 @@ async def test_rebuild_embed_only(qa_corpus: Dataset, temp_db_path): assert doc.id in processed_ids + # DoclingDocument JSON should be unchanged (embed-only doesn't touch documents) + doc_after = await client.document_repository.get_by_id(doc.id) + assert doc_after is not None + assert doc_after.docling_document_json == original_docling_json + chunks_after = await client.chunk_repository.get_by_document_id(doc.id) chunk_ids_after = {c.id for c in chunks_after} @@ -99,6 +112,7 @@ async def test_rebuild_rechunk(qa_corpus: Dataset, temp_db_path): async with HaikuRAG(temp_db_path, create=True) as client: doc = await client.create_document(content=qa_corpus["document_extracted"][0]) assert doc.id is not None + assert doc.docling_document_json is not None # Set a fake URI to simulate a document that came from a file doc.uri = "file:///nonexistent/path.txt" @@ -115,10 +129,12 @@ async def test_rebuild_rechunk(qa_corpus: Dataset, temp_db_path): assert doc.id in processed_ids - # Document content should be unchanged + # Document content should be unchanged, but docling JSON should be updated doc_after = await client.document_repository.get_by_id(doc.id) assert doc_after is not None assert doc_after.content == content_before + assert doc_after.docling_document_json is not None + assert doc_after.docling_version is not None chunks_after = await client.chunk_repository.get_by_document_id(doc.id) assert len(chunks_after) > 0