From eb11a165b810a23541b4f1775a6d7e6745df29a8 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 6 Aug 2026 12:28:18 +0300 Subject: [PATCH] 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.