Merge pull request #526 from ggozad/fix/light-document-lookups

Load document blobs only when asked

CHANGELOG.md

@@ -1,6 +1,10 @@
 # Changelog
 ## [Unreleased]
 
+### Changed
+
+- `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
 
 - Ingester jobs failing with an `obstore` `PermissionDeniedError`, `UnauthenticatedError`, `UnknownConfigurationKeyError` or `InvalidPathError` are dead-lettered instead of retried to `max_attempts`.

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)

evaluations/evaluations/benchmark.py

@@ -230,23 +230,15 @@ 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
-                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
 

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()
@@ -455,16 +469,89 @@ 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 _stub_spec(
+            retrieval_loader=lambda: [{"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_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 [
+                    _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:
+            mock_haiku.return_value.__aenter__.return_value = FakeRag()
+            result = await run_retrieval_benchmark(
+                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"] == 0.5
+
+
 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:

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)
+    # 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=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.
 
@@ -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=include_blobs
+        )
         if doc is None:
             continue
         assert doc.id is not None
@@ -197,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:
@@ -209,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
 
 
@@ -822,7 +832,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

haiku_rag_slim/haiku/rag/store/repositories/document.py

@@ -182,19 +182,28 @@ class DocumentRepository:
             raise
         return documents
 
-    async def get_by_id(self, entity_id: str) -> Document | None:
-        """Get a document by its ID."""
-        safe_id = escape_sql_string(entity_id)
-        results = await query_to_pydantic(
-            self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1),
-            DocumentRecord,
-        )
+    _LIGHT_COLUMNS = ["id", "content"]
 
-        if not results:
+    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."""
+        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)."""
@@ -373,7 +382,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(
@@ -387,15 +398,11 @@ class DocumentRepository:
             return None
 
         meta = meta_results[0]
-        safe_id = escape_sql_string(meta.id)
-        doc_results = await query_to_pydantic(
-            self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1),
-            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."""

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"

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
 ):

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

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
@@ -545,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.
 
@@ -707,7 +741,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 +771,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 +833,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 +1176,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 +1472,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)