From a4839c0ad2b02d51aa746de4951e8f727c3d8133 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 6 Aug 2026 11:02:21 +0300
Subject: [PATCH 1/4] Score retrieval evals from search results

---
 evaluations/evaluations/benchmark.py | 11 +---
 evaluations/tests/test_benchmark.py  | 87 ++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 9 deletions(-)

diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py
index 2541ca05..1249660c 100644
--- a/evaluations/evaluations/benchmark.py
+++ b/evaluations/evaluations/benchmark.py
@@ -230,20 +230,13 @@ async def run_retrieval_benchmark(
     async with HaikuRAG(db, config=config, read_only=True) as rag:
 
         async def retrieval_target(question: str) -> list[str]:
-            chunks = await rag.search(query=question, limit=5)
+            chunks = await rag.search(query=question, limit=5, include_images=False)
 
             seen = set()
             identifiers = []
             for result in chunks:
-                if result.document_id is None:
-                    continue
-                doc = await rag.get_document_by_id(result.document_id)
-                if doc is None:
-                    continue
                 # Use arxiv_id from metadata if present, otherwise use URI
-                doc_id = doc.metadata.get("arxiv_id") if doc.metadata else None
-                if doc_id is None:
-                    doc_id = doc.uri
+                doc_id = result.document_meta.get("arxiv_id") or result.document_uri
                 if doc_id and doc_id not in seen:
                     identifiers.append(doc_id)
                     seen.add(doc_id)
diff --git a/evaluations/tests/test_benchmark.py b/evaluations/tests/test_benchmark.py
index dcdc27da..782a9c5e 100644
--- a/evaluations/tests/test_benchmark.py
+++ b/evaluations/tests/test_benchmark.py
@@ -455,6 +455,93 @@ class TestLoadCaseIds:
         assert _load_case_ids(None) is None
 
 
+class TestRetrievalTarget:
+    def _spec(self) -> DatasetSpec:
+        from evaluations.config import RetrievalSample
+        from evaluations.evaluators import MAPEvaluator
+
+        return DatasetSpec(
+            key="test",
+            db_filename="test.lancedb",
+            document_loader=lambda: None,  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+            document_mapper=lambda doc: None,
+            qa_loader=lambda: [],  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+            qa_case_builder=lambda idx, doc: None,  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+            retrieval_loader=lambda: [  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+                {"q": "What is X?", "uris": ("uri-x",)},
+            ],
+            retrieval_mapper=lambda d: RetrievalSample(
+                question=d["q"], expected_uris=d["uris"]
+            ),
+            retrieval_evaluator=MAPEvaluator(),
+        )
+
+    @pytest.mark.asyncio
+    async def test_scores_from_search_results_without_reading_documents(
+        self, tmp_path: Path
+    ) -> None:
+        from haiku.rag.store.models.chunk import SearchResult
+
+        from evaluations.benchmark import run_retrieval_benchmark
+
+        searches: list[dict] = []
+
+        class FakeRag:
+            async def search(self, **kwargs) -> list[SearchResult]:
+                searches.append(kwargs)
+                return [
+                    SearchResult(
+                        content="x",
+                        score=1.0,
+                        document_id="doc-1",
+                        document_uri="uri-x",
+                    )
+                ]
+
+            async def get_document_by_id(self, document_id: str) -> None:
+                raise AssertionError(
+                    "retrieval scoring must not read whole document rows"
+                )
+
+        fake = FakeRag()
+        with patch("evaluations.benchmark.HaikuRAG") as mock_haiku:
+            mock_haiku.return_value.__aenter__.return_value = fake
+            result = await run_retrieval_benchmark(
+                self._spec(), AppConfig(), db_path=tmp_path / "test.lancedb"
+            )
+
+        assert result is not None
+        assert result["map"] == 1.0
+        assert searches[0]["include_images"] is False
+
+    @pytest.mark.asyncio
+    async def test_prefers_metadata_identifier_over_uri(self, tmp_path: Path) -> None:
+        from haiku.rag.store.models.chunk import SearchResult
+
+        from evaluations.benchmark import run_retrieval_benchmark
+
+        class FakeRag:
+            async def search(self, **kwargs) -> list[SearchResult]:
+                return [
+                    SearchResult(
+                        content="x",
+                        score=1.0,
+                        document_id="doc-1",
+                        document_uri="uri-other",
+                        document_meta={"arxiv_id": "uri-x"},
+                    )
+                ]
+
+        with patch("evaluations.benchmark.HaikuRAG") as mock_haiku:
+            mock_haiku.return_value.__aenter__.return_value = FakeRag()
+            result = await run_retrieval_benchmark(
+                self._spec(), AppConfig(), db_path=tmp_path / "test.lancedb"
+            )
+
+        assert result is not None
+        assert result["map"] == 1.0
+
+
 class TestEvaluateDatasetCaseIds:
     def _spec(self) -> DatasetSpec:
         return DatasetSpec(

From 788a2fe731d4f5e35f64277fb64cb4e645fc6358 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 6 Aug 2026 11:06:13 +0300
Subject: [PATCH 2/4] Score retrieval evals by document URI

---
 evaluations/evaluations/benchmark.py |  9 ++++-----
 evaluations/tests/test_benchmark.py  | 19 ++++++++++---------
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py
index 1249660c..a550bc9f 100644
--- a/evaluations/evaluations/benchmark.py
+++ b/evaluations/evaluations/benchmark.py
@@ -235,11 +235,10 @@ async def run_retrieval_benchmark(
             seen = set()
             identifiers = []
             for result in chunks:
-                # Use arxiv_id from metadata if present, otherwise use URI
-                doc_id = result.document_meta.get("arxiv_id") or result.document_uri
-                if doc_id and doc_id not in seen:
-                    identifiers.append(doc_id)
-                    seen.add(doc_id)
+                uri = result.document_uri
+                if uri and uri not in seen:
+                    identifiers.append(uri)
+                    seen.add(uri)
 
             return identifiers
 
diff --git a/evaluations/tests/test_benchmark.py b/evaluations/tests/test_benchmark.py
index 782a9c5e..2e49a836 100644
--- a/evaluations/tests/test_benchmark.py
+++ b/evaluations/tests/test_benchmark.py
@@ -515,21 +515,21 @@ class TestRetrievalTarget:
         assert searches[0]["include_images"] is False
 
     @pytest.mark.asyncio
-    async def test_prefers_metadata_identifier_over_uri(self, tmp_path: Path) -> None:
+    async def test_ranks_each_document_once(self, tmp_path: Path) -> None:
         from haiku.rag.store.models.chunk import SearchResult
 
         from evaluations.benchmark import run_retrieval_benchmark
 
+        def _result(uri: str, score: float) -> SearchResult:
+            return SearchResult(content="x", score=score, document_uri=uri)
+
         class FakeRag:
             async def search(self, **kwargs) -> list[SearchResult]:
                 return [
-                    SearchResult(
-                        content="x",
-                        score=1.0,
-                        document_id="doc-1",
-                        document_uri="uri-other",
-                        document_meta={"arxiv_id": "uri-x"},
-                    )
+                    _result("uri-other", 1.0),
+                    _result("uri-x", 0.9),
+                    _result("uri-other", 0.8),
+                    _result("uri-x", 0.7),
                 ]
 
         with patch("evaluations.benchmark.HaikuRAG") as mock_haiku:
@@ -538,8 +538,9 @@ class TestRetrievalTarget:
                 self._spec(), AppConfig(), db_path=tmp_path / "test.lancedb"
             )
 
+        # uri-x is the only relevant document and ranks second of two
         assert result is not None
-        assert result["map"] == 1.0
+        assert result["map"] == 0.5
 
 
 class TestEvaluateDatasetCaseIds:

From ac9b2cbf817be1bfec30e0b860fd74e61606589a Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 6 Aug 2026 11:24:47 +0300
Subject: [PATCH 3/4] Load document blobs only when asked

---
 CHANGELOG.md                                  |  4 ++
 docs/python.md                                |  8 +++
 haiku_rag_slim/haiku/rag/client/documents.py  |  6 +-
 haiku_rag_slim/haiku/rag/client/rebuild.py    |  8 ++-
 .../haiku/rag/store/repositories/document.py  | 18 +++--
 tests/store/test_v0_58_0_migration.py         |  2 +-
 tests/test_client.py                          | 38 +++++++++-
 tests/test_document.py                        | 69 +++++++++++++++++++
 tests/test_rebuild.py                         | 26 ++++---
 9 files changed, 157 insertions(+), 22 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f4d1e816..08483dfa 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,10 @@
 # Changelog
 ## [Unreleased]
 
+### Changed
+
+- `get_document_by_id` / `get_document_by_uri` return content and the mutable attributes only; the docling structure and page-image blobs are no longer loaded. Load them with `DocumentRepository.get_docling_data` / `get_pages_data`, or pass `include_blobs=True` to the repository method.
+
 ### Fixed
 
 - Ingester jobs failing with an `obstore` `PermissionDeniedError`, `UnauthenticatedError`, `UnknownConfigurationKeyError` or `InvalidPathError` are dead-lettered instead of retried to `max_attempts`.
diff --git a/docs/python.md b/docs/python.md
index abaecc1e..48075426 100644
--- a/docs/python.md
+++ b/docs/python.md
@@ -94,6 +94,14 @@ By URI:
 doc = await client.get_document_by_uri("file:///path/to/document.pdf")
 ```
 
+Both return content, uri, title and metadata. The multi-MB docling blobs are
+loaded separately:
+
+```python
+docling = await client.document_repository.get_docling_data(doc.id)
+pages = await client.document_repository.get_pages_data(doc.id)
+```
+
 List all documents:
 ```python
 docs = await client.list_documents(limit=10, offset=0)
diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py
index 8cb916e4..f0455198 100644
--- a/haiku_rag_slim/haiku/rag/client/documents.py
+++ b/haiku_rag_slim/haiku/rag/client/documents.py
@@ -856,7 +856,11 @@ async def update_document(
             "Provide one or the other, not both."
         )
 
-    existing_doc = await client.get_document_by_id(document_id)
+    # An update that only replaces content writes the record back as-is, so
+    # the blobs have to make the round trip.
+    existing_doc = await client.document_repository.get_by_id(
+        document_id, include_blobs=True
+    )
     if existing_doc is None:
         raise ValueError(f"Document with ID {document_id} not found")
 
diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py
index 6ba9ca7b..639572de 100644
--- a/haiku_rag_slim/haiku/rag/client/rebuild.py
+++ b/haiku_rag_slim/haiku/rag/client/rebuild.py
@@ -187,7 +187,9 @@ async def _hydrate(
     """
     for light_doc in light_docs:
         assert light_doc.id is not None
-        doc = await client.get_document_by_id(light_doc.id)
+        doc = await client.document_repository.get_by_id(
+            light_doc.id, include_blobs=True
+        )
         if doc is None:
             continue
         assert doc.id is not None
@@ -822,7 +824,9 @@ async def _rebuild_full(
 
         # Fallback: rebuild from stored content. Now we need the full
         # record (content + docling_pages for the round-trip write).
-        doc = await client.get_document_by_id(light_doc.id)
+        doc = await client.document_repository.get_by_id(
+            light_doc.id, include_blobs=True
+        )
         if doc is None:
             continue
         assert doc.id is not None
diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py
index 8f332f0b..c31cdcae 100644
--- a/haiku_rag_slim/haiku/rag/store/repositories/document.py
+++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py
@@ -182,11 +182,16 @@ class DocumentRepository:
             raise
         return documents
 
-    async def get_by_id(self, entity_id: str) -> Document | None:
-        """Get a document by its ID."""
+    _LIGHT_COLUMNS = ["id", "content"]
+
+    async def get_by_id(
+        self, entity_id: str, include_blobs: bool = False
+    ) -> Document | None:
+        """Get a document by its ID. `include_blobs` adds the docling blobs."""
         safe_id = escape_sql_string(entity_id)
+        query = self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1)
         results = await query_to_pydantic(
-            self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1),
+            query if include_blobs else query.select(self._LIGHT_COLUMNS),
             DocumentRecord,
         )
 
@@ -373,7 +378,9 @@ class DocumentRepository:
         """Count documents with optional filtering (over document_meta columns)."""
         return await self.store.document_meta_table.count_rows(filter=filter)
 
-    async def get_by_uri(self, uri: str) -> Document | None:
+    async def get_by_uri(
+        self, uri: str, include_blobs: bool = False
+    ) -> Document | None:
         """Get a document by its URI (resolved via document_meta)."""
         escaped_uri = escape_sql_string(uri)
         meta_results = await query_to_pydantic(
@@ -388,8 +395,9 @@ class DocumentRepository:
 
         meta = meta_results[0]
         safe_id = escape_sql_string(meta.id)
+        query = self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1)
         doc_results = await query_to_pydantic(
-            self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1),
+            query if include_blobs else query.select(self._LIGHT_COLUMNS),
             DocumentRecord,
         )
         if not doc_results:
diff --git a/tests/store/test_v0_58_0_migration.py b/tests/store/test_v0_58_0_migration.py
index 286b4d05..30fe510e 100644
--- a/tests/store/test_v0_58_0_migration.py
+++ b/tests/store/test_v0_58_0_migration.py
@@ -59,7 +59,7 @@ class TestV0_58_0Migration:
 
             # Full hydration still works (content + metadata + blobs intact).
             repo = DocumentRepository(store)
-            doc = await repo.get_by_id("doc-1")
+            doc = await repo.get_by_id("doc-1", include_blobs=True)
             assert doc is not None
             assert doc.content == "body one"
             assert doc.uri == "s3://b/one"
diff --git a/tests/test_client.py b/tests/test_client.py
index 48de4ad5..db2da16f 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -984,13 +984,17 @@ async def test_metadata_only_update_does_not_advance_documents_table(temp_db_pat
         # The light document_meta table absorbs the updates.
         assert await client.store.document_meta_table.version() > meta_v0
 
-        # Reads still hydrate the full document (content + blobs + metadata).
+        # Reads still hydrate content and the mutable attributes together.
         fetched = await client.get_document_by_id(created.id)
         assert fetched is not None
         assert fetched.metadata["source_revision"] == "rev-5"
         assert fetched.title == "Title 5"
         assert fetched.content == "Body text"
-        assert fetched.get_docling_document() is not None
+
+        # And the untouched docling blob is still there.
+        docling = await client.document_repository.get_docling_data(created.id)
+        assert docling is not None
+        assert docling.get_docling_document() is not None
 
 
 async def test_delete_marks_vacuum_dirty(temp_db_path):
@@ -1202,7 +1206,9 @@ async def test_client_create_document_from_file_stores_docling_json(temp_db_path
             assert doc.docling_version is not None
 
             # Verify the stored document also has the JSON
-            retrieved = await client.get_document_by_id(doc.id)
+            retrieved = await client.document_repository.get_by_id(
+                doc.id, include_blobs=True
+            )
             assert retrieved is not None
             assert retrieved.docling_document == doc.docling_document
             assert retrieved.docling_version == doc.docling_version
@@ -2200,6 +2206,32 @@ async def test_update_document_with_url_prefixed_content(temp_db_path, monkeypat
         assert "New heading" in updated.content
 
 
+async def test_update_document_with_chunks_keeps_page_images(temp_db_path, monkeypatch):
+    """Replacing content and chunks without a docling document writes the stored
+    record back as-is, so its page rasters must survive the round trip."""
+    _patch_embed_chunks(monkeypatch)
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        doc = await client.create_document(content="initial body", uri="test://pages")
+        assert doc.id is not None
+
+        sentinel_pages = b"\x80SENTINEL_PAGE_BYTES"
+        await client.store.documents_table.update(
+            {"docling_pages": sentinel_pages}, where=f"id = '{doc.id}'"
+        )
+
+        await client.update_document(
+            doc.id,
+            content="replacement body",
+            chunks=[Chunk(content="replacement body")],
+        )
+
+        stored = await client.document_repository.get_by_id(doc.id, include_blobs=True)
+        assert stored is not None
+        assert stored.content == "replacement body"
+        assert stored.docling_pages == sentinel_pages
+
+
 async def test_rebuild_rechunk_with_url_prefixed_stored_content(
     temp_db_path, monkeypatch
 ):
diff --git a/tests/test_document.py b/tests/test_document.py
index a45fd03d..4f10b9f2 100644
--- a/tests/test_document.py
+++ b/tests/test_document.py
@@ -457,6 +457,75 @@ async def test_get_pages_data_loads_only_pages_column(
         assert await doc_repo.get_pages_data("nonexistent-id") is None
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("include_blobs", [False, True])
+async def test_document_get_by_id_docling_blobs(temp_db_path, include_blobs):
+    """get_by_id leaves the docling blobs out unless asked for them: a single
+    document's page rasters run to hundreds of MB."""
+    async with Store(temp_db_path, create=True) as store:
+        doc_repo = DocumentRepository(store)
+
+        created = await doc_repo.create(
+            Document(
+                content="the text",
+                uri="https://example.com/doc.pdf",
+                title="Test Document",
+                metadata={"key": "value"},
+                docling_document=b"structure-blob",
+                docling_pages=b"page-raster-blob",
+                docling_version="2.1.0",
+            )
+        )
+        assert created.id is not None
+
+        doc = await doc_repo.get_by_id(created.id, include_blobs=include_blobs)
+
+        assert doc is not None
+        assert doc.id == created.id
+        assert doc.content == "the text"
+        assert doc.uri == "https://example.com/doc.pdf"
+        assert doc.title == "Test Document"
+        assert doc.metadata == {"key": "value"}
+        if include_blobs:
+            assert doc.docling_document == b"structure-blob"
+            assert doc.docling_pages == b"page-raster-blob"
+            assert doc.docling_version == "2.1.0"
+        else:
+            assert doc.docling_document is None
+            assert doc.docling_pages is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("include_blobs", [False, True])
+async def test_document_get_by_uri_docling_blobs(temp_db_path, include_blobs):
+    """get_by_uri has the same projection as get_by_id."""
+    async with Store(temp_db_path, create=True) as store:
+        doc_repo = DocumentRepository(store)
+
+        created = await doc_repo.create(
+            Document(
+                content="the text",
+                uri="https://example.com/doc.pdf",
+                docling_document=b"structure-blob",
+                docling_pages=b"page-raster-blob",
+            )
+        )
+
+        doc = await doc_repo.get_by_uri(
+            "https://example.com/doc.pdf", include_blobs=include_blobs
+        )
+
+        assert doc is not None
+        assert doc.id == created.id
+        assert doc.content == "the text"
+        if include_blobs:
+            assert doc.docling_document == b"structure-blob"
+            assert doc.docling_pages == b"page-raster-blob"
+        else:
+            assert doc.docling_document is None
+            assert doc.docling_pages is None
+
+
 @pytest.mark.asyncio
 async def test_document_get_by_uri_with_special_characters(
     qa_corpus: list[dict[str, str]], temp_db_path
diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py
index f76c5592..7e4a5ce6 100644
--- a/tests/test_rebuild.py
+++ b/tests/test_rebuild.py
@@ -36,7 +36,9 @@ async def test_rebuild_full(qa_corpus: list[dict[str, str]], 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)
+        doc_after = await client.document_repository.get_by_id(
+            doc.id, include_blobs=True
+        )
         assert doc_after is not None
         assert doc_after.docling_document is not None
         assert doc_after.docling_version is not None
@@ -70,7 +72,9 @@ async def test_rebuild_embed_only(qa_corpus: list[dict[str, str]], 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)
+        doc_after = await client.document_repository.get_by_id(
+            doc.id, include_blobs=True
+        )
         assert doc_after is not None
         assert doc_after.docling_document == original_docling_json
 
@@ -498,7 +502,9 @@ async def test_rebuild_rechunk(qa_corpus: list[dict[str, str]], temp_db_path):
         assert doc.id in processed_ids
 
         # Document content should be unchanged, but docling JSON should be updated
-        doc_after = await client.document_repository.get_by_id(doc.id)
+        doc_after = await client.document_repository.get_by_id(
+            doc.id, include_blobs=True
+        )
         assert doc_after is not None
         assert doc_after.content == content_before
         assert doc_after.docling_document is not None
@@ -707,7 +713,7 @@ async def test_rebuild_descriptions_patches_blob_and_chunks(temp_db_path, monkey
         )
 
         from_blob = (
-            await rag.document_repository.get_by_id(created.id)
+            await rag.document_repository.get_by_id(created.id, include_blobs=True)
         ).get_docling_document()  # type: ignore[union-attr]
         assert from_blob is not None and from_blob.pictures
         # No description in the freshly-ingested doc
@@ -737,7 +743,7 @@ async def test_rebuild_descriptions_patches_blob_and_chunks(temp_db_path, monkey
         assert created.id in processed
 
         # The stored docling blob now has the description
-        after = await rag.document_repository.get_by_id(created.id)
+        after = await rag.document_repository.get_by_id(created.id, include_blobs=True)
         assert after is not None
         after_doc = after.get_docling_document()
         assert after_doc is not None and after_doc.pictures
@@ -799,7 +805,7 @@ async def test_rebuild_descriptions_skips_already_described(temp_db_path, monkey
         # VLM was never called for this picture (it already had a description)
         assert called_with == [] or all(not d for d in called_with)
 
-        after = await rag.document_repository.get_by_id(created.id)
+        after = await rag.document_repository.get_by_id(created.id, include_blobs=True)
         assert after is not None
         after_doc = after.get_docling_document()
         assert after_doc is not None
@@ -1142,10 +1148,10 @@ async def test_hydrate_skips_documents_deleted_mid_rebuild(temp_db_path):
             Document(content="body", uri="test://gone")
         )
 
-        async def vanished(_document_id):
+        async def vanished(_document_id, include_blobs=False):
             return None
 
-        client.get_document_by_id = vanished  # type: ignore[method-assign]
+        client.document_repository.get_by_id = vanished  # type: ignore[method-assign]
 
         assert [doc async for doc in _hydrate(client, [stored])] == []
 
@@ -1438,10 +1444,10 @@ async def test_rebuild_full_skips_document_deleted_mid_rebuild(temp_db_path):
         doc = await client.create_document(content="doc that disappears")
         assert doc.id is not None
 
-        async def vanished(_document_id):
+        async def vanished(_document_id, include_blobs=False):
             return None
 
-        client.get_document_by_id = vanished  # type: ignore[method-assign]
+        client.document_repository.get_by_id = vanished  # type: ignore[method-assign]
 
         processed = [
             doc_id async for doc_id in client.rebuild_database(mode=RebuildMode.FULL)

From eb11a165b810a23541b4f1775a6d7e6745df29a8 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 6 Aug 2026 12:28:18 +0300
Subject: [PATCH 4/4] Stop loading page rasters on the title and update paths

---
 CHANGELOG.md                                  |  2 +-
 evaluations/tests/test_benchmark.py           | 35 +++++++++----------
 haiku_rag_slim/haiku/rag/client/documents.py  |  6 ++--
 haiku_rag_slim/haiku/rag/client/rebuild.py    | 20 +++++++----
 .../haiku/rag/store/repositories/document.py  | 33 +++++++++--------
 tests/test_rebuild.py                         | 28 +++++++++++++++
 6 files changed, 79 insertions(+), 45 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 08483dfa..6a79c487 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,7 +3,7 @@
 
 ### Changed
 
-- `get_document_by_id` / `get_document_by_uri` return content and the mutable attributes only; the docling structure and page-image blobs are no longer loaded. Load them with `DocumentRepository.get_docling_data` / `get_pages_data`, or pass `include_blobs=True` to the repository method.
+- `get_document_by_id` / `get_document_by_uri` no longer load the docling structure and page-image blobs. Load them with `DocumentRepository.get_docling_data` / `get_pages_data`, or `get_by_id(..., include_blobs=True)`.
 
 ### Fixed
 
diff --git a/evaluations/tests/test_benchmark.py b/evaluations/tests/test_benchmark.py
index 2e49a836..23e36b46 100644
--- a/evaluations/tests/test_benchmark.py
+++ b/evaluations/tests/test_benchmark.py
@@ -15,6 +15,20 @@ from evaluations.config import DatasetSpec
 from haiku.rag.config.models import AppConfig, ModelConfig
 
 
+def _stub_spec(**overrides) -> DatasetSpec:
+    """A DatasetSpec whose loaders/mappers are inert, for tests that only
+    exercise the surrounding plumbing."""
+    return DatasetSpec(
+        key="test",
+        db_filename="test.lancedb",
+        document_loader=lambda: None,  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        document_mapper=lambda doc: None,
+        qa_loader=lambda: [],  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        qa_case_builder=lambda idx, doc: None,  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
+        **overrides,
+    )
+
+
 class TestBuildExperimentMetadata:
     def test_basic_metadata(self) -> None:
         config = AppConfig()
@@ -460,16 +474,8 @@ class TestRetrievalTarget:
         from evaluations.config import RetrievalSample
         from evaluations.evaluators import MAPEvaluator
 
-        return DatasetSpec(
-            key="test",
-            db_filename="test.lancedb",
-            document_loader=lambda: None,  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
-            document_mapper=lambda doc: None,
-            qa_loader=lambda: [],  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
-            qa_case_builder=lambda idx, doc: None,  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
-            retrieval_loader=lambda: [  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
-                {"q": "What is X?", "uris": ("uri-x",)},
-            ],
+        return _stub_spec(
+            retrieval_loader=lambda: [{"q": "What is X?", "uris": ("uri-x",)}],
             retrieval_mapper=lambda d: RetrievalSample(
                 question=d["q"], expected_uris=d["uris"]
             ),
@@ -545,14 +551,7 @@ class TestRetrievalTarget:
 
 class TestEvaluateDatasetCaseIds:
     def _spec(self) -> DatasetSpec:
-        return DatasetSpec(
-            key="test",
-            db_filename="test.lancedb",
-            document_loader=lambda: None,  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
-            document_mapper=lambda doc: None,
-            qa_loader=lambda: [],  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
-            qa_case_builder=lambda idx, doc: None,  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
-        )
+        return _stub_spec()
 
     @pytest.mark.asyncio
     async def test_threads_case_ids_to_qa_benchmark(self) -> None:
diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py
index f0455198..8facaba2 100644
--- a/haiku_rag_slim/haiku/rag/client/documents.py
+++ b/haiku_rag_slim/haiku/rag/client/documents.py
@@ -856,10 +856,10 @@ async def update_document(
             "Provide one or the other, not both."
         )
 
-    # An update that only replaces content writes the record back as-is, so
-    # the blobs have to make the round trip.
+    # Caller-supplied chunks without a docling document replace neither blob,
+    # and the row is written back whole, so they have to make the round trip.
     existing_doc = await client.document_repository.get_by_id(
-        document_id, include_blobs=True
+        document_id, include_blobs=chunks is not None and docling_document is None
     )
     if existing_doc is None:
         raise ValueError(f"Document with ID {document_id} not found")
diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py
index 639572de..9374a2f4 100644
--- a/haiku_rag_slim/haiku/rag/client/rebuild.py
+++ b/haiku_rag_slim/haiku/rag/client/rebuild.py
@@ -176,7 +176,7 @@ async def _set_embedder(client: "HaikuRAG") -> None:
 
 
 async def _hydrate(
-    client: "HaikuRAG", light_docs: list[Document]
+    client: "HaikuRAG", light_docs: list[Document], include_blobs: bool = True
 ) -> AsyncGenerator[Document, None]:
     """Yield fully-loaded documents one at a time from a light listing.
 
@@ -188,7 +188,7 @@ async def _hydrate(
     for light_doc in light_docs:
         assert light_doc.id is not None
         doc = await client.document_repository.get_by_id(
-            light_doc.id, include_blobs=True
+            light_doc.id, include_blobs=include_blobs
         )
         if doc is None:
             continue
@@ -199,9 +199,18 @@ async def _hydrate(
 async def _rebuild_title_only(
     client: "HaikuRAG", documents: list[Document]
 ) -> AsyncGenerator[str, None]:
-    """Generate titles for documents that don't have one."""
+    """Generate titles for documents that don't have one.
+
+    A title comes from the content or the docling structure, never the page
+    rasters, so those are left out of the per-document load.
+    """
+    repo = client.document_repository
     untitled = [d for d in documents if d.title is None]
-    async for doc in _hydrate(client, untitled):
+    async for doc in _hydrate(client, untitled, include_blobs=False):
+        assert doc.id is not None
+        structure = await repo.get_docling_data(doc.id)
+        if structure is not None:
+            doc.docling_document = structure.docling_document
         try:
             title = await client.generate_title(doc)
         except Exception:
@@ -211,8 +220,7 @@ async def _rebuild_title_only(
             continue
         if title is not None:
             doc.title = title
-            await client.document_repository.update_meta(doc)
-            assert doc.id is not None
+            await repo.update_meta(doc)
             yield doc.id
 
 
diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py
index c31cdcae..9bfd9b57 100644
--- a/haiku_rag_slim/haiku/rag/store/repositories/document.py
+++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py
@@ -184,22 +184,26 @@ class DocumentRepository:
 
     _LIGHT_COLUMNS = ["id", "content"]
 
+    async def _record_by_id(
+        self, doc_id: str, include_blobs: bool
+    ) -> DocumentRecord | None:
+        safe_id = escape_sql_string(doc_id)
+        query = self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1)
+        if not include_blobs:
+            query = query.select(self._LIGHT_COLUMNS)
+        results = await query_to_pydantic(query, DocumentRecord)
+        return results[0] if results else None
+
     async def get_by_id(
         self, entity_id: str, include_blobs: bool = False
     ) -> Document | None:
         """Get a document by its ID. `include_blobs` adds the docling blobs."""
-        safe_id = escape_sql_string(entity_id)
-        query = self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1)
-        results = await query_to_pydantic(
-            query if include_blobs else query.select(self._LIGHT_COLUMNS),
-            DocumentRecord,
-        )
-
-        if not results:
+        record = await self._record_by_id(entity_id, include_blobs)
+        if record is None:
             return None
 
         meta = await self._meta_by_id(entity_id)
-        return self._merge_to_document(results[0], meta)
+        return self._merge_to_document(record, meta)
 
     async def get_content(self, entity_id: str) -> str | None:
         """Get only the text content of a document (skips docling blobs)."""
@@ -394,16 +398,11 @@ class DocumentRepository:
             return None
 
         meta = meta_results[0]
-        safe_id = escape_sql_string(meta.id)
-        query = self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1)
-        doc_results = await query_to_pydantic(
-            query if include_blobs else query.select(self._LIGHT_COLUMNS),
-            DocumentRecord,
-        )
-        if not doc_results:
+        record = await self._record_by_id(meta.id, include_blobs)
+        if record is None:
             return None
 
-        return self._merge_to_document(doc_results[0], meta)
+        return self._merge_to_document(record, meta)
 
     async def delete_all(self) -> None:
         """Delete all documents from the database."""
diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py
index 7e4a5ce6..bcc41f4b 100644
--- a/tests/test_rebuild.py
+++ b/tests/test_rebuild.py
@@ -551,6 +551,34 @@ async def test_rebuild_full_with_accessible_source(temp_db_path):
             assert "Fresh content" in new_doc.content
 
 
+async def test_rebuild_title_only_reads_structural_title(temp_db_path):
+    """TITLE_ONLY takes the title from the stored docling structure, so it never
+    reaches the LLM for a document that carries one."""
+    from docling_core.types.doc.document import DoclingDocument
+    from docling_core.types.doc.labels import DocItemLabel
+
+    from haiku.rag.store.models.document import Document
+
+    docling_doc = DoclingDocument(name="structured")
+    docling_doc.add_text(label=DocItemLabel.TITLE, text="The Stored Title")
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        doc = Document(content="body text", metadata={})
+        doc.set_docling(docling_doc)
+        created = await client.document_repository.create(doc)
+        assert created.id is not None
+
+        processed_ids = [
+            doc_id
+            async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY)
+        ]
+
+        assert processed_ids == [created.id]
+        refreshed = await client.get_document_by_id(created.id)
+        assert refreshed is not None
+        assert refreshed.title == "The Stored Title"
+
+
 async def test_rebuild_title_only_handles_llm_failure(temp_db_path, monkeypatch):
     """TITLE_ONLY: a failure on one document does not abort the generator.