diff --git a/CHANGELOG.md b/CHANGELOG.md index f4d1e816..6a79c487 100644 --- a/CHANGELOG.md +++ b/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`. 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/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index 2541ca05..a550bc9f 100644 --- a/evaluations/evaluations/benchmark.py +++ b/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 diff --git a/evaluations/tests/test_benchmark.py b/evaluations/tests/test_benchmark.py index dcdc27da..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() @@ -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: diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index 8cb916e4..8facaba2 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) + # 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") diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index 6ba9ca7b..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. @@ -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 diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 8f332f0b..9bfd9b57 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/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.""" 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..bcc41f4b 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 @@ -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)