From e95ac2e25d087d526824589d3cd80b0e43daa942 Mon Sep 17 00:00:00 2001
From: bryan davis <bryan@monkeytronics.org>
Date: Tue, 23 Jun 2026 15:55:24 -0500
Subject: [PATCH] remove redundant serialization

---
 haiku_rag_slim/haiku/rag/store/compression.py | 20 +++++++++-
 .../haiku/rag/store/models/document.py        |  4 +-
 tests/test_document.py                        | 39 +++++++++++++++++++
 3 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/haiku_rag_slim/haiku/rag/store/compression.py b/haiku_rag_slim/haiku/rag/store/compression.py
index 2202cba8..97edf1e9 100644
--- a/haiku_rag_slim/haiku/rag/store/compression.py
+++ b/haiku_rag_slim/haiku/rag/store/compression.py
@@ -35,7 +35,20 @@ def decompress_json(data: bytes) -> str:
 
 
 def compress_docling_split(json_str: str) -> tuple[bytes, bytes | None]:
-    """Split a DoclingDocument JSON into structure and pages, compress both with zstd.
+    """Parse a DoclingDocument JSON string and compress it.
+
+    Thin wrapper over :func:`compress_docling_data` for callers that only hold
+    the serialized string — store migrations and rebuild-from-blob, neither of
+    which is speed-sensitive. The ingestion hot path should call
+    ``compress_docling_data`` with ``DoclingDocument.model_dump(mode="json")``
+    instead, to avoid serializing the document to a full JSON string only to
+    parse it straight back into a dict.
+    """
+    return compress_docling_data(json.loads(json_str))
+
+
+def compress_docling_data(data: dict) -> tuple[bytes, bytes | None]:
+    """Split a DoclingDocument dict into structure and pages, compress both with zstd.
 
     Picture image URIs are stripped from the structure blob — they are stored on
     the corresponding ``document_items.picture_data`` rows and don't need to be
@@ -43,11 +56,14 @@ def compress_docling_split(json_str: str) -> tuple[bytes, bytes | None]:
     field is present, so each picture's ``image`` is set to ``None`` rather than
     partially mutated to keep the JSON re-validating cleanly.
 
+    Mutates ``data`` in place (pops ``pages``, nulls picture images); callers
+    pass a freshly built dict (``model_dump`` / ``json.loads`` output), so this
+    never touches a live DoclingDocument.
+
     Returns:
         Tuple of (structure_bytes, pages_bytes). pages_bytes is None if the
         document has no page images.
     """
-    data = json.loads(json_str)
     pages = data.pop("pages", None)
 
     for picture in data.get("pictures") or []:
diff --git a/haiku_rag_slim/haiku/rag/store/models/document.py b/haiku_rag_slim/haiku/rag/store/models/document.py
index 1f38d197..48dbf432 100644
--- a/haiku_rag_slim/haiku/rag/store/models/document.py
+++ b/haiku_rag_slim/haiku/rag/store/models/document.py
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
 
 from pydantic import BaseModel, Field
 
-from haiku.rag.store.compression import compress_docling_split, decompress_json
+from haiku.rag.store.compression import compress_docling_data, decompress_json
 
 if TYPE_CHECKING:
     from docling_core.types.doc.document import DoclingDocument, PageItem
@@ -32,7 +32,7 @@ class Document(BaseModel):
         Sets docling_document (zstd-compressed structure without pages),
         docling_pages (zstd-compressed page images), and docling_version.
         """
-        structure, pages = compress_docling_split(docling_doc.model_dump_json())
+        structure, pages = compress_docling_data(docling_doc.model_dump(mode="json"))
         self.docling_document = structure
         self.docling_pages = pages
         self.docling_version = docling_doc.version
diff --git a/tests/test_document.py b/tests/test_document.py
index eb41a9a4..317aeb46 100644
--- a/tests/test_document.py
+++ b/tests/test_document.py
@@ -281,6 +281,45 @@ def test_set_docling_with_page_images():
     assert "1" in pages
 
 
+def test_set_docling_dict_path_matches_string_path():
+    """The ingestion dict path must produce the same stored bytes as the
+    legacy string path.
+
+    ``set_docling`` feeds ``compress_docling_data`` the dict from
+    ``model_dump(mode="json")`` instead of round-tripping through
+    ``model_dump_json()`` + ``json.loads``. The on-disk format must not change
+    — in particular int-keyed ``pages`` must still serialize to string keys.
+    """
+    import json
+
+    from docling_core.types.doc.base import Size
+    from docling_core.types.doc.document import DoclingDocument, PageItem
+    from docling_core.types.doc.labels import DocItemLabel
+
+    from haiku.rag.store.compression import (
+        compress_docling_data,
+        compress_docling_split,
+        decompress_json,
+    )
+
+    docling_doc = DoclingDocument(name="equivalence_test")
+    docling_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Hello world")
+    docling_doc.pages[1] = PageItem(size=Size(width=612, height=792), page_no=1)
+
+    struct_dict, pages_dict = compress_docling_data(
+        docling_doc.model_dump(mode="json")
+    )
+    struct_str, pages_str = compress_docling_split(docling_doc.model_dump_json())
+
+    assert json.loads(decompress_json(struct_dict)) == json.loads(
+        decompress_json(struct_str)
+    )
+    assert pages_dict is not None and pages_str is not None
+    assert json.loads(decompress_json(pages_dict)) == json.loads(
+        decompress_json(pages_str)
+    )
+
+
 def test_get_page_images():
     """get_page_images returns requested pages from docling_pages blob."""
     import json