Stop loading page rasters on the title and update paths

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
 

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(
+        return _stub_spec(
-            key="test",
+            retrieval_loader=lambda: [{"q": "What is X?", "uris": ("uri-x",)}],
-            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"]
             ),
@@ -545,14 +551,7 @@ class TestRetrievalTarget:
 
 class TestEvaluateDatasetCaseIds:
     def _spec(self) -> DatasetSpec:
-        return DatasetSpec(
+        return _stub_spec()
-            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]
-        )
 
     @pytest.mark.asyncio
     async def test_threads_case_ids_to_qa_benchmark(self) -> None:

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
+    # Caller-supplied chunks without a docling document replace neither blob,
-    # the blobs have to make the round trip.
+    # 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")

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)
+            await repo.update_meta(doc)
-            assert doc.id is not None
             yield doc.id
 
 

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)
+        record = await self._record_by_id(entity_id, include_blobs)
-        query = self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1)
+        if record is None:
-        results = await query_to_pydantic(
-            query if include_blobs else query.select(self._LIGHT_COLUMNS),
-            DocumentRecord,
-        )
-
-        if not results:
             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)
+        record = await self._record_by_id(meta.id, include_blobs)
-        query = self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1)
+        if record is None:
-        doc_results = await query_to_pydantic(
-            query if include_blobs else query.select(self._LIGHT_COLUMNS),
-            DocumentRecord,
-        )
-        if not doc_results:
             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."""

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.