diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 48592048..fe8fd11f 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -17,6 +17,7 @@ import httpx from haiku.rag.config import AppConfig, Config from haiku.rag.converters import get_converter from haiku.rag.reranking import get_reranker +from haiku.rag.store.compression import compress_json from haiku.rag.store.engine import Store from haiku.rag.store.models.chunk import Chunk, SearchResult from haiku.rag.store.models.document import Document @@ -372,7 +373,7 @@ class HaikuRAG: embedded_chunks = await embed_chunks(chunks, self._config) # Store markdown export as content for better display/readability - # The original content is preserved in docling_document_json + # The original content is preserved in docling_document stored_content = docling_document.export_to_markdown() # Create document model @@ -381,7 +382,7 @@ class HaikuRAG: uri=uri, title=title, metadata=metadata or {}, - docling_document_json=docling_document.model_dump_json(), + docling_document=compress_json(docling_document.model_dump_json()), docling_version=docling_document.version, ) @@ -417,7 +418,7 @@ class HaikuRAG: uri=uri, title=title, metadata=metadata or {}, - docling_document_json=docling_document.model_dump_json(), + docling_document=compress_json(docling_document.model_dump_json()), docling_version=docling_document.version, ) @@ -550,7 +551,9 @@ class HaikuRAG: # Update existing document and rechunk 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_document = compress_json( + docling_document.model_dump_json() + ) existing_doc.docling_version = docling_document.version if title is not None: existing_doc.title = title @@ -564,7 +567,7 @@ class HaikuRAG: uri=uri, title=title, metadata=metadata, - docling_document_json=docling_document.model_dump_json(), + docling_document=compress_json(docling_document.model_dump_json()), docling_version=docling_document.version, ) return await self._store_document_with_chunks(document, embedded_chunks) @@ -657,7 +660,9 @@ class HaikuRAG: # Update existing document and rechunk 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_document = compress_json( + docling_document.model_dump_json() + ) existing_doc.docling_version = docling_document.version if title is not None: existing_doc.title = title @@ -671,7 +676,7 @@ class HaikuRAG: uri=url, title=title, metadata=metadata, - docling_document_json=docling_document.model_dump_json(), + docling_document=compress_json(docling_document.model_dump_json()), docling_version=docling_document.version, ) return await self._store_document_with_chunks(document, embedded_chunks) @@ -788,7 +793,9 @@ class HaikuRAG: # Store docling data if provided if docling_document is not None: existing_doc.content = docling_document.export_to_markdown() - existing_doc.docling_document_json = docling_document.model_dump_json() + existing_doc.docling_document = compress_json( + docling_document.model_dump_json() + ) existing_doc.docling_version = docling_document.version elif content is not None: existing_doc.content = content @@ -798,7 +805,9 @@ class HaikuRAG: # DoclingDocument provided without chunks - chunk and embed using primitives if docling_document is not None: existing_doc.content = docling_document.export_to_markdown() - existing_doc.docling_document_json = docling_document.model_dump_json() + existing_doc.docling_document = compress_json( + docling_document.model_dump_json() + ) existing_doc.docling_version = docling_document.version new_chunks = await self.chunk(docling_document) @@ -810,7 +819,9 @@ class HaikuRAG: # Content provided without chunks - convert, chunk, and embed using primitives existing_doc.content = content # type: ignore[assignment] converted_docling = await self.convert(existing_doc.content) - existing_doc.docling_document_json = converted_docling.model_dump_json() + existing_doc.docling_document = compress_json( + converted_docling.model_dump_json() + ) existing_doc.docling_version = converted_docling.version new_chunks = await self.chunk(converted_docling) @@ -1485,7 +1496,7 @@ class HaikuRAG: uri=doc.uri, title=doc.title, metadata=json.dumps(doc.metadata), - docling_document_json=doc.docling_document_json, + docling_document=doc.docling_document, docling_version=doc.docling_version, created_at=doc.created_at.isoformat() if doc.created_at else now, updated_at=now, @@ -1523,7 +1534,7 @@ class HaikuRAG: embedded_chunks = await embed_chunks(chunks, self._config) # Update document fields - doc.docling_document_json = docling_document.model_dump_json() + doc.docling_document = compress_json(docling_document.model_dump_json()) doc.docling_version = docling_document.version # Prepare chunks with document_id and order @@ -1602,7 +1613,7 @@ class HaikuRAG: chunks = await self.chunk(docling_document) embedded_chunks = await embed_chunks(chunks, self._config) - doc.docling_document_json = docling_document.model_dump_json() + doc.docling_document = compress_json(docling_document.model_dump_json()) doc.docling_version = docling_document.version # Prepare chunks with document_id and order diff --git a/tests/test_client.py b/tests/test_client.py index bfe27c85..b8a8e877 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -8,6 +8,7 @@ from datasets import Dataset from haiku.rag.client import HaikuRAG from haiku.rag.config import Config +from haiku.rag.store.compression import decompress_json from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document @@ -790,13 +791,15 @@ async def test_client_create_document_stores_docling_json(temp_db_path): ) assert doc.id is not None - assert doc.docling_document_json is not None + assert doc.docling_document 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) + from haiku.rag.store.compression import decompress_json + + parsed = json.loads(decompress_json(doc.docling_document)) assert "version" in parsed assert parsed["version"] == doc.docling_version @@ -824,7 +827,8 @@ async def test_client_import_document_stores_docling_data(temp_db_path): assert doc.id is not None assert "Content from docling document" in doc.content - assert doc.docling_document_json == docling_doc.model_dump_json() + assert doc.docling_document is not None + assert decompress_json(doc.docling_document) == docling_doc.model_dump_json() assert doc.docling_version == docling_doc.version @@ -840,13 +844,13 @@ async def test_client_create_document_from_file_stores_docling_json(temp_db_path assert isinstance(doc, Document) assert doc.id is not None - assert doc.docling_document_json is not None + assert doc.docling_document 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_document == doc.docling_document assert retrieved.docling_version == doc.docling_version @@ -857,17 +861,17 @@ async def test_client_update_document_stores_docling_json(temp_db_path): # Create initial document doc = await client.create_document(content="Initial content") assert doc.id is not None - original_json = doc.docling_document_json + original_json = doc.docling_document # Update content via update_document updated_doc = await client.update_document( document_id=doc.id, content="New content via fields update" ) - assert updated_doc.docling_document_json is not None + assert updated_doc.docling_document 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 + assert updated_doc.docling_document != original_json @pytest.mark.vcr() @@ -879,7 +883,7 @@ async def test_client_update_document_with_custom_chunks_no_docling_json( # Create initial document doc = await client.create_document(content="Initial content") assert doc.id is not None - original_json = doc.docling_document_json + original_json = doc.docling_document # Update with custom chunks custom_chunks = [Chunk(content="Custom chunk", order=0)] @@ -888,7 +892,7 @@ async def test_client_update_document_with_custom_chunks_no_docling_json( ) # Docling JSON should remain unchanged (no conversion when custom chunks provided) - assert updated_doc.docling_document_json == original_json + assert updated_doc.docling_document == original_json @pytest.mark.vcr() @@ -942,7 +946,11 @@ async def test_client_update_document_with_docling_rechunks(temp_db_path): # 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_document is not None + assert ( + decompress_json(updated_doc.docling_document) + == docling_doc.model_dump_json() + ) assert updated_doc.docling_version == docling_doc.version # Chunks should be regenerated @@ -981,7 +989,11 @@ async def test_client_update_document_docling_with_chunks(temp_db_path): # 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() + assert updated_doc.docling_document is not None + assert ( + decompress_json(updated_doc.docling_document) + == 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) @@ -1001,7 +1013,7 @@ async def test_client_file_update_stores_docling_json(temp_db_path): # Create initial document doc1 = await client.create_document_from_source(temp_path) assert isinstance(doc1, Document) - original_json = doc1.docling_document_json + original_json = doc1.docling_document original_version = doc1.docling_version # Modify file @@ -1013,8 +1025,8 @@ async def test_client_file_update_stores_docling_json(temp_db_path): 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_document is not None + assert doc2.docling_document != original_json assert doc2.docling_version == original_version # Version stays same @@ -1038,7 +1050,7 @@ async def test_client_visualize_chunk_no_bounding_boxes(temp_db_path): ) assert doc.id is not None - assert doc.docling_document_json is not None + assert doc.docling_document is not None chunks = await client.chunk_repository.get_by_document_id(doc.id) assert len(chunks) >= 1 @@ -1095,7 +1107,7 @@ async def test_client_visualize_chunk_with_pdf(temp_db_path): doc = await client.create_document_from_source(pdf_path) assert isinstance(doc, Document) assert doc.id is not None - assert doc.docling_document_json is not None + assert doc.docling_document is not None chunks = await client.chunk_repository.get_by_document_id(doc.id) assert len(chunks) > 0 @@ -1345,7 +1357,7 @@ async def test_client_create_document_with_html_format(temp_db_path): ) assert doc.id is not None - assert doc.docling_document_json is not None + assert doc.docling_document is not None # Verify the DoclingDocument has proper structure docling_doc = doc.get_docling_document() diff --git a/tests/test_context_enhancement.py b/tests/test_context_enhancement.py index 8ec6abe8..54f48f54 100644 --- a/tests/test_context_enhancement.py +++ b/tests/test_context_enhancement.py @@ -613,7 +613,7 @@ This is paragraph four about topic C. ) assert doc.id is not None - assert doc.docling_document_json is not None + assert doc.docling_document is not None # Get chunks which should have doc_item_refs chunks = await client.chunk_repository.get_by_document_id(doc.id) diff --git a/tests/test_converters.py b/tests/test_converters.py index 1e832c70..f8281f60 100644 --- a/tests/test_converters.py +++ b/tests/test_converters.py @@ -20,7 +20,7 @@ def vcr_cassette_dir(): return str(Path(__file__).parent / "cassettes" / "test_converters") -def create_mock_docling_document_json(name: str = "test") -> dict: +def create_mock_docling_document(name: str = "test") -> dict: """Create a minimal valid DoclingDocument JSON structure for mocking.""" return { "schema_name": "DoclingDocument", @@ -476,7 +476,7 @@ class TestDoclingServeConverter: @pytest.mark.asyncio async def test_convert_text_success(self, converter): """Test successful text conversion via docling-serve async workflow.""" - doc_json = create_mock_docling_document_json("test") + doc_json = create_mock_docling_document("test") submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json) with patch("httpx.AsyncClient") as mock_client_class: @@ -498,7 +498,7 @@ class TestDoclingServeConverter: config.providers.docling_serve.api_key = "test-key" converter = DoclingServeConverter(config) - doc_json = create_mock_docling_document_json("test") + doc_json = create_mock_docling_document("test") submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json) with patch("httpx.AsyncClient") as mock_client_class: @@ -527,7 +527,7 @@ class TestDoclingServeConverter: config.processing.conversion_options.images_scale = 3.0 converter = DoclingServeConverter(config) - doc_json = create_mock_docling_document_json("test") + doc_json = create_mock_docling_document("test") submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json) with patch("httpx.AsyncClient") as mock_client_class: @@ -622,7 +622,7 @@ class TestDoclingServeConverter: @pytest.mark.asyncio async def test_convert_file_pdf(self, converter): """Test converting PDF file via docling-serve async workflow.""" - doc_json = create_mock_docling_document_json("test") + doc_json = create_mock_docling_document("test") submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json) with patch("httpx.AsyncClient") as mock_client_class: @@ -645,7 +645,7 @@ class TestDoclingServeConverter: @pytest.mark.asyncio async def test_convert_file_text(self, converter): """Test converting text file (reads locally, sends to docling-serve).""" - doc_json = create_mock_docling_document_json("test") + doc_json = create_mock_docling_document("test") submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json) with patch("httpx.AsyncClient") as mock_client_class: @@ -734,7 +734,7 @@ class TestDoclingServeConverterPictureDescription: config.prompts.picture_description = "Test prompt for picture description" converter = DoclingServeConverter(config) - doc_json = create_mock_docling_document_json("test") + doc_json = create_mock_docling_document("test") submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json) with patch("httpx.AsyncClient") as mock_client_class: @@ -767,7 +767,7 @@ class TestDoclingServeConverterPictureDescription: """Test that picture description is disabled by default.""" converter = DoclingServeConverter(config) - doc_json = create_mock_docling_document_json("test") + doc_json = create_mock_docling_document("test") submit_resp, poll_resp, result_resp = create_async_workflow_mocks(doc_json) with patch("httpx.AsyncClient") as mock_client_class: diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index b9737d6b..912a5fc1 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -10,7 +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 + assert doc.docling_document is not None chunks_before = await client.chunk_repository.get_by_document_id(doc.id) assert len(chunks_before) > 0 @@ -23,7 +23,7 @@ async def test_rebuild_full(qa_corpus: Dataset, temp_db_path): # 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_document is not None assert doc_after.docling_version is not None chunks_after = await client.chunk_repository.get_by_document_id(doc.id) @@ -40,7 +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 + original_docling_json = doc.docling_document chunks_before = await client.chunk_repository.get_by_document_id(doc.id) assert len(chunks_before) > 0 @@ -57,7 +57,7 @@ async def test_rebuild_embed_only(qa_corpus: Dataset, temp_db_path): # 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 + assert doc_after.docling_document == 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} @@ -112,7 +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 + assert doc.docling_document is not None # Set a fake URI to simulate a document that came from a file doc.uri = "file:///nonexistent/path.txt" @@ -133,7 +133,7 @@ async def test_rebuild_rechunk(qa_corpus: Dataset, temp_db_path): 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_document is not None assert doc_after.docling_version is not None chunks_after = await client.chunk_repository.get_by_document_id(doc.id)